diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d4308445cb57..95fc389d655f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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}") diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index c53e7a085048..c74cdbad0bfd 100644 --- a/tensorrt_llm/_utils.py +++ b/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") diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 0d0a1e3b3da3..332419ce3a98 100644 --- a/tensorrt_llm/commands/serve.py +++ b/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])) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index f67b7cb2ccce..8da785d45a2f 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/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() diff --git a/tensorrt_llm/executor/rpc/rpc_client.py b/tensorrt_llm/executor/rpc/rpc_client.py index aaf9e8e2a64c..87f6c36a91ba 100644 --- a/tensorrt_llm/executor/rpc/rpc_client.py +++ b/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 diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index ca3af1b6a34f..ba9dd3ac72dd 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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. \ diff --git a/tensorrt_llm/llmapi/utils.py b/tensorrt_llm/llmapi/utils.py index 6765617a1804..ef501cd0894f 100644 --- a/tensorrt_llm/llmapi/utils.py +++ b/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,6 +316,7 @@ 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 @@ -335,26 +324,27 @@ def __init__(self, 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.")