[PD] Graceful shutdown for disaggregation with RDMA cleanup

[Kangyan-Zhou]

Summary

Combines #16484 and #19810 into a single PR for graceful shutdown of SGLang PD disaggregation pods.

When Kubernetes sends SIGTERM to SGLang pods, the sigterm_watchdog previously sent SIGKILL to scheduler child processes, bypassing all Python cleanup (atexit, __del__, C++ destructors). This left Mooncake/NIXL RDMA memory regions registered, causing pod sandbox teardown to hang with FailedKillPod errors.

Parent side (tokenizer_manager)

• Add graceful_kill_process_tree() that sends SIGTERM first, waits up to a configurable timeout (SGLANG_CHILD_PROCESS_SHUTDOWN_TIMEOUT, default 10s), then SIGKILL for stragglers
• Use os._exit(0) instead of sys.exit(0) to avoid SystemExit being caught by the asyncio event loop

Child side (scheduler + KV managers)

• Register SIGTERM handler in scheduler that calls sys.exit(143) so atexit handlers run
• Register atexit handler that calls kv_manager.shutdown() for RDMA cleanup
• Add CommonKVManager.shutdown() — closes ZMQ socket (linger=0) and terminates context, unblocking threads on recv_multipart()
• Add MooncakeKVManager.shutdown() — shuts down thread pool executors, then deregisters all RDMA memory (kv, aux, state buffers)
• Add NixlKVManager.shutdown() — deregisters RDMA memory via NIXL agent
• Mark bootstrap_thread, decode_thread, heartbeat_checker as daemon=True

Test plan

• Send SIGTERM to a running PD disaggregation deployment and verify RDMA memory is deregistered (check logs for "MooncakeKVManager shutdown complete" / "NixlKVManager shutdown complete")
• Verify pod terminates cleanly without FailedKillPod errors in Kubernetes
• Verify non-disaggregation mode is unaffected (no SIGTERM handler registered)
• Run existing disaggregation tests to ensure no regression

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a graceful shutdown mechanism for disaggregation components to ensure RDMA memory is deregistered and ZMQ sockets are closed, preventing Kubernetes pod teardown hangs. Key changes include adding shutdown methods to KV managers, implementing a graceful_kill_process_tree utility, and registering atexit and SIGTERM handlers in the scheduler. Feedback focuses on ensuring the idempotency of the shutdown process and preventing potential duplicate atexit registrations for the cleanup handler.

[gemini-code-assist]

The shutdown method should ideally be idempotent. If shutdown is called multiple times, the second call might raise an exception when attempting to close the socket or terminate the context again. Adding a check to ensure the socket and context are still valid or using a flag would make this more robust.

[gemini-code-assist]

The _shutdown_kv_manager function is defined inside the loop and registered with atexit. If run_scheduler_process is called multiple times or if the scheduler is re-initialized, this could lead to multiple atexit registrations. It is better to register the cleanup handler once or ensure it is cleaned up properly.