make gc tracer singleton

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

Author: Superjomn

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -20,9 +20,8 @@
 
 from tensorrt_llm._torch.pyexecutor.resource_manager import (
     ResourceManagerType, request_context)
-from tensorrt_llm._utils import (customized_gc_thresholds, gc_nvtx_watcher,
-                                 is_trace_enabled, mpi_disabled, nvtx_range,
-                                 trace_func)
+from tensorrt_llm._utils import (customized_gc_thresholds, is_trace_enabled,
+                                 mpi_disabled, nvtx_range, trace_func)
 from tensorrt_llm.bindings.executor import (DisServingRequestStats,
                                             FinishReason, InflightBatchingStats,
                                             IterationStats, KvCacheStats,
@@ -139,7 +138,6 @@ 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()
 
         # related modules
         self.resource_manager = resource_manager

tensorrt_llm/_utils.py

@@ -1218,15 +1218,30 @@ class _GCNvtxHandle:
     """Handle object for GC NVTX watcher to keep it alive."""
 
 
-def gc_nvtx_watcher() -> Optional[_GCNvtxHandle]:
+# Singleton for the GC NVTX watcher handle.
+_gc_watcher_handle: Optional[_GCNvtxHandle] = None
+
+
+def _setup_gc_nvtx_profiling() -> Optional[_GCNvtxHandle]:
     """
-    Set up NVTX range markers for Python garbage collection events.
+    Set up NVTX range markers for Python garbage collection events (singleton).
     This helps in profiling to visualize when GC occurs during execution.
 
+    This function is called automatically at module import time. The environment
+    variable TLLM_PROFILE_RECORD_GC must be set before importing this module.
+
+    This is an internal function and should not be called directly by users.
+
     Returns:
         _GCNvtxHandle or None: A handle object that keeps the GC callback alive,
                                or None if GC profiling is not enabled.
     """
+    global _gc_watcher_handle
+
+    # Return existing handle if already initialized
+    if _gc_watcher_handle is not None:
+        return _gc_watcher_handle
+
     enabled = os.environ.get(PROFILE_RECORD_GC_ENV_VAR_NAME, None)
     if not enabled:
         return None
@@ -1253,4 +1268,10 @@ def gc_cleanup(callback):
 
     handle = _GCNvtxHandle()
     weakref.finalize(handle, gc_cleanup, gc_callback)
+
+    _gc_watcher_handle = handle
     return handle
+
+
+# Initialize GC NVTX profiling singleton at module import time
+_setup_gc_nvtx_profiling()

tensorrt_llm/llmapi/llm.py

@@ -21,7 +21,7 @@
 from tensorrt_llm.llmapi import tracing
 from tensorrt_llm.metrics.enums import MetricNames
 
-from .._utils import gc_nvtx_watcher, nvtx_range_debug
+from .._utils import nvtx_range_debug
 from ..bindings import executor as tllm
 from ..bindings import steady_clock_now
 from ..builder import EngineConfig
@@ -131,8 +131,6 @@ def __init__(self,
         self._orchestrator_type = kwargs.get("orchestrator_type", None)
         self._llm_id = None
 
-        self._gc_nvtx_watcher_handle = gc_nvtx_watcher()
-
         log_level = logger.level
         logger.set_level("info")  # force display the backend