feat: gc.freeze hot processes and disable gc during cudagraph capture

lightllm/common/basemodel/basemodel.py

@@ -24,6 +24,7 @@
 from lightllm.common.quantization import Quantcfg
 from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.gc_utils import freeze_gc, gc_frozen_and_disabled
 from lightllm.utils.dist_utils import get_dp_world_size
 from lightllm.utils.envs_utils import get_env_start_args, get_llm_data_type, get_added_mtp_kv_layer_num
 from lightllm.distributed.communication_op import dist_group_manager
@@ -127,11 +128,13 @@ def __init__(self, kvargs):
 
         self._autotune_warmup()
         self._init_padded_req()
-        self._init_cudagraph()
-        self._init_prefill_cuda_graph()
+        with gc_frozen_and_disabled("cudagraph-capture"):
+            self._init_cudagraph()
+            self._init_prefill_cuda_graph()
         self._check_max_len_infer()
         torch.cuda.empty_cache()
         set_model_init_status(True)
+        freeze_gc("model-infer-worker")
         return
 
     def _init_config(self):
@@ -604,7 +607,6 @@ def _decode(
 
     @final
     def _context_forward(self, infer_state: InferStateInfo):
-
         input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight)
         if self.args.enable_dp_prefill_balance:
             assert not self.args.enable_prefill_cudagraph, "not support now"

lightllm/server/api_http.py

@@ -47,6 +47,7 @@
 from .api_lightllm import lightllm_get_score
 from lightllm.utils.envs_utils import get_env_start_args, get_lightllm_websocket_max_message_size
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.gc_utils import freeze_gc
 from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError
 from lightllm.server.metrics.manager import MetricClient
 from lightllm.utils.envs_utils import get_unique_server_name
@@ -497,4 +498,5 @@ async def startup_event():
     g_objs.set_args(get_env_start_args())
     loop.create_task(g_objs.httpserver_manager.handle_loop())
     logger.info(f"server start up ok, loop use is {asyncio.get_event_loop()}")
+    freeze_gc("httpserver")
     return

lightllm/server/detokenization/manager.py

@@ -16,6 +16,7 @@
 import pickle
 import time
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.gc_utils import freeze_gc
 from lightllm.utils.envs_utils import get_unique_server_name
 
 logger = init_logger(__name__)
@@ -179,5 +180,6 @@ def start_detokenization_process(args, pipe_writer):
         pipe_writer.send(str(e))
         raise
     pipe_writer.send("init ok")
+    freeze_gc("detokenization")
     manager.handle_loop()
     return

lightllm/server/router/manager.py

@@ -26,6 +26,7 @@
 from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient
 from lightllm.server.core.objs.shm_objs_io_buffer import ShmObjsIOBuffer
 from lightllm.utils.log_utils import init_logger, log_time_ready
+from lightllm.utils.gc_utils import freeze_gc
 from lightllm.utils.profiler import ProfilerCmd
 from lightllm.server.router.token_load import TokenLoad
 from lightllm.server.metrics.manager import MetricClient
@@ -569,5 +570,6 @@ def handle_exception(loop, context):
         raise
 
     pipe_writer.send("init ok")
+    freeze_gc("router")
     loop.run_until_complete(router.loop_for_fwd())
     return

lightllm/utils/gc_utils.py

@@ -0,0 +1,23 @@
+import gc
+from contextlib import contextmanager
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def freeze_gc(tag: str = "") -> None:
+    gc.collect()
+    gc.freeze()
+    logger.info(f"gc.freeze done ({tag}): frozen={gc.get_freeze_count()}")
+
+
+@contextmanager
+def gc_frozen_and_disabled(tag: str = ""):
+    freeze_gc(tag)
+    was_enabled = gc.isenabled()
+    gc.disable()
+    try:
+        yield
+    finally:
+        if was_enabled:
+            gc.enable()