MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / close

Method close

tensorrt_llm/executor/rpc/rpc_client.py:158–214  ·  view source on GitHub ↗

Gracefully close the client, cleaning up background tasks.

(self)

Source from the content-addressed store, hash-verified

156 self._server_stopped = True
157
158 def close(self):
159 """Gracefully close the client, cleaning up background tasks."""
160
161 if self._closed:
162 return
163 self._closed = True
164
165 logger_debug("[client] RPC Client closing")
166
167 # Notify any active streaming consumers so they don't hang waiting for
168 # further data. This must be done *before* shutting down the event
169 # loop/executor because they may depend on the loop to complete.
170 self._broadcast_streaming_error(RPCCancelled("RPC client closed"))
171
172 # 1. Cancel the reader task
173 if self._reader_task and not self._reader_task.done():
174 if self._loop and self._loop.is_running(
175 ) and self._reader_asyncio_task:
176 try:
177
178 async def cancel_reader_task():
179 if self._reader_asyncio_task and not self._reader_asyncio_task.done(
180 ):
181 self._reader_asyncio_task.cancel()
182 try:
183 await self._reader_asyncio_task
184 except asyncio.CancelledError:
185 pass
186
187 cancel_future = asyncio.run_coroutine_threadsafe(
188 cancel_reader_task(), self._loop)
189 cancel_future.result(timeout=2.0)
190 logger_debug("[client] Reader task cancelled successfully")
191 except concurrent.futures.TimeoutError:
192 logger.warning("Reader task did not exit gracefully")
193 except Exception as e:
194 logger_debug(f"[client] Reader task cleanup: {e}")
195
196 # 2. Stop the event loop
197 if self._loop and self._loop.is_running():
198 self._loop.call_soon_threadsafe(self._loop.stop)
199
200 # 3. Join the event loop thread
201 if self._loop_thread:
202 self._loop_thread.join(timeout=2.0)
203 if self._loop_thread.is_alive():
204 logger.warning("Event loop thread did not exit gracefully")
205
206 # 4. Shutdown the executor
207 if self._executor:
208 self._executor.shutdown(wait=True)
209
210 # 5. Close the socket
211 if self._client_socket:
212 self._client_socket.close()
213
214 logger_debug("[client] RPC Client closed")
215

Callers 5

__exit__Method · 0.95
shutdownMethod · 0.45
run_loopMethod · 0.45

Calls 6

logger_debugFunction · 0.85
RPCCancelledClass · 0.85
resultMethod · 0.45
warningMethod · 0.45
shutdownMethod · 0.45

Tested by 2