[PD] Graceful shutdown for disaggregation with RDMA cleanup

python/sglang/srt/disaggregation/common/conn.py

@@ -121,9 +121,9 @@ def __init__(
         )
 
         # bind zmq socket
-        context = zmq.Context()
+        self._zmq_context = zmq.Context()
         self.rank_port, self.server_socket = get_zmq_socket_on_host(
-            context, zmq.PULL, host=self.local_ip
+            self._zmq_context, zmq.PULL, host=self.local_ip
         )
         logger.debug(f"kv manager bind to {self.local_ip}:{self.rank_port}")
 
@@ -175,6 +175,20 @@ def __init__(
                 f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
             )
 
+    def shutdown(self):
+        """Close ZMQ socket and context to unblock threads waiting on recv_multipart().
+
+        Subclasses should call super().shutdown() before performing backend-specific cleanup.
+        """
+        try:
+            self.server_socket.close(linger=0)
+        except Exception as e:
+            logger.warning(f"Failed to close ZMQ server socket: {e}")
+        try:
+            self._zmq_context.term()
+        except Exception as e:
+            logger.warning(f"Failed to terminate ZMQ context: {e}")
+
     def check_status(self, bootstrap_room: int) -> KVPoll:
         return self.request_status[bootstrap_room]
 

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -276,6 +276,47 @@ def register_buffer_to_engine(self):
                 self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
             )
 
+    def shutdown(self):
+        """Deregister RDMA memory and close ZMQ sockets for clean process exit.
+
+        Without explicit deregistration, Mooncake's C++ TransferEngine destructor
+        must tear down RDMA resources during process exit. In Kubernetes, this can
+        cause pod sandbox teardown to hang with FailedKillPod errors.
+        """
+        logger.info("MooncakeKVManager shutting down...")
+
+        # 1. Close ZMQ sockets to unblock threads waiting on recv_multipart()
+        super().shutdown()
+
+        # 2. Shut down thread pool executors first to stop in-flight transfers
+        #    before deregistering the memory they may be reading from.
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            for executor in getattr(self, "executors", []):
+                try:
+                    executor.shutdown(wait=False)
+                except Exception as e:
+                    logger.warning(f"Failed to shut down thread pool executor: {e}")
+
+        # 3. Deregister RDMA memory regions (mirror registration conditions)
+        if self.engine is not None:
+            try:
+                if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
+                    self.engine.batch_deregister(self.kv_args.kv_data_ptrs)
+            except Exception as e:
+                logger.warning(f"Failed to deregister KV data buffers: {e}")
+            try:
+                if self.kv_args.aux_data_ptrs and self.kv_args.aux_data_lens:
+                    self.engine.batch_deregister(self.kv_args.aux_data_ptrs)
+            except Exception as e:
+                logger.warning(f"Failed to deregister aux data buffers: {e}")
+            try:
+                if self.kv_args.state_data_ptrs and self.kv_args.state_data_lens:
+                    self.engine.batch_deregister(self.kv_args.state_data_ptrs)
+            except Exception as e:
+                logger.warning(f"Failed to deregister state data buffers: {e}")
+
+        logger.info("MooncakeKVManager shutdown complete.")
+
     # ------------------------------------------------------------------
     # Staging buffer methods (all delegate to staging_handler.py)
     # ------------------------------------------------------------------
@@ -1420,7 +1461,7 @@ def bootstrap_thread():
                     if len(self.transfer_infos[room]) == required_dst_info_num:
                         self.update_status(room, KVPoll.WaitingForInput)
 
-        threading.Thread(target=bootstrap_thread).start()
+        threading.Thread(target=bootstrap_thread, daemon=True).start()
 
     def start_decode_thread(self):
         def decode_thread():
@@ -1548,8 +1589,8 @@ def heartbeat_checker():
                             if bootstrap_addr in self.session_pool:
                                 del self.session_pool[bootstrap_addr]
 
-        threading.Thread(target=decode_thread).start()
-        threading.Thread(target=heartbeat_checker).start()
+        threading.Thread(target=decode_thread, daemon=True).start()
+        threading.Thread(target=heartbeat_checker, daemon=True).start()
 
     def add_transfer_request(
         self,

python/sglang/srt/disaggregation/nixl/conn.py

@@ -882,7 +882,26 @@ def bootstrap_thread():
                     logger.debug(f"{room=} is bootstrapped")
                     self.update_status(room, KVPoll.WaitingForInput)
 
-        threading.Thread(target=bootstrap_thread).start()
+        threading.Thread(target=bootstrap_thread, daemon=True).start()
+
+    def shutdown(self):
+        """Deregister NIXL RDMA memory and close ZMQ sockets for clean process exit."""
+        logger.info("NixlKVManager shutting down...")
+
+        # 1. Close ZMQ sockets to unblock threads waiting on recv_multipart()
+        super().shutdown()
+
+        # 2. Deregister RDMA memory regions via NIXL agent
+        if hasattr(self, "agent") and self.agent is not None:
+            for desc_name in ("kv_descs", "aux_descs", "state_descs"):
+                descs = getattr(self, desc_name, None)
+                if descs is not None:
+                    try:
+                        self.agent.deregister_memory(descs)
+                    except Exception as e:
+                        logger.warning(f"Failed to deregister {desc_name}: {e}")
+
+        logger.info("NixlKVManager shutdown complete.")
 
 
 class NixlKVSender(CommonKVSender):

python/sglang/srt/managers/scheduler.py

@@ -13,6 +13,7 @@
 # ==============================================================================
 """A scheduler that manages a tensor parallel GPU worker."""
 
+import atexit
 import faulthandler
 import logging
 import os
@@ -3609,6 +3610,35 @@ def run_scheduler_process(
         # Send initialization info back to the parent process
         pipe_writer.send(scheduler.get_init_info())
 
+        # Register RDMA cleanup for disaggregation modes only.
+        # This ensures RDMA memory is deregistered and ZMQ sockets are closed
+        # when the scheduler process exits, preventing Kubernetes pod sandbox
+        # teardown hangs (FailedKillPod).
+        if scheduler.disaggregation_mode != DisaggregationMode.NULL:
+
+            def _shutdown_kv_manager():
+                try:
+                    kv_mgr = None
+                    if hasattr(scheduler, "disagg_prefill_bootstrap_queue"):
+                        kv_mgr = scheduler.disagg_prefill_bootstrap_queue.kv_manager
+                    elif hasattr(scheduler, "disagg_decode_prealloc_queue"):
+                        kv_mgr = scheduler.disagg_decode_prealloc_queue.kv_manager
+                    if kv_mgr is not None and hasattr(kv_mgr, "shutdown"):
+                        kv_mgr.shutdown()
+                except Exception as e:
+                    logger.warning(f"Error during KV manager shutdown: {e}")
+
+            atexit.register(_shutdown_kv_manager)
+
+            # Register SIGTERM handler so atexit handlers run on graceful shutdown.
+            # Without this, SIGTERM uses the default handler (immediate exit) which
+            # skips atexit/destructors, leaving RDMA resources unreleased.
+            def _sigterm_handler(signum, frame):
+                logger.info("Scheduler received SIGTERM, exiting gracefully...")
+                sys.exit(143)  # 128 + SIGTERM(15)
+
+            signal.signal(signal.SIGTERM, _sigterm_handler)
+
         # Run the event loop (blocks until shutdown)
         scheduler.run_event_loop()
 

python/sglang/srt/managers/tokenizer_manager.py

@@ -103,6 +103,7 @@
     freeze_gc,
     get_bool_env_var,
     get_or_create_event_loop,
+    graceful_kill_process_tree,
     kill_process_tree,
 )
 from sglang.srt.utils.aio_rwlock import RWLock
@@ -122,6 +123,10 @@
 
 logger = logging.getLogger(__name__)
 
+# Constants for graceful shutdown configuration
+CHILD_PROCESS_SHUTDOWN_TIMEOUT_ENV = "SGLANG_CHILD_PROCESS_SHUTDOWN_TIMEOUT"
+DEFAULT_CHILD_PROCESS_SHUTDOWN_TIMEOUT = 10.0
+
 _INCREMENTAL_STREAMING_META_INFO_KEYS = (
     "output_token_logprobs",
     "output_top_logprobs",
@@ -2199,8 +2204,22 @@ async def sigterm_watchdog(self):
                 self.dump_requests_before_crash()
                 break
 
-        kill_process_tree(os.getpid(), include_parent=True)
-        sys.exit(0)
+        # Gracefully terminate child processes (e.g., scheduler, detokenizer)
+        # by first sending SIGTERM to allow cleanup (RDMA deregistration, etc.),
+        # then SIGKILL after timeout if they don't exit.
+        shutdown_timeout = float(
+            os.environ.get(
+                CHILD_PROCESS_SHUTDOWN_TIMEOUT_ENV,
+                DEFAULT_CHILD_PROCESS_SHUTDOWN_TIMEOUT,
+            )
+        )
+        graceful_kill_process_tree(
+            os.getpid(), include_parent=False, timeout=shutdown_timeout
+        )
+        # Use os._exit() instead of sys.exit() to avoid SystemExit exception
+        # being caught by asyncio event loop, which would interrupt FastAPI's
+        # lifespan cleanup and prevent graceful shutdown of uvicorn server.
+        os._exit(0)
 
     def force_exit_handler(self):
         """Put some custom force exit logic here."""

python/sglang/srt/utils/common.py

@@ -1070,6 +1070,89 @@ def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = N
             pass
 
 
+def graceful_kill_process_tree(
+    parent_pid=None,
+    include_parent: bool = False,
+    skip_pid: int = None,
+    timeout: float = 10.0,
+):
+    """Gracefully kill process tree: SIGTERM first, then SIGKILL after timeout.
+
+    Useful for processes like Mooncake/NIXL that need time to deregister RDMA
+    memory and run atexit/C++ destructors when receiving SIGTERM.
+
+    Args:
+        parent_pid: The parent process ID. If None, uses current process.
+        include_parent: Whether to kill the parent process as well.
+        skip_pid: Process ID to skip.
+        timeout: Time in seconds to wait for graceful shutdown before SIGKILL.
+    """
+    # Remove sigchld handler to avoid spammy logs.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGCHLD, signal.SIG_DFL)
+
+    if parent_pid is None:
+        parent_pid = os.getpid()
+        include_parent = False
+
+    try:
+        itself = psutil.Process(parent_pid)
+    except psutil.NoSuchProcess:
+        return
+
+    children = itself.children(recursive=True)
+    if not children:
+        return
+
+    # Step 1: Send SIGTERM to all children for graceful shutdown
+    alive_children = []
+    for child in children:
+        if child.pid == skip_pid:
+            continue
+        try:
+            logger.info(
+                f"Sending SIGTERM to child process {child.pid} ({child.name()})"
+            )
+            child.terminate()  # Send SIGTERM
+            alive_children.append(child)
+        except psutil.NoSuchProcess:
+            pass
+
+    # Step 2: Wait for processes to terminate gracefully
+    if alive_children:
+        logger.info(
+            f"Waiting up to {timeout}s for {len(alive_children)} "
+            "child processes to terminate gracefully..."
+        )
+        gone, alive = psutil.wait_procs(alive_children, timeout=timeout)
+
+        if gone:
+            logger.info(f"{len(gone)} child processes terminated gracefully")
+
+        # Step 3: Force kill any remaining processes
+        if alive:
+            logger.warning(
+                f"{len(alive)} child processes did not terminate gracefully, "
+                "sending SIGKILL"
+            )
+            for child in alive:
+                try:
+                    logger.info(f"Sending SIGKILL to child process {child.pid}")
+                    child.kill()
+                except psutil.NoSuchProcess:
+                    pass
+
+            # Wait a bit for SIGKILL to take effect
+            psutil.wait_procs(alive, timeout=3)
+
+    # Handle parent process if requested
+    if include_parent:
+        try:
+            itself.kill()
+        except psutil.NoSuchProcess:
+            pass
+
+
 def monkey_patch_p2p_access_check():
     """
     Monkey patch the slow p2p access check.