fix(sglang): remove TP worker sockets left behind by SIGINT - #553

Summary

On SIGINT, python -m sglang.launch_server and sglang serve leave /tmp/kvcached-tp-<name>-<hash>/ behind with the schedulers' w<rank>.sock entries, which RixinLiu filed as #548 after the 2x L4 runs at tp size 2. This is the SGLang half of the socket finding from the #482 validation that became #510 and #543 for vLLM. #546's shutdown owner, which releases the control segments once every recorded child process is gone, now releases the deployment's worker sockets as well, through the WorkerSocketCleanup #543 added for vLLM, and waits, bounded, for a killed scheduler to release a socket that is still open in it.

Stacked on #546 at 61b97e5. Below it are ursasi's commits, the last two are this change, and the diff shrinks to those two once #546 merges. The branch is force-pushed on every re-stack.

Root cause

launch_server.py:86 and cli/serve.py:207 end with their own kill_process_tree(os.getpid(), include_parent=False), bound from sglang.srt.utils into their own namespace. It SIGKILLs every scheduler and waits until none holds resources. Exiting on its own, a scheduler runs the atexit hook that unlinks its socket, and that hook printed "IPC listener stopped" in the 0.5B runs of the issue. Killed before that hook, it exits with -9 and its socket node and directory stay. On SIGTERM the path sends ShutdownReq and waits before killing, and the hook usually runs there. #546 at 61b97e5 binds the launch_server, cli.serve, http_server and __main__ names itself. The SIGINT call reaches its owner there, and the utility wrap from the first round of this PR is gone.

shipiyouniao's review of the first round found the second half. SGLang's own wait at 0.5.20 returns at the zombie report without reaping, and so does the owner's exit check. psutil reports a killed scheduler as a zombie while the rest of its threads are still being torn down, and its listener accepts connections until the last of them has released the socket. So the final kill call of launch_server.py saw a served socket, kept it, and left the owner queued for a call that never comes. Reproduced on Linux with the spawn start method, 2 failed / 2 passed twice on his side and 5 of 5 runs here.

Changes

kvcached/integration/sglang/shutdown.py. Each scheduler's ready message now carries _kvcached_worker_sockets, the socket paths its IPC listeners bound (tp_ipc_util.worker_socket_paths()), next to the pool names, and the DP controller forwards both keys the same way. _capture_owner remembers the reported nodes while the schedulers are alive, as a WorkerSocketCleanup.for_paths entry among the owner's segments, and claims nothing outside the layout tp_ipc_util binds, w<rank>.sock in a kvcached-tp-* directory or one of its pp<k> stage directories. At TP=1 there is no listener and the scheduler does not report any. In the first round the parent derived the paths from the tp and pp sizes and the reported pool names, and that derivation is gone.

_ShutdownOwner.cleanup retries the entries that are not done yet, at the exit check's interval and until the same deadline, and the kill wrapper grants that deadline unless the call skips a child. Before, it granted it only when wait_timeout was None, and the upstream wait with a timeout returns at the zombie report as well. Identity and connect checks run on every attempt. Sockets a live worker serves are still kept and the next kill call retries the owner, a skipped live child keeps every file without a wait, and the "Keeping worker socket" line can appear once or twice before the "Removed" line while a killed scheduler lets go.

kvcached/utils.py adds WorkerSocketCleanup.for_paths(paths) and get_tp_socket_dir_of(path), the inverse of get_tp_worker_socket_path, so the deployment root and a stage directory still go once empty. Nothing changes for the vLLM constructor. kvcached/tp_ipc_util.py adds worker_socket_paths() over the listener registry.

Validation

tests/test_sglang_shutdown_sockets.py, 20 cases in the fake-module style of test_sglang_shutdown.py, covering the readiness report from the listener registry and without the extension, the SIGINT sequence through the sglang.launch_server binding, Engine.shutdown, pp stage sockets, the self-exit path removing the sockets before the parent's own kill, a skipped live scheduler, a served socket kept in bounded time with the owner retried, a same-name replacement left alone served or dead, reported nodes absent at readiness never claimed, repeated kill calls leaving a replacement directory, an empty report, four reported paths outside the layout claiming nothing, the DP forwarding of both keys, the socket released after the exit report removed by the one final call, the identity check during the wait, and kvcached disabled. Of these, the released-after-exit case fails at 88be58f (the first commit) and passes at the head.

tests/test_sglang_shutdown_sockets_process.py, 4 cases over spawn and fork, with real SIGKILLed children that serve the socket from a thread like the listener does, a kill that polls to the zombie report the way 0.5.20's wait does and never reaps, plus a survivor that keeps every socket until the next call. In a Linux container from the python:3.11 image (3.11.17, psutil 7.1.0, aarch64, kernel 6.8), 5 of 5 runs at 88be58f fail with Keeping worker socket ...: a worker still accepts connections on it, a spawn case in every run and a fork case once, and 10 of 10 runs at the head pass. macOS passes the 4 cases at the head as well. tests/test_tp_socket_cleanup.py gets two cases for for_paths and the inverse. Both new files are in the CPU manifest, following #546's move of test_sglang_shutdown_process.py there.

#546's own files pass unchanged, test_sglang_shutdown.py and test_sglang_shutdown_process.py among them. Full CPU manifest on macOS with Python 3.11, 1562 passed and 14 skipped, failure set byte-identical to #546's head at 61b97e5 with 1536 passed there, the known shm tracker, ipc timeout, listener device and pip target root cases. ruff 0.11.7, isort 6.0.1, codespell 2.4.1 clean. CI-exact mypy 1.11.1, 3.10 to 3.13, error set identical to #546's head on kvcached and tests. Classification check valid.

CPU only. No SGLang GPU run here. Still open is the issue's own protocol, python -m sglang.launch_server at tp size 2 on 0.5.20 with kvcached, serve a few requests, SIGINT the process group, wait for every process to exit, inspect /tmp/kvcached-tp-<name>-*, and the same for SIGTERM and for Engine.shutdown(). On main the directory stays for the models in the issue's table.

Fixes #548