add gc trace in LLM

Signed-off-by: Superjomn <[email protected]>

Author: Superjomn

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1,13 +1,11 @@
 import dataclasses
 import datetime
 import functools
-import gc
 import os
 import pickle  # nosec B403
 import threading
 import time
 import traceback
-import weakref
 from contextlib import contextmanager
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
@@ -22,8 +20,9 @@
 
 from tensorrt_llm._torch.pyexecutor.resource_manager import (
     ResourceManagerType, request_context)
-from tensorrt_llm._utils import (customized_gc_thresholds, is_trace_enabled,
-                                 mpi_disabled, nvtx_range, trace_func)
+from tensorrt_llm._utils import (customized_gc_thresholds, gc_nvtx_watcher,
+                                 is_trace_enabled, mpi_disabled, nvtx_range,
+                                 trace_func)
 from tensorrt_llm.bindings.executor import (DisServingRequestStats,
                                             FinishReason, InflightBatchingStats,
                                             IterationStats, KvCacheStats,
@@ -59,10 +58,6 @@
 # Format: "start1-stop1,start2-stop2,..." or single iterations "iter1,iter2,..."
 PROFILE_START_STOP_ENV_VAR_NAME = "TLLM_PROFILE_START_STOP"
 
-# Environment variable to enable garbage collection profiling.
-# Set to "1" to enable recording of garbage collection events during profiling.
-PROFILE_RECORD_GC_ENV_VAR_NAME = "TLLM_PROFILE_RECORD_GC"
-
 # Environment variable to enable PyTorch profiler tracing.
 # Set to a path to save detailed tracing of PyTorch operations.
 PROFILE_TRACE_ENV_VAR_NAME = "TLLM_TORCH_PROFILE_TRACE"
@@ -97,40 +92,6 @@ def _load_iteration_indexes(env_var: str):
     return frozenset(starts), frozenset(stops)
 
 
-class _GCNvtxHandle:
-    pass
-
-
-def _gc_nvtx_watcher():
-    enabled = os.environ.get(PROFILE_RECORD_GC_ENV_VAR_NAME, None)
-    if not enabled:
-        return None
-
-    range_id: Optional[int] = None
-
-    def gc_callback(phase, _):
-        nonlocal range_id
-        if phase == "start":
-            assert range_id is None, "Unexpected state in GC callback: another GC while last GC not finished?"
-            range_id = torch.cuda.nvtx.range_start("Python GC")
-        elif phase == "stop":
-            assert range_id is not None, "Unexpected state in GC callback: no active GC but got GC finished?"
-            torch.cuda.nvtx.range_end(range_id)
-            range_id = None
-
-    gc.callbacks.append(gc_callback)
-
-    def gc_cleanup(callback):
-        try:
-            gc.callbacks.remove(callback)
-        except ValueError:
-            pass
-
-    handle = _GCNvtxHandle()
-    weakref.finalize(handle, gc_cleanup, gc_callback)
-    return handle
-
-
 @dataclasses.dataclass
 class BatchState:
     sample_state: SampleState
@@ -178,7 +139,7 @@ def __init__(self,
         # profile config
         self.profile_start_iters, self.profile_stop_iters = _load_iteration_indexes(
             PROFILE_START_STOP_ENV_VAR_NAME)
-        self.gc_nvtx_watcher_handle = _gc_nvtx_watcher()
+        self.gc_nvtx_watcher_handle = gc_nvtx_watcher()
 
         # related modules
         self.resource_manager = resource_manager

tensorrt_llm/_utils.py

@@ -1191,3 +1191,50 @@ def is_device_integrated() -> bool:
     if not torch.cuda.is_available():
         return False
     return torch.cuda.get_device_properties().is_integrated
+
+
+# Environment variable to enable garbage collection profiling.
+# Set to "1" to enable recording of garbage collection events during profiling.
+PROFILE_RECORD_GC_ENV_VAR_NAME = "TLLM_PROFILE_RECORD_GC"
+
+
+class _GCNvtxHandle:
+    """Handle object for GC NVTX watcher to keep it alive."""
+
+
+def gc_nvtx_watcher() -> Optional[_GCNvtxHandle]:
+    """
+    Set up NVTX range markers for Python garbage collection events.
+    This helps in profiling to visualize when GC occurs during execution.
+
+    Returns:
+        _GCNvtxHandle or None: A handle object that keeps the GC callback alive,
+                               or None if GC profiling is not enabled.
+    """
+    enabled = os.environ.get(PROFILE_RECORD_GC_ENV_VAR_NAME, None)
+    if not enabled:
+        return None
+
+    range_id: Optional[int] = None
+
+    def gc_callback(phase, _):
+        nonlocal range_id
+        if phase == "start":
+            assert range_id is None, "Unexpected state in GC callback: another GC while last GC not finished?"
+            range_id = torch.cuda.nvtx.range_start("Python GC")
+        elif phase == "stop":
+            assert range_id is not None, "Unexpected state in GC callback: no active GC but got GC finished?"
+            torch.cuda.nvtx.range_end(range_id)
+            range_id = None
+
+    gc.callbacks.append(gc_callback)
+
+    def gc_cleanup(callback):
+        try:
+            gc.callbacks.remove(callback)
+        except ValueError:
+            pass
+
+    handle = _GCNvtxHandle()
+    weakref.finalize(handle, gc_cleanup, gc_callback)
+    return handle

tensorrt_llm/executor/rpc/rpc_client.py

@@ -8,7 +8,9 @@
 
 from tensorrt_llm._utils import nvtx_mark_debug
 
-from ...llmapi.utils import AsyncQueue, _SyncQueue, logger_debug
+from ..._utils import nvtx_range_debug
+from ...llmapi.utils import (AsyncQueue, _SyncQueue, enable_llmapi_debug,
+                             logger_debug)
 from ...logger import logger
 from ..ipc import ZeroMqQueue
 from .rpc_common import (RPCCancelled, RPCParams, RPCRequest, RPCResponse,
@@ -175,8 +177,10 @@ def _handle_streaming_response(self, response: RPCResponse):
             # put to the sync queue, as the current event loop is
             # different from the one in call_async or call_streaming
             assert isinstance(queue, AsyncQueue)
-            logger_debug(
-                f"RPC Client putting response to AsyncQueue: {response}")
+            if enable_llmapi_debug() or logger.level == 'debug':
+                logger_debug(
+                    f"RPC Client putting response to AsyncQueue: status={response.stream_status}, request_id={response.request_id}"
+                )
             queue.sync_q.put(response)
             # Clean up if stream ended
             if response.stream_status in ['end', 'error']:
@@ -188,32 +192,29 @@ def _handle_regular_response(self, response: RPCResponse):
         Args:
             response: The response to handle
         """
-        logger_debug(
-            f"Handling regular response for request_id: {response.request_id}")
-
         if future_info := self._pending_futures.get(response.request_id):
             future, target_loop = future_info
-            logger_debug(
-                f"Found future for request_id: {response.request_id}, future done: {future.done()}"
-            )
 
             if not future.done():
                 if response.error is None:
-                    logger_debug(
-                        f"Setting result for request_id: {response.request_id}, result: {response.result}"
-                    )
+                    if enable_llmapi_debug() or logger.level == 'debug':
+                        logger_debug(
+                            f"Setting result for request_id: {response.request_id}"
+                        )
                     target_loop.call_soon_threadsafe(future.set_result,
                                                      response.result)
                 else:
                     # Use the original RPCError from the response
-                    logger_debug(
-                        f"Setting exception for request_id: {response.request_id}, error: {response.error}"
-                    )
+                    if enable_llmapi_debug() or logger.level == 'debug':
+                        logger_debug(
+                            f"Setting exception for request_id: {response.request_id}, error: {response.error}"
+                        )
                     target_loop.call_soon_threadsafe(future.set_exception,
                                                      response.error)
         else:
-            logger_debug(
-                f"No future found for request_id: {response.request_id}")
+            if enable_llmapi_debug() or logger.level == 'debug':
+                logger_debug(
+                    f"No future found for request_id: {response.request_id}")
 
         self._pending_futures.pop(response.request_id, None)
 
@@ -256,35 +257,43 @@ async def _response_reader(self):
         logger_debug("Response reader started")
 
         while not self._stop_event.is_set():
-            try:
-                response = await self._wait_for_response()
-                if response is None:
-                    continue
-
-                nvtx_mark_debug(
-                    f"RPC.response.{'streaming' if response.is_streaming else 'sync'}",
-                    color="black",
-                    category="RPC")
-
-                logger_debug(f"RPC Client received response: {response}")
-                logger_debug(
-                    f"Response request_id: {response.request_id}, is_streaming: {response.is_streaming}"
-                )
-                logger_debug(
-                    f"Pending futures: {list(self._pending_futures.keys())}")
-
-                if response.is_streaming:
-                    self._handle_streaming_response(response)
-                else:
-                    self._handle_regular_response(response)
+            with nvtx_range_debug("response_reader",
+                                  color="cyan",
+                                  category="RPC"):
+                try:
+                    response = await self._wait_for_response()
+
+                    if response is None:
+                        continue
+
+                    nvtx_mark_debug(
+                        f"RPC.response.{'streaming' if response.is_streaming else 'sync'}",
+                        color="black",
+                        category="RPC")
 
-            except asyncio.CancelledError:
-                # Still handle cancellation for backward compatibility
-                logger_debug("Response reader cancelled")
-                break
-            except Exception as e:
-                await self._handle_reader_exception(e)
-                break
+                    # Optimize: Check debug flag before expensive string operations
+                    # This avoids holding GIL for f-string evaluation when debug is disabled
+                    if enable_llmapi_debug() or logger.level == 'debug':
+                        logger_debug(
+                            f"RPC Client received response: request_id={response.request_id}, "
+                            f"is_streaming={response.is_streaming}, "
+                            f"pending_futures={len(self._pending_futures)}")
+
+                    with nvtx_range_debug("handle_response",
+                                          color="purple",
+                                          category="RPC"):
+                        if response.is_streaming:
+                            self._handle_streaming_response(response)
+                        else:
+                            self._handle_regular_response(response)
+
+                except asyncio.CancelledError:
+                    # Still handle cancellation for backward compatibility
+                    logger_debug("Response reader cancelled")
+                    break
+                except Exception as e:
+                    await self._handle_reader_exception(e)
+                    break
 
         logger_debug("Response reader exiting gracefully")
         self._reader_task = None
@@ -310,9 +319,8 @@ async def _call_async(self, method_name, *args, **kwargs):
         Returns:
             The result of the remote method call
         """
-        logger_debug(
-            f"RPC client calling method: {method_name} with args: {args} and kwargs: {kwargs}"
-        )
+        if enable_llmapi_debug() or logger.level == 'debug':
+            logger_debug(f"RPC client calling method: {method_name}")
         nvtx_mark_debug(f"RPC.async.{method_name}",
                         color="yellow",
                         category="RPC")
@@ -331,36 +339,24 @@ async def _call_async(self, method_name, *args, **kwargs):
                              kwargs,
                              need_response,
                              timeout=timeout)
-        logger_debug(f"RPC client sending request: {request}")
         await self._client_socket.put_async(request)
 
         if not need_response:
             return None
 
         loop = asyncio.get_running_loop()
         future = loop.create_future()
-        logger_debug(
-            f"RPC Client _call_async: Created future for request_id: {request_id} in loop: {id(loop)}"
-        )
         self._pending_futures[request_id] = (future, loop)
-        logger_debug(
-            f"RPC Client _call_async: Stored future in pending_futures")
 
         try:
             # If timeout, the remote call should return a timeout error timely,
             # so we add 1 second to the timeout to ensure the client can get
             # that result.
-            logger_debug(
-                f"RPC Client _call_async: Awaiting future for request_id: {request_id}"
-            )
             if timeout is None:
                 res = await future
             else:
                 # Add 1 second to the timeout to ensure the client can get
                 res = await asyncio.wait_for(future, timeout)
-            logger_debug(
-                f"RPC Client _call_async: Got result for request_id: {request_id}: {res}"
-            )
             return res
         except RPCCancelled:
             self._server_stopped = True
@@ -394,22 +390,15 @@ def run_loop():
 
     def _call_sync(self, method_name, *args, **kwargs):
         """Synchronous version of RPC call."""
-        logger_debug(
-            f"RPC Client calling method: {method_name} with args: {args} and kwargs: {kwargs}"
-        )
+        if enable_llmapi_debug() or logger.level == 'debug':
+            logger_debug(f"RPC Client calling method: {method_name}")
         nvtx_mark_debug(f"RPC.sync.{method_name}",
                         color="green",
                         category="RPC")
         self._ensure_event_loop()
-        logger_debug(
-            f"RPC Client _call_sync: Creating future for {method_name}")
         future = asyncio.run_coroutine_threadsafe(
             self._call_async(method_name, *args, **kwargs), self._loop)
-        logger_debug(
-            f"RPC Client _call_sync: Waiting for result of {method_name}")
         result = future.result()
-        logger_debug(
-            f"RPC Client _call_sync: Got result for {method_name}: {result}")
         return result
 
     def _call_future(self, name: str, *args,
@@ -478,24 +467,21 @@ async def _call_streaming(self, name: str, *args,
 
             # Read streaming responses
             while True:
-                logger_debug(f"RPC Client _call_streaming waiting for response",
-                             color="green")
                 if timeout is None:
                     response = await queue.get()
                 else:
                     response = await asyncio.wait_for(queue.get(),
                                                       timeout=timeout)
 
-                logger_debug(
-                    f"RPC Client _call_streaming received [{response.stream_status}] response: {response}",
-                    color="green")
+                if enable_llmapi_debug() or logger.level == 'debug':
+                    logger_debug(
+                        f"RPC Client _call_streaming received [{response.stream_status}] response",
+                        color="green")
+
                 if response.stream_status == 'start':
                     # Start of stream
                     continue
                 elif response.stream_status == 'data':
-                    logger_debug(
-                        f"RPC Client _call_streaming received data: {response.result}",
-                        color="green")
                     yield response.result
                 elif response.stream_status == 'end':
                     # End of stream