diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 31e5d9dec0..ee04a91e85 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4114,6 +4114,9 @@ class LlamaCppBackend: env = env, **_windows_hidden_subprocess_kwargs(), ) + # Record the PID so a later startup can reap this server if Studio dies + # abruptly (SIGHUP/SIGKILL/OOM) before _kill_process clears it. + self._record_server_pid(self._process.pid) # Start background thread to drain stdout and prevent pipe deadlock self._stdout_thread = threading.Thread( @@ -5424,6 +5427,7 @@ class LlamaCppBackend: env = env, **_windows_hidden_subprocess_kwargs(), ) + self._record_server_pid(self._process.pid) # Background thread to drain stdout (prevents pipe deadlock) self._stdout_thread = threading.Thread( @@ -6129,6 +6133,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 @@ -6146,6 +6151,101 @@ 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.""" + path = cls._server_pidfile_path() + if path is None: + return + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(str(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 + + @classmethod + def _reap_recorded_pid(cls) -> int: + """Kill the exact llama-server PID recorded at spawn if it is still alive and + still a llama-server, then clear the pidfile. Path-independent, so it 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.""" + path = cls._server_pidfile_path() + if path is None or not path.exists(): + return 0 + killed = 0 + try: + pid = int(path.read_text().strip()) + except Exception: + pid = -1 + if pid > 0 and pid != os.getpid() and cls._pid_is_llama_server(pid): + try: + os.kill(pid, signal.SIGKILL) + 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}") + try: + path.unlink(missing_ok = True) + except Exception: + pass + return killed + @staticmethod def _kill_orphaned_servers() -> int: """Kill orphaned llama-server processes started by studio. @@ -6163,7 +6263,9 @@ class LlamaCppBackend: Returns the count of processes killed; callers arm the VRAM-settle wait on a positive count. """ - killed = 0 + # Reap the exact PID we recorded at spawn first (path-independent); the + # root-gated enumeration below stays as a fallback for pre-pidfile servers. + killed = LlamaCppBackend._reap_recorded_pid() try: # -- Build the ownership allowlist -------------------------------- # exact_binaries -- env var overrides (exact path match). diff --git a/studio/backend/main.py b/studio/backend/main.py index 54946fa992..0b119055e2 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -478,10 +478,18 @@ async def lifespan(app: FastAPI): await _close_llama_http() + def _kill_llama_server_on_shutdown(): + # Mirror run.py's _graceful_shutdown step 5 so a direct-uvicorn shutdown + # (which bypasses the signal handler) also kills the GPU child. + from routes.inference import _llama_cpp_backend + if _llama_cpp_backend is not None: + _llama_cpp_backend._kill_process() + await run_lifespan_shutdown( terminate_hub_downloads, clear_unsloth_compiled_cache, _hw_module, + kill_llama_server = _kill_llama_server_on_shutdown, ) diff --git a/studio/backend/tests/test_lifespan_shutdown.py b/studio/backend/tests/test_lifespan_shutdown.py index 0e3a34441d..b4b1b22af1 100644 --- a/studio/backend/tests/test_lifespan_shutdown.py +++ b/studio/backend/tests/test_lifespan_shutdown.py @@ -132,3 +132,40 @@ def test_run_lifespan_shutdown_preserves_contextvars(): assert seen == ["bound-value"], "terminate must run with the caller's contextvars" assert clear_box["n"] == 1 assert hw.DEVICE is None + + +def test_run_lifespan_shutdown_kills_llama_server_first(): + """The injected kill_llama_server runs once, before hardware/cache teardown, + so a direct-uvicorn shutdown cannot orphan the GPU child.""" + order = [] + hw = types.SimpleNamespace(DEVICE = "cuda:0") + + asyncio.run( + run_lifespan_shutdown( + lambda: order.append("terminate"), + lambda: order.append("clear"), + hw, + kill_llama_server = lambda: order.append("kill"), + ) + ) + + assert order and order[0] == "kill", "llama-server must be killed before other teardown" + assert order.count("kill") == 1 + assert "terminate" in order and "clear" in order + assert hw.DEVICE is None + + +def test_run_lifespan_shutdown_swallows_kill_llama_server_errors(): + """A kill_llama_server failure must not block the remaining cleanup.""" + term_box, terminate = _counter() + clear_box, clear = _counter() + hw = types.SimpleNamespace(DEVICE = "cuda:0") + + def _boom(): + raise RuntimeError("kill failed") + + asyncio.run(run_lifespan_shutdown(terminate, clear, hw, kill_llama_server = _boom)) + + assert term_box["n"] == 1 + assert clear_box["n"] == 1 + assert hw.DEVICE is None 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 91979e3e0f..fd7d9199b4 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 @@ -386,3 +386,96 @@ 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" + + +# --------------------------------------------------------------------------- +# Pidfile-tracked orphan reaping: a server PID recorded at spawn is reaped on +# the next startup regardless of install path, but only if it is still a +# llama-server (PID-reuse guard). Fixes sticky-CPU-after-overflow. +# --------------------------------------------------------------------------- + + +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): + """The recorded PID 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_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_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_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_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 diff --git a/studio/backend/utils/lifespan_shutdown.py b/studio/backend/utils/lifespan_shutdown.py index 6e8371e53d..a88ffe4a7c 100644 --- a/studio/backend/utils/lifespan_shutdown.py +++ b/studio/backend/utils/lifespan_shutdown.py @@ -13,7 +13,7 @@ unit-tested without the heavy backend import graph. import asyncio import contextvars import types -from typing import Callable +from typing import Callable, Optional import structlog @@ -24,8 +24,18 @@ async def run_lifespan_shutdown( terminate_downloads: Callable[[], None], clear_compiled_cache: Callable[[], None], hw_module: types.ModuleType, + kill_llama_server: Optional[Callable[[], None]] = None, ) -> None: """Run each shutdown step guarded so one failure can't skip the others; never raise.""" + # Kill the llama-server child first, before clearing hardware/cache state, so a + # direct-uvicorn shutdown (which bypasses run.py's signal handler) cannot orphan a + # GPU process. The signal / _graceful_shutdown path already covers SIGTERM/SIGINT. + if kill_llama_server is not None: + try: + kill_llama_server() + except Exception as exc: + logger.warning("kill_llama_server failed at shutdown: %s", exc) + loop = asyncio.get_running_loop() # Copy context for parity with asyncio.to_thread. Schedule and await # separately so a dead executor (raises at submit) runs inline, while a