From 7fecce4e490656eb5173a7ea5e318a47f1b67b50 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 05:52:18 -0700 Subject: [PATCH] Studio: cross-session backstop to reap a leftover llama-server on startup (#6431) * Reap Studio child processes when the parent dies abnormally Standalone `unsloth studio` launches orphaned cloudflared and llama-server when the parent exited without running the cooperative shutdown path (terminal-window close, Task Manager End Task, SIGKILL): the children reparented to init and kept running, leaving an authenticated Cloudflare tunnel up for days. Add utils/process_lifetime.py: a parent-owned Windows Job Object (JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, children auto-inherit) plus Linux PR_SET_PDEATHSIG, behind a best-effort helper that mirrors the desktop app's windows_job.rs. initialize_parent_lifetime() runs at the top of run_server; long-lived spawns (cloudflared, llama-server, RAG embedder, llama.cpp updater) get the PDEATHSIG preexec, multiprocessing workers are adopted into the job, and _graceful_shutdown plus atexit gain a terminate_all() backstop sweep. The cooperative shutdown path is otherwise unchanged. Verified on Linux: killing the parent now reaps cloudflared and llama-server within ~2s instead of orphaning them. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: add real Windows kill-on-job-close integration test Spawn a parent that installs the job and a child that inherits it, terminate the parent, and assert the child is reaped. Skipped off Windows. Also make the liveness probe Windows-safe (os.kill(pid, 0) terminates on Windows). * Fix Win64 handle truncation in the Job Object calls Set explicit argtypes so the 64-bit job/process handles are not marshaled as c_int (which truncated them on Win64, failing AssignProcessToJobObject). Assert install success in the Windows integration test. * Bind multiprocessing workers to parent death; harden the sweep Review follow-ups: - Multiprocessing workers (inference/export/training/data-recipe/Xet) cannot be given a preexec_fn by the parent, so adopt_pid alone left them orphanable on a Linux SIGKILL. They now bind themselves with PR_SET_PDEATHSIG at startup via bind_current_process_to_parent_lifetime(), wired into the shared run_without_native_path_secret entrypoint and the Xet child entry. - Wire the previously-missed data-recipe worker through adopt_pid. - terminate_all now honors its timeout: SIGTERM, wait, then SIGKILL the survivors, so cooperative children can exit cleanly. - Track adopted pids with a /proc starttime identity and add forget_pid, so the shutdown sweep never signals a recycled pid. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: cross-session backstop to reap a leftover llama-server on startup Builds on the parent-lifetime reaper: the Windows Job Object / PR_SET_PDEATHSIG path kills children when the parent dies, and terminate_all() sweeps the living parent's children. Neither covers an orphan left by an already-dead Studio: terminate_all()'s registry is in-memory, PR_SET_PDEATHSIG has no macOS equivalent, and both are best-effort. This records the spawned llama-server PID to a pidfile under the active studio root (removed on _kill_process). The startup reaper kills that exact PID first, verifying it is still a llama-server to guard against PID reuse, then clears the pidfile. It is path-independent, so it also catches an orphan the install-root match would miss; the pidfile only ever names a Studio-spawned server, so unrelated user processes (vllm, games) are never candidates. The existing root-gated enumeration stays as a further fallback. Adds tests: kills a recorded live server (real subprocess, verifies the actual SIGKILL), skips a reused non-llama PID, cleans a stale/missing pidfile, and clears the pidfile on kill. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: only reap a recorded llama-server when it is a true orphan Harden the pidfile cross-session reaper so it cannot kill a live server and does not crash on Windows. - Reap the recorded PID only when its parent is gone (a genuine orphan), so constructing a second LlamaCppBackend in-process (the helper and advisor paths each build one) can never kill the active chat server. The check is topology independent: it holds whether the sweep runs in the main process or a worker. - Record pid:starttime and verify the start-time identity before killing, so a PID recycled to another process is never reaped. - Fall back to SIGTERM when signal.SIGKILL is undefined (Windows), where os.kill maps it to TerminateProcess, instead of raising and leaving the orphan alive while clearing the record. Update and extend the pidfile tests: a live server with a running parent is spared and its record kept, an identity mismatch is skipped, the record-to-reap round trip kills a matching orphan, and the Windows SIGKILL fallback uses SIGTERM. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Michael Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 203 +++++++++++++++++- .../test_llama_cpp_wait_for_vram_settle.py | 203 ++++++++++++++++++ 2 files changed, 405 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index db0ce71f61..3c77487efe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4117,6 +4117,9 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), **_child_popen_kwargs(), ) + # Cross-session backstop: record the PID so a later startup can reap this + # server if parent-death cleanup did not run (macOS / best-effort failure). + self._record_server_pid(self._process.pid) # Start background thread to drain stdout and prevent pipe deadlock self._stdout_thread = threading.Thread( @@ -5450,6 +5453,7 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), **_child_popen_kwargs(), ) + self._record_server_pid(self._process.pid) # Background thread to drain stdout (prevents pipe deadlock) self._stdout_thread = threading.Thread( @@ -6203,6 +6207,7 @@ class LlamaCppBackend: self._stats_logger.stop() self._stats_logger = None self._process = None + self._clear_server_pid() # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False @@ -6221,6 +6226,198 @@ class LlamaCppBackend: pass self._llama_log_fh = None + @staticmethod + def _server_pidfile_path() -> Optional[Path]: + """Pidfile recording the live llama-server PID, under the active studio root + (per-root, so concurrent Studios with distinct UNSLOTH_STUDIO_HOME stay + isolated, mirroring the reaper's custom-root isolation).""" + try: + from utils.paths.storage_roots import studio_root # noqa: WPS433 + return studio_root() / "llama-server.pid" + except Exception: + return None + + @classmethod + def _record_server_pid(cls, pid: int) -> None: + """Best-effort record of the spawned llama-server PID for orphan reaping. + + Stores ``pid:starttime`` so a later startup can reject a PID that has + since been recycled to a different process (see ``_pid_start_identity``). + A bare ``pid`` (no identity) is still accepted on read for compatibility. + """ + path = cls._server_pidfile_path() + if path is None: + return + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(f"{pid}:{cls._pid_start_identity(pid)}") + except Exception as e: + logger.debug(f"Could not write llama-server pidfile: {e}") + + @classmethod + def _clear_server_pid(cls) -> None: + """Best-effort removal of the llama-server pidfile.""" + path = cls._server_pidfile_path() + if path is None: + return + try: + path.unlink(missing_ok = True) + except Exception as e: + logger.debug(f"Could not remove llama-server pidfile: {e}") + + @staticmethod + def _pid_is_llama_server(pid: int) -> bool: + """True only if pid is a live process whose binary is a llama-server. Guards + against PID reuse before killing a recorded orphan; returns False on any + uncertainty so an unrelated process is never killed.""" + try: + import psutil + try: + proc = psutil.Process(pid) + if (proc.name() or "").lower().startswith("llama-server"): + return True + return Path(proc.exe() or "").name.lower().startswith("llama-server") + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return False + except ImportError: + pass + if sys.platform != "linux": + return False + try: + if Path(os.readlink(f"/proc/{pid}/exe")).name.lower().startswith("llama-server"): + return True + except OSError: + pass + try: + with open(f"/proc/{pid}/cmdline", "rb") as fh: + tokens = fh.read().split(b"\x00") + first = tokens[0].decode("utf-8", "replace") if tokens else "" + return Path(first).name.lower().startswith("llama-server") + except OSError: + return False + + @staticmethod + def _pid_start_identity(pid: int) -> str: + """Stable per-PID identity (process start time) guarding against PID reuse. + + Returns a token string, or "" when it cannot be determined (the caller + then falls back to the llama-server name check only).""" + try: + import psutil + try: + return str(psutil.Process(pid).create_time()) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return "" + except ImportError: + pass + if sys.platform == "linux": + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + data = fh.read() + # field 22 (starttime), counted from after the ")" that closes comm. + return data[data.rfind(b")") + 2 :].split()[19].decode() + except (OSError, IndexError): + return "" + return "" + + @staticmethod + def _pid_parent_is_alive(pid: int) -> bool: + """True if the recorded server's parent is still running, i.e. the server is + NOT orphaned. Lets the cross-session reap kill only a true orphan (parent + gone) and never a live server owned by a running Studio, regardless of which + process performs the sweep. Biased toward "alive" on uncertainty so a live + server is never mistakenly reaped.""" + try: + import psutil + + try: + ppid = psutil.Process(pid).ppid() + except psutil.NoSuchProcess: + return False # the recorded server itself is gone + except psutil.Error: + return True # cannot tell -- never risk killing a live server + if ppid <= 1: + return False # reparented to init -> orphan + return psutil.pid_exists(ppid) + except ImportError: + pass + if sys.platform == "linux": + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + data = fh.read() + ppid = int(data[data.rfind(b")") + 2 :].split()[1]) + except (OSError, IndexError, ValueError): + return False + if ppid <= 1: + return False + return Path(f"/proc/{ppid}").exists() + return False + + @staticmethod + def _unlink_pidfile(path: Path) -> None: + """Best-effort removal of a resolved pidfile path.""" + try: + path.unlink(missing_ok = True) + except Exception: + pass + + @classmethod + def _reap_recorded_pid(cls) -> int: + """Kill the exact llama-server PID recorded at spawn, but only when it is a + genuine orphan -- its parent (the Studio that spawned it) is gone. This is + the cross-session backstop the parent-death reaper (Job Object / + PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio + (macOS, a best-effort failure, or a pre-existing orphan). Path-independent, + so it also catches an orphan the install-root match would miss. + + A live server whose parent is still running is never reaped, so constructing + a second backend in-process (the helper / advisor paths each build a + LlamaCppBackend) cannot kill the active chat server. A recorded PID that has + been recycled to a different process is rejected by the start-time identity + and the llama-server name check, so unrelated user processes are never + touched. SIGKILL falls back to SIGTERM on Windows, where os.kill maps it to + TerminateProcess and SIGKILL is undefined.""" + path = cls._server_pidfile_path() + if path is None or not path.exists(): + return 0 + + pid = -1 + identity = "" + try: + pid_str, _, identity = path.read_text().strip().partition(":") + pid = int(pid_str) + except Exception: + pid = -1 + + if pid <= 0: + cls._unlink_pidfile(path) # garbage record + return 0 + if pid == os.getpid(): + return 0 # never our own pid; leave the record alone + + if cls._pid_parent_is_alive(pid): + # Live server with a running parent -> not an orphan; keep the record so + # a later startup can still reap it if that parent later dies abnormally. + return 0 + + # Parent is gone: candidate orphan. Reject a PID recycled to something else. + if identity and cls._pid_start_identity(pid) != identity: + cls._unlink_pidfile(path) + return 0 + + killed = 0 + if cls._pid_is_llama_server(pid): + try: + os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM)) + killed = 1 + logger.info(f"Killed orphaned llama-server from pidfile (pid={pid})") + except (ProcessLookupError, PermissionError): + pass + except Exception as e: + logger.debug(f"Could not kill recorded llama-server pid {pid}: {e}") + cls._unlink_pidfile(path) + return killed + @staticmethod def _kill_orphaned_servers() -> int: """Kill orphaned llama-server processes started by studio. @@ -6238,7 +6435,11 @@ class LlamaCppBackend: Returns the count of processes killed; callers arm the VRAM-settle wait on a positive count. """ - killed = 0 + # Cross-session backstop first: reap the exact PID we recorded at spawn, + # but only if it is a true orphan whose parent is gone (so a helper backend + # built while a chat server is live can never kill it). The root-gated + # enumeration below stays as a fallback. + killed = LlamaCppBackend._reap_recorded_pid() try: # -- Build the ownership allowlist -------------------------------- # exact_binaries -- env var overrides (exact path match). diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index c103c6394a..3c21f41701 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -406,3 +406,206 @@ def test_startup_reaper_arms_settle_timestamp(): assert ( backend_cold._last_kill_monotonic == 0.0 ), "no reap must leave the cold-start sentinel so the wait is skipped" + + +# --------------------------------------------------------------------------- +# Cross-session backstop: a server PID recorded at spawn is reaped on the next +# startup even when parent-death cleanup did not run (macOS, a best-effort +# PR_SET_PDEATHSIG / Job Object failure, or a pre-existing orphan), but ONLY when +# it is a true orphan (its parent is gone), it still is a llama-server, and its +# start-time identity matches. A live server (parent still running) is spared so a +# helper backend built in-process can never kill the active chat server. +# --------------------------------------------------------------------------- + + +class _FakeKillProc: + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def test_kill_process_clears_pidfile(tmp_path): + """A real kill removes the recorded pidfile so a clean eject leaves no orphan marker.""" + pidfile = tmp_path / "llama-server.pid" + pidfile.write_text("12345") + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = _FakeKillProc() + backend._healthy = False + backend._stdout_thread = None + backend._llama_log_fh = None + backend._last_kill_monotonic = 0.0 + backend._stats_logger = None + with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)): + backend._kill_process() + assert not pidfile.exists() + + +def test_reap_recorded_pid_kills_recorded_server(tmp_path): + """An orphaned recorded PID (parent gone) is killed and the pidfile cleared + when it is still a llama-server.""" + import subprocess + + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + pidfile = tmp_path / "llama-server.pid" + pidfile.write_text(str(proc.pid)) + try: + with ( + patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object( + LlamaCppBackend, + "_pid_is_llama_server", + staticmethod(lambda pid: pid == proc.pid), + ), + ): + n = LlamaCppBackend._reap_recorded_pid() + assert n == 1 + assert not pidfile.exists() + proc.wait(timeout = 5) + assert proc.poll() is not None + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout = 5) + + +def test_record_then_reap_round_trip_identity_matches(tmp_path): + """Full round trip: _record_server_pid writes pid:starttime, and an orphaned + reap whose recorded identity still matches DOES kill it.""" + import subprocess + + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + pidfile = tmp_path / "llama-server.pid" + try: + with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)): + LlamaCppBackend._record_server_pid(proc.pid) + assert ":" in pidfile.read_text(), "a start-time identity must be recorded" + with ( + patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + ): + n = LlamaCppBackend._reap_recorded_pid() + assert n == 1, "a matching identity on a true orphan must be reaped" + proc.wait(timeout = 5) + assert proc.poll() is not None + assert not pidfile.exists() + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout = 5) + + +def test_reap_recorded_pid_spares_live_server(tmp_path): + """A recorded server whose parent is still alive (the running Studio) is NEVER + reaped, and its pidfile is kept. This is the finding-3 guard: a helper backend + constructed in-process must not kill the active chat server. Uses the REAL + _pid_parent_is_alive (the child's parent is this live test process).""" + import subprocess + + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + pidfile = tmp_path / "llama-server.pid" + pidfile.write_text(str(proc.pid)) + try: + with ( + patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + # Force the name check True so ONLY the parent-alive guard can spare it. + patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + ): + n = LlamaCppBackend._reap_recorded_pid() + assert n == 0, "a live server with a running parent must not be reaped" + assert proc.poll() is None, "the live server must still be running" + assert pidfile.exists(), "the record is kept so a later orphan reap still works" + finally: + proc.kill() + proc.wait(timeout = 5) + + +def test_reap_recorded_pid_skips_pid_reuse(tmp_path): + """A recorded PID recycled to a non-llama-server must NOT be killed (only the + stale pidfile is cleaned), so the user's vllm/games are never touched.""" + import subprocess + + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + pidfile = tmp_path / "llama-server.pid" + pidfile.write_text(str(proc.pid)) + try: + with ( + patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: False)), + ): + n = LlamaCppBackend._reap_recorded_pid() + assert n == 0 + assert proc.poll() is None, "an unrelated reused PID must not be killed" + assert not pidfile.exists(), "stale pidfile is cleaned up" + finally: + proc.kill() + proc.wait(timeout = 5) + + +def test_reap_recorded_pid_skips_identity_mismatch(tmp_path): + """An orphaned PID whose recorded start-time identity no longer matches has been + recycled; it must NOT be killed even if it now looks like a llama-server.""" + import subprocess + + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + pidfile = tmp_path / "llama-server.pid" + pidfile.write_text(f"{proc.pid}:0.0") # stale identity that cannot match + try: + with ( + patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + ): + n = LlamaCppBackend._reap_recorded_pid() + assert n == 0, "a PID whose start-time identity changed must not be killed" + assert proc.poll() is None, "the recycled process must survive" + assert not pidfile.exists(), "stale pidfile is cleaned up" + finally: + proc.kill() + proc.wait(timeout = 5) + + +def test_reap_recorded_pid_windows_sigkill_fallback(tmp_path, monkeypatch): + """On Windows signal.SIGKILL is undefined; the reaper must fall back to SIGTERM + (os.kill -> TerminateProcess) instead of crashing and leaving the orphan.""" + import os as _os + import signal as _signal + + monkeypatch.delattr(_signal, "SIGKILL", raising = False) + captured = {} + + def _fake_kill(pid, sig): + captured["pid"] = pid + captured["sig"] = sig # recorded; do not actually signal anything + + pidfile = tmp_path / "llama-server.pid" + pidfile.write_text("424242") + with ( + patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), + patch.object(LlamaCppBackend, "_pid_is_llama_server", staticmethod(lambda pid: True)), + patch.object(_os, "kill", _fake_kill), + ): + n = LlamaCppBackend._reap_recorded_pid() + assert n == 1 + assert ( + captured.get("sig") == _signal.SIGTERM + ), "must fall back to SIGTERM when SIGKILL is absent" + assert not pidfile.exists() + + +def test_reap_recorded_pid_no_pidfile(tmp_path): + """No pidfile -> nothing reaped, no error.""" + pidfile = tmp_path / "llama-server.pid" # never created + with patch.object(LlamaCppBackend, "_server_pidfile_path", staticmethod(lambda: pidfile)): + assert LlamaCppBackend._reap_recorded_pid() == 0