Log more info

Author: HuiGao-NV

tensorrt_llm/executor/base_worker.py

@@ -638,6 +638,7 @@ def submit(self, request: GenerationRequest) -> GenerationResult:
         return result
 
     def shutdown(self):
+        print(f"========================== Shutting down worker {self.rank}")
         if self.doing_shutdown:
             return
         else:

tensorrt_llm/executor/proxy.py

@@ -290,6 +290,9 @@ def pre_shutdown(self):
             self.request_queue.put_noblock(None, retry=4)
 
     def shutdown(self):
+        print(
+            f"====================== shutdown in GenerationExecutorProxy is called pid:  {os.getpid()}"
+        )
         if not self.workers_started:
             return
 
@@ -325,6 +328,10 @@ def shutdown(self):
         self.result_queue.close()
 
         self.workers_started = False
+        print(
+            f"====================== shutdown in GenerationExecutorProxy 2 is called pid:  {os.getpid()}"
+        )
+
         self.mpi_session.shutdown()
 
         # Process the errors in-case error during shutting down the threads

tensorrt_llm/executor/ray_executor.py

@@ -294,7 +294,9 @@ def shutdown(self):
             self._shutdown_event.set()
 
         logger_debug(f"Shutting down RayExecutor", color="yellow")
-
+        print(
+            f"====================== shutdown in RayExecutor is called pid:  {os.getpid()}"
+        )
         if hasattr(self, 'main_loop') and self.main_loop and hasattr(
                 self, 'main_loop_task_obj') and self.main_loop_task_obj:
             logger_debug("Cancelling main loop task.", color="yellow")

tensorrt_llm/executor/ray_gpu_worker.py

@@ -151,6 +151,9 @@ def call_worker_method(self, method_name: str, *args, **kwargs):
                 f"The RayGPUWorker has no method called '{method_name}'.")
 
     def shutdown(self):
+        print(
+            f"====================== shutdown in RayWorkerWrapper is called pid:  {os.getpid()}"
+        )
         if hasattr(self, 'worker'):
             self.worker.shutdown()
 
@@ -298,6 +301,9 @@ def shutdown(self):
             return
         else:
             self.doing_shutdown = True
+        print(
+            f"====================== shutdown in RayWorkerWrapper is called pid:  {os.getpid()}"
+        )
 
         logger.debug(f'Worker {self.rank} shutting down...')
 

tensorrt_llm/executor/rpc/rpc_server.py

@@ -134,6 +134,9 @@ def shutdown(self, is_remote_call: bool = False) -> None:
         # Set the stop event to True, this will trigger immediate shutdown
         self._stop_event.set()
 
+        print(
+            f"==================================== RPCServer shutdown called, is_remote_call={is_remote_call}"
+        )
         # Log pending requests that will be cancelled
         logger_debug(
             f"[server] RPCServer is shutting down: {self._num_pending_requests} pending requests will be cancelled"

tensorrt_llm/executor/rpc_proxy.py

@@ -191,12 +191,12 @@ def shutdown(self):
         traceback.print_stack()
         import os
         print(
-            f"====================== shutdown in generator is called pid:  {os.getpid()}"
+            f"====================== shutdown in GenerationExecutorRpcProxy is called pid:  {os.getpid()}"
         )
         if self._shutdown_event.is_set():
             return
         print(
-            f"====================== shutdown in generator 2 is called pid:  {os.getpid()}"
+            f"====================== shutdown in GenerationExecutorRpcProxy 2 is called pid:  {os.getpid()}"
         )
         self._shutdown_event.set()
         logger_debug(f"Shutting down GenerationExecutorRpcProxy",

tensorrt_llm/executor/rpc_worker.py

@@ -102,6 +102,9 @@ def shutdown(self):
         logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutting down",
                      color="yellow")
         self.shutdown_event.set()
+        print(
+            f"====================== shutdown in RpcWorker is called pid:  {os.getpid()}"
+        )
         super().shutdown()
         logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutdown",
                      color="yellow")

tensorrt_llm/executor/utils.py

@@ -101,6 +101,9 @@ def submit_sync(self, task: Callable, *args, **kwargs) -> List[Any]:
         return [future.result() for future in futures]
 
     def shutdown(self):
+        print(
+            f"==================================== shutdown ProcessPoolExecutor session"
+        )
         self.mpi_pool.shutdown(wait=True)
 
 

tensorrt_llm/executor/worker.py

@@ -101,6 +101,9 @@ def shutdown(self):
                     self.await_response_thread.stop()
                     self.await_response_thread.join()
 
+            print(
+                f"====================== GenerationExecutorWorker engine shutdown is called pid:  {os.getpid()}"
+            )
             self.engine.shutdown()
             self.engine = None
 
@@ -320,6 +323,9 @@ def notify_proxy_threads_to_quit():
                     else:
                         raise ValueError(f"Unknown request type: {type(req)}")
 
+                print(
+                    f"====================== Worker {mpi_rank()} received shutdown signal from proxy process."
+                )
                 notify_proxy_threads_to_quit()
 
         except GenerationExecutorWorker.WorkerExit as e:
@@ -328,7 +334,11 @@ def notify_proxy_threads_to_quit():
 
         except Exception as e:  # other critical errors
             if is_leader:
+                print(
+                    f"====================== Worker {mpi_rank()} received shutdown signal from proxy process."
+                )
                 notify_proxy_threads_to_quit()
             logger.error(traceback.format_exc())
             # This will be captured by mpi4py and handled by future.done_callback
             raise e
+    logger_debug(f"Worker {mpi_rank()} exiting worker_main...\n", "green")

tensorrt_llm/llmapi/mpi_session.py

@@ -160,6 +160,7 @@ def submit_sync(self, task: Callable[..., T], *args, **kwargs) -> List[T]:
         return [future.result() for future in futures]
 
     def shutdown(self, wait=True):
+        print(f"==================================== shutdown MPI pool session")
         if self.mpi_pool is not None:
             self.mpi_pool.shutdown(wait=wait)
             self.mpi_pool = None
@@ -237,6 +238,9 @@ def submit_sync(self, task: Callable[..., T], *args, **kwargs) -> List[T]:
     def shutdown(self, wait=True):
         # Only shutdown the mpi_pool if this instance created it
         # For shared global mpi_pool, we don't shut it down
+        print(
+            f"==================================== shutdown is called MPI comm session"
+        )
         if self.mpi_pool is not None and self.owns_mpi_pool:
             self.mpi_pool.shutdown(wait=wait)
         self.mpi_pool = None