From e6722c751b0f5d846c779a1eb2dea5ba6133b411 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 21 Jun 2026 15:08:10 -0700 Subject: [PATCH] [Feature] Add graceful scheduler shutdown; free hisparse host buffer on exit (#28779) --- .../srt/managers/hisparse_coordinator.py | 7 +++++ python/sglang/srt/managers/io_struct.py | 7 +++++ python/sglang/srt/managers/scheduler.py | 28 +++++++++++++++++- .../sglang/srt/managers/tokenizer_manager.py | 10 +++++++ python/sglang/srt/mem_cache/pool_host/base.py | 29 ++++++++++++++++++- .../sglang/srt/mem_cache/pool_host/common.py | 13 +++++++++ .../models_e2e/test_dsa_glm5_hisparse.py | 15 ++++++++++ 7 files changed, 107 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index cc1ae09f7..b6db923e9 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -189,6 +189,13 @@ class HiSparseCoordinator: def set_decode_producer_stream(self, stream) -> None: self.decode_producer_stream = stream + def destroy(self) -> None: + # Drain in-flight transfers so the buffer is idle, then unregister it. + # See HostKVCache.destroy for why the explicit unregister matters. + self.write_staging_stream.synchronize() + self.decode_backup_stream.synchronize() + self.mem_pool_host.destroy() + def get_token_stats(self) -> HiSparseTokenStats: device_allocator = self.token_to_kv_pool_allocator.hisparse_attn_allocator device_capacity = device_allocator.size diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 074ab43a2..951f35495 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1814,6 +1814,13 @@ class FreezeGCReq(BaseReq): pass +@dataclass +class ShutdownReq(BaseReq): + # Broadcast across TP ranks via the normal recv path, so all ranks break + # the scheduler loop on the same iteration. + pass + + @dataclass class ConfigureLoggingReq(BaseReq): log_requests: Optional[bool] = None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 6d546ee41..d32d44cbc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -135,6 +135,7 @@ from sglang.srt.managers.io_struct import ( SendWeightsToRemoteInstanceReqOutput, SetInternalStateReq, SetInternalStateReqOutput, + ShutdownReq, SlowDownReqInput, SlowDownReqOutput, TokenizedEmbeddingReqInput, @@ -352,6 +353,9 @@ class Scheduler( self.enable_hisparse = server_args.enable_hisparse self.hisparse_coordinator: Optional[HiSparseCoordinator] = None + # Set by the ShutdownReq handler to break the event loop for graceful shutdown. + self.gracefully_exit = False + # Distributed rank info attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = ( compute_dp_attention_world_info( @@ -1389,6 +1393,7 @@ class Scheduler( lambda req: self.profiler_manager._profile(req), ), (FreezeGCReq, self.handle_freeze_gc), + (ShutdownReq, self.handle_shutdown), (GetInternalStateReq, self.get_internal_state), (SetInternalStateReq, self.set_internal_state), (RpcReqInput, self.handle_rpc_request), @@ -1445,6 +1450,12 @@ class Scheduler( return result_dict + def release_host_resources(self) -> None: + # Release pinned host buffers in userspace on graceful shutdown; see + # HostKVCache.destroy. Called from run_scheduler_process's finally. + if self.hisparse_coordinator is not None: + self.hisparse_coordinator.destroy() + def run_event_loop(self) -> None: """Run the scheduler's event loop. @@ -1473,6 +1484,9 @@ class Scheduler( def event_loop_normal(self): """A normal scheduler loop.""" while True: + if self.gracefully_exit: + break + # Receive requests recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) @@ -1509,6 +1523,9 @@ class Scheduler( self.process_batch_result(tmp_batch, tmp_result) while True: + if self.gracefully_exit: + break + # Receive requests recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) @@ -4002,6 +4019,11 @@ class Scheduler( self.ipc_channels.send_to_detokenizer.send_output(recv_req, recv_req) return None + def handle_shutdown(self, recv_req: ShutdownReq): + # Break the event loop; the finally in run_scheduler_process releases resources. + self.gracefully_exit = True + return None + def configure_logging(self, recv_req: ConfigureLoggingReq): if recv_req.log_level is not None: logging.getLogger().setLevel(recv_req.log_level.upper()) @@ -4182,7 +4204,7 @@ def run_scheduler_process( # Send initialization info back to the parent process pipe_writer.send(scheduler.get_init_info()) - # Run the event loop (blocks until shutdown) + # Run the event loop (blocks until a ShutdownReq sets gracefully_exit) scheduler.run_event_loop() except Exception: @@ -4201,3 +4223,7 @@ def run_scheduler_process( # FPM has a background ZMQ publisher thread that needs explicit # teardown to flush queued metrics and close the socket cleanly. scheduler.metrics_reporter._shutdown_fpm() + # Graceful path only: on the exception path the GPU may be wedged + # and the synchronize() in destroy() could itself hang. + if scheduler.gracefully_exit: + scheduler.release_host_resources() diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index f9e6cb932..bf932611a 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -70,6 +70,7 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqOutput, PauseGenerationReqInput, SessionParams, + ShutdownReq, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, UpdateWeightFromDiskReqInput, @@ -2642,6 +2643,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): else: break + # Stop the watchdog: child exits are expected during shutdown, not crashes. + if self._subprocess_watchdog is not None: + self._subprocess_watchdog.stop() + # Ask schedulers to release resources in userspace and exit (see + # ShutdownReq), then wait for them before hard-killing the rest. + self.send_to_scheduler.send_pyobj(ShutdownReq()) + deadline = time.monotonic() + 15 + while time.monotonic() < deadline and collect_scheduler_processes(): + time.sleep(0.1) kill_process_tree(os.getpid(), include_parent=True) sys.exit(0) diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index bafe065db..0329fc7ca 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -10,10 +10,17 @@ import psutil import torch from sglang.srt.mem_cache.memory_pool import KVCache -from sglang.srt.mem_cache.pool_host.common import get_allocator_from_storage +from sglang.srt.mem_cache.pool_host.common import ( + _cuda_host_unregister, + get_allocator_from_storage, +) +from sglang.srt.utils import is_cuda, is_hip logger = logging.getLogger(__name__) +_is_cuda = is_cuda() +_is_hip = is_hip() + # Host RAM to leave free when sizing HiCache pools (OS, other processes). HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3) @@ -86,6 +93,26 @@ class HostKVCache(abc.ABC): self.lock = threading.RLock() self.clear() + def destroy(self): + """Unregister pinned host buffers in userspace before process exit. + + Large cudaHostRegister'd buffers are otherwise unpinned by the kernel + during SIGKILL reclaim, which can stall teardown in uninterruptible + sleep for tens of seconds. Idempotent. (Only the host_register path + needs this; npu/musa pin_memory buffers are freed by torch.) + """ + if getattr(self, "_destroyed", False): + return + self._destroyed = True + buffers = getattr(self, "kv_buffer", None) + if buffers is not None and self.pin_memory and (_is_cuda or _is_hip): + if not isinstance(buffers, (list, tuple)): + buffers = [buffers] + for buf in buffers: + if buf is not None: + _cuda_host_unregister(buf) + self.kv_buffer = None + @abc.abstractmethod def get_size_per_token(self): raise NotImplementedError() diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index 3e90b3d24..db17af43a 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/python/sglang/srt/mem_cache/pool_host/common.py @@ -57,6 +57,19 @@ def _cuda_host_register(buffer: torch.Tensor) -> None: ) +def _cuda_host_unregister(buffer: torch.Tensor) -> None: + cudart = torch.cuda.cudart() + rc = cudart.cudaHostUnregister(buffer.data_ptr()) + if int(rc) != 0: + # Best-effort on shutdown: warn, don't raise -- a leak is reclaimed at exit. + logger.warning( + "cudaHostUnregister failed (rc=%d, %s) for ptr=%#x", + int(rc), + cudart.cudaGetErrorString(rc), + buffer.data_ptr(), + ) + + def alloc_with_host_register( dims: tuple, dtype: torch.dtype, diff --git a/test/registered/models_e2e/test_dsa_glm5_hisparse.py b/test/registered/models_e2e/test_dsa_glm5_hisparse.py index 2e5c3c819..a846d11ae 100644 --- a/test/registered/models_e2e/test_dsa_glm5_hisparse.py +++ b/test/registered/models_e2e/test_dsa_glm5_hisparse.py @@ -1,5 +1,8 @@ +import subprocess +import time import unittest +from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.eval_accuracy_kit import GSM8KMixin from sglang.test.server_fixtures.default_fixture import DefaultServerBase @@ -49,6 +52,18 @@ class TestGLM5HiSparse(DefaultServerBase, GSM8KMixin): gsm8k_num_threads = 100 gsm8k_num_shots = 24 + @classmethod + def tearDownClass(cls): + # HiSparse's large pinned host buffer stalls an external SIGKILL teardown + # (kernel unpin). Drive the server's own graceful shutdown so each rank + # unregisters in userspace; hard-kill as a fallback. + cls.process.terminate() + try: + cls.process.wait(timeout=90) + except subprocess.TimeoutExpired: + kill_process_tree(cls.process.pid, wait_timeout=60) + time.sleep(2) + if __name__ == "__main__": unittest.main()