[None][feat] support calling gc.freeze

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -20,7 +20,7 @@
 except ImportError:
     from cuda import cudart
 
-from tensorrt_llm._utils import (customized_gc_thresholds, is_trace_enabled,
+from tensorrt_llm._utils import (customized_gc_configuration, is_trace_enabled,
                                  mpi_comm, mpi_disabled, nvtx_range,
                                  set_thread_local_mpi_comm, trace_func)
 from tensorrt_llm.bindings.executor import (DisServingRequestStats,
@@ -296,6 +296,7 @@ def __init__(
             kv_cache_transceiver: Optional[KvCacheTransceiver] = None,
             guided_decoder: Optional[GuidedDecoder] = None,
             garbage_collection_gen0_threshold: Optional[int] = None,
+            garbage_collection_freeze_after_init: bool = False,
             start_worker: bool = True,
             kv_connector_manager: Optional[KvCacheConnectorManager] = None,
             max_seq_len: Optional[int] = None,
@@ -596,6 +597,7 @@ def on_detected():
                     "Drafting is not supported for selected executor loop. "
                     "Please disable disagg/pipeline parallelism scheduler.")
         self.garbage_collection_gen0_threshold = garbage_collection_gen0_threshold
+        self.garbage_collection_freeze_after_init = garbage_collection_freeze_after_init
         self.max_seq_len = max_seq_len
 
         self.worker_started = False
@@ -693,7 +695,7 @@ def _event_loop_wrapper(self):
             enable_profiler = bool(os.environ.get(
                 "TLLM_LINE_PROFILER_PATH")) and not self.is_warmup
             with host_profiler_context(enable=enable_profiler), \
-                 customized_gc_thresholds(self.garbage_collection_gen0_threshold):
+                 customized_gc_configuration(gen0_threshold=self.garbage_collection_gen0_threshold, freeze=self.garbage_collection_freeze_after_init):
                 self.event_loop()
         except Exception as e:
             logger.error(f"Error in event loop: {e}")

tensorrt_llm/_utils.py

@@ -942,11 +942,31 @@ def customized_gc_thresholds(gen0_threshold: Optional[int] = None):
             )
 
 
+@contextmanager
+def maybe_gc_freeze(freeze: bool):
+    try:
+        if freeze:
+            gc.collect(2)
+            gc.freeze()
+        yield
+    finally:
+        if freeze:
+            gc.unfreeze()
+
+
 @contextmanager
 def _null_context_manager():
     yield
 
 
+@contextmanager
+def customized_gc_configuration(*,
+                                gen0_threshold: Optional[int] = None,
+                                freeze: bool = False):
+    with maybe_gc_freeze(freeze), customized_gc_thresholds(gen0_threshold):
+        yield
+
+
 _T = TypeVar("_T")
 
 

tensorrt_llm/commands/serve.py

@@ -348,6 +348,13 @@ def launch_server(
         if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1":
             gc.disable()
 
+        # Optionally freeze GC (default: no freeze) to reduce latency spikes due to repeated attempts
+        # at garbage collection of long-lived resources, at the expense of incomplete resource
+        # reclamantion.
+        if os.getenv("TRTLLM_SERVER_FREEZE_GC", "0") == "1":
+            gc.collect(2)
+            gc.freeze()
+
         asyncio.run(server(host, port, sockets=[s]))
 
 

tensorrt_llm/executor/proxy.py

@@ -13,7 +13,7 @@
 
 from tensorrt_llm.logger import logger
 
-from .._utils import customized_gc_thresholds, mpi_rank, nvtx_range_debug
+from .._utils import customized_gc_configuration, mpi_rank, nvtx_range_debug
 from ..llmapi.mpi_session import (MpiCommSession, MpiPoolSession, MpiSession,
                                   RemoteMpiCommSessionClient)
 from ..llmapi.tracer import enable_llm_tracer, get_tracer, global_tracer
@@ -89,9 +89,9 @@ def __init__(
 
         self.model_world_size = model_world_size
 
-        self.garbage_collection_gen0_threshold = worker_kwargs[
-            "llm_args"].garbage_collection_gen0_threshold if worker_kwargs.get(
-                "llm_args", None) is not None else None
+        llm_args = worker_kwargs.get("llm_args")
+        self.garbage_collection_gen0_threshold = llm_args.garbage_collection_gen0_threshold if llm_args is not None else None
+        self.garbage_collection_freeze_after_init: bool = llm_args.garbage_collection_freeze_after_init if llm_args is not None else False
 
         # Generate RPC address and key for stats RPC
         self.rpc_addr = get_unique_ipc_addr()
@@ -269,9 +269,8 @@ def abort_request(self, request_id: int) -> None:
     def dispatch_result_task(self) -> bool:
         # TODO[chunweiy]: convert the dispatch_result_task to async, that should
         # benefit from zmq.asyncio.Context
-        with customized_gc_thresholds(self.garbage_collection_gen0_threshold):
-            if (res := self.result_queue.get()) is None:
-                return False  # shutdown the thread
+        if (res := self.result_queue.get()) is None:
+            return False  # shutdown the thread
 
         async_queues = []
         event_loop = None
@@ -325,7 +324,12 @@ def _start_dispatch_threads(self):
             self.dispatch_result_thread = ManagedThread(
                 weakref.WeakMethod(self.dispatch_result_task),
                 error_queue=self._error_queue,
-                name="proxy_dispatch_result_thread")
+                name="proxy_dispatch_result_thread",
+                context=customized_gc_configuration(
+                    gen0_threshold=self.garbage_collection_gen0_threshold,
+                    freeze=self.garbage_collection_freeze_after_init,
+                ),
+            )
 
             self.dispatch_result_thread.start()
 

tensorrt_llm/executor/rpc/rpc_client.py

@@ -9,7 +9,7 @@
 
 import zmq
 
-from tensorrt_llm._utils import (customized_gc_thresholds, nvtx_mark_debug,
+from tensorrt_llm._utils import (customized_gc_configuration, nvtx_mark_debug,
                                  nvtx_range_debug)
 
 from ...llmapi.utils import (AsyncQueue, _SyncQueue, enable_llmapi_debug,
@@ -345,7 +345,7 @@ async def _response_reader(self):
             await asyncio.sleep(0.1)
             logger_debug("[client] Response reader ready to process messages")
 
-            with customized_gc_thresholds(10000):
+            with customized_gc_configuration(gen0_threshold=10000):
                 last_alive_log = time.time()
                 while not self._closed:
                     # Periodic alive logging for debugging

tensorrt_llm/llmapi/llm_args.py

@@ -3715,6 +3715,13 @@ class TorchLlmArgs(BaseLlmArgs):
         "Lower values trigger more frequent garbage collection.",
         status="beta")
 
+    garbage_collection_freeze_after_init: int = Field(
+        default=False,
+        description=
+        "Call gc.freeze after engine initialization to prevent attempting garbage collection of long-lived " \
+        "executor resources. This can reduce latency spikes, at the expense of incomplete resource reclamation.",
+        status="beta")
+
     cuda_graph_config: Optional[CudaGraphConfig] = Field(
         default_factory=CudaGraphConfig,
         description="CUDA graph config. If true, use CUDA graphs for decoding. \

tensorrt_llm/llmapi/utils.py

@@ -15,11 +15,12 @@
 import traceback
 import warnings
 import weakref
+from contextlib import nullcontext
 from functools import wraps
 from pathlib import Path
 from queue import Queue
-from typing import (Any, Callable, Iterable, List, Optional, Tuple, Type,
-                    get_type_hints)
+from typing import (Any, Callable, ContextManager, Iterable, List, Optional,
+                    Tuple, Type, get_type_hints)
 
 import filelock
 import huggingface_hub
@@ -171,19 +172,6 @@ def get_gpu_arch(device: int = 0) -> int:
     return torch.cuda.get_device_properties(device).major
 
 
-class ContextManager:
-    ''' A helper to create a context manager for a resource. '''
-
-    def __init__(self, resource):
-        self.resource = resource
-
-    def __enter__(self):
-        return self.resource.__enter__()
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        return self.resource.__exit__(exc_type, exc_value, traceback)
-
-
 def is_directory_empty(directory: Path) -> bool:
     return not any(directory.iterdir())
 
@@ -328,33 +316,35 @@ def __init__(self,
                  error_queue: Queue,
                  name: Optional[str] = None,
                  stop_event: Optional[threading.Event] = None,
+                 context: Optional[ContextManager[Any]] = None,
                  **kwargs):
         super().__init__(name=name)
         self.task = task
         self.error_queue = error_queue
         self.kwargs = kwargs
         self.daemon = True
         self.stop_event = stop_event or threading.Event()
+        self.context = context or nullcontext()
 
     def run(self):
-
-        while not self.stop_event.is_set():
-            task = self.task
-            if isinstance(task, weakref.WeakMethod):
-                task = task()
-                if task is None:
-                    # Normally, this should not happen.
-                    logger.warning("WeakMethod is expired.")
-                    break
-
-            try:
-                if not task(**self.kwargs):
-                    break
-            except Exception as e:
-                logger.error(
-                    f"Error in thread {self.name}: {e}\n{traceback.format_exc()}"
-                )
-                self.error_queue.put(e)
+        with self.context:
+            while not self.stop_event.is_set():
+                task = self.task
+                if isinstance(task, weakref.WeakMethod):
+                    task = task()
+                    if task is None:
+                        # Normally, this should not happen.
+                        logger.warning("WeakMethod is expired.")
+                        break
+
+                try:
+                    if not task(**self.kwargs):
+                        break
+                except Exception as e:
+                    logger.error(
+                        f"Error in thread {self.name}: {e}\n{traceback.format_exc()}"
+                    )
+                    self.error_queue.put(e)
 
         logger.info(f"Thread {self.name} stopped.")