From 8e0d082c92dc7125bb08564196370974dc6df70b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 05:51:22 -0700 Subject: [PATCH] Reap Studio child processes when the parent dies abnormally (#6425) * 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 --------- Co-authored-by: Michael Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/cloudflare_tunnel.py | 11 + .../backend/core/data_recipe/jobs/manager.py | 3 + studio/backend/core/export/orchestrator.py | 3 + studio/backend/core/inference/llama_cpp.py | 23 +- studio/backend/core/inference/orchestrator.py | 3 + studio/backend/core/rag/embed_llama_server.py | 2 + studio/backend/core/training/training.py | 6 + studio/backend/run.py | 16 + studio/backend/tests/test_process_lifetime.py | 319 ++++++++++++++++++ studio/backend/utils/hf_xet_fallback.py | 10 + studio/backend/utils/llama_cpp_update.py | 2 + studio/backend/utils/native_path_leases.py | 9 + studio/backend/utils/process_lifetime.py | 293 ++++++++++++++++ 13 files changed, 684 insertions(+), 16 deletions(-) create mode 100644 studio/backend/tests/test_process_lifetime.py create mode 100644 studio/backend/utils/process_lifetime.py diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index e5dba69452..ef7bacba67 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -49,6 +49,16 @@ def _windows_hidden_kwargs() -> dict: return {"creationflags": flags} if flags else {} +def _lifetime_kwargs() -> dict: + """Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy + + best-effort so this module still loads standalone (storage_roots-style).""" + try: + from utils.process_lifetime import child_popen_kwargs + return child_popen_kwargs() + except Exception: + return {} + + def _asset_name() -> Optional[Tuple[str, bool]]: """(release asset filename, is_tgz) for this OS/arch, or None if unsupported.""" system = platform.system().lower() @@ -233,6 +243,7 @@ class CloudflareTunnel: errors = "replace", bufsize = 1, **_windows_hidden_kwargs(), + **_lifetime_kwargs(), ) self._proc = proc threading.Thread( diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 32523e469e..c238c250bd 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -176,6 +176,9 @@ class JobManager: daemon = True, ) proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) self._mp_q = mp_q self._proc = proc diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index d78866a494..a9694bb8e9 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -147,6 +147,9 @@ class ExportOrchestrator: daemon = True, ) self._proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Export subprocess started (pid=%s)", self._proc.pid) def _shutdown_subprocess(self, timeout: float = 10.0) -> None: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e7030b77ce..db0ce71f61 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -56,6 +56,7 @@ from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs from core.inference.tool_call_parser import ( RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, @@ -3342,28 +3343,16 @@ class LlamaCppBackend: except OSError as e: logger.debug(f"Could not open diffusion runner log file: {e}") - # PR_SET_PDEATHSIG: the shim (and its visual server) die with this backend - # process, so a Studio crash/restart never orphans a GPU process. - popen_kwargs = dict(_windows_hidden_subprocess_kwargs()) - if sys.platform.startswith("linux"): # prctl/libc.so.6 are Linux-only - - def _pdeathsig(): - try: - import ctypes - import signal as _signal - ctypes.CDLL("libc.so.6", use_errno = True).prctl(1, _signal.SIGTERM) - except Exception: - pass - - popen_kwargs["preexec_fn"] = _pdeathsig - + # The shim (and its visual server) die with this backend process, so a + # Studio crash/restart never orphans a GPU process. self._process = subprocess.Popen( cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, env = env, - **popen_kwargs, + **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) self._stdout_thread = threading.Thread( target = self._drain_stdout, daemon = True, name = "diffusion-stdout" @@ -4126,6 +4115,7 @@ class LlamaCppBackend: text = True, env = env, **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) # Start background thread to drain stdout and prevent pipe deadlock @@ -5458,6 +5448,7 @@ class LlamaCppBackend: text = True, env = env, **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) # Background thread to drain stdout (prevents pipe deadlock) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index b81f217539..a0fe8623a3 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -177,6 +177,9 @@ class InferenceOrchestrator: daemon = True, ) self._proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Inference subprocess started (pid=%s)", self._proc.pid) def _cancel_generation(self) -> None: diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index 2fadbcc6c6..c2e4ecc740 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -30,6 +30,7 @@ import numpy as np from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import windows_hidden_subprocess_kwargs +from utils.process_lifetime import child_popen_kwargs from . import config @@ -267,6 +268,7 @@ class LlamaServerBackend: text = True, env = env, **windows_hidden_subprocess_kwargs(), + **child_popen_kwargs(), ) self._process = proc self._port = port diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index d8126363fe..73744466b6 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -367,6 +367,9 @@ class TrainingBackend: daemon = True, ) proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) except Exception: logger.error("Failed to start training subprocess", exc_info = True) return False @@ -529,6 +532,9 @@ class TrainingBackend: daemon = True, ) new_proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep) except Exception: logger.error("Failed to respawn training subprocess", exc_info = True) with self._lock: diff --git a/studio/backend/run.py b/studio/backend/run.py index a9d1f1db26..469365f588 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -642,6 +642,13 @@ def _graceful_shutdown(server = None): except Exception as e: logger.warning("Error stopping Cloudflare tunnel: %s", e) + # 7. Backstop sweep for any adopted child the steps above missed. + try: + from utils.process_lifetime import terminate_all + terminate_all() + except Exception as e: + logger.warning("Error in process-lifetime sweep: %s", e) + logger.info("All subprocesses cleaned up") @@ -876,6 +883,12 @@ def run_server( """ global _server, _shutdown_event + # Reap every child if the parent dies abnormally (terminal close, Task + # Manager kill, SIGKILL); must run before any child can spawn. + from utils.process_lifetime import initialize_parent_lifetime + + initialize_parent_lifetime() + # --secure exposes only the Cloudflare link: force a loopback bind so the raw # port is never public (even with -H 0.0.0.0), and reject the contradictory combo. if secure and not cloudflare: @@ -1083,6 +1096,9 @@ def run_server( import atexit atexit.register(_remove_pid_file) + from utils.process_lifetime import terminate_all + + atexit.register(terminate_all) # Output port for Tauri (api-only), only after sockets bind and startup done. if api_only: diff --git a/studio/backend/tests/test_process_lifetime.py b/studio/backend/tests/test_process_lifetime.py new file mode 100644 index 0000000000..c36cd75405 --- /dev/null +++ b/studio/backend/tests/test_process_lifetime.py @@ -0,0 +1,319 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the parent-lifetime reaper (utils/process_lifetime). + +The Linux PDEATHSIG cases spawn real processes and assert actual liveness; the +Windows Job Object path is exercised with a mocked kernel32 so it runs on CI. +""" + +from __future__ import annotations + +import os +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +import utils.process_lifetime as pl # noqa: E402 + +IS_POSIX = os.name == "posix" +IS_LINUX = sys.platform.startswith("linux") + + +@pytest.fixture(autouse = True) +def _reset_module_state(): + pl._tracked_pids.clear() + pl._win_job_handle = None + pl._initialized = False + yield + pl._tracked_pids.clear() + pl._win_job_handle = None + pl._initialized = False + + +def _alive(pid: int) -> bool: + if sys.platform == "win32": + return _win_alive(pid) + try: + os.kill(pid, 0) # POSIX existence probe (on Windows this would terminate it) + return True + except OSError: + return False + + +def _win_alive(pid: int) -> bool: + import ctypes + + PROCESS_QUERY_LIMITED_INFORMATION, STILL_ACTIVE = 0x1000, 259 + kernel32 = ctypes.windll.kernel32 + handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid) + if not handle: + return False + code = ctypes.c_ulong() + kernel32.GetExitCodeProcess(handle, ctypes.byref(code)) + kernel32.CloseHandle(handle) + return code.value == STILL_ACTIVE + + +def _wait_dead(pid: int, timeout: float) -> bool: + end = time.time() + timeout + while time.time() < end: + if not _alive(pid): + return True + time.sleep(0.05) + return not _alive(pid) + + +# ── No-op safety / composition ── + + +def test_initialize_idempotent_and_noop_on_posix(): + pl.initialize_parent_lifetime() + pl.initialize_parent_lifetime() # second call short-circuits + if IS_POSIX: + assert pl._win_job_handle is None # POSIX installs no job + + +def test_adopt_pid_tolerates_none_and_dead_pid(): + pl.adopt_pid(None) # ignored + pl.adopt_pid(2**31 - 1) # almost-certainly-dead pid: recorded, never raises + assert None not in pl._tracked_pids + + +def test_child_popen_kwargs_linux_vs_other(monkeypatch): + monkeypatch.setattr(pl, "_is_linux", lambda: True) + assert "preexec_fn" in pl.child_popen_kwargs() + monkeypatch.setattr(pl, "_is_linux", lambda: False) + assert pl.child_popen_kwargs() == {} # Windows/macOS add nothing here + + +def test_compose_preexec_runs_pdeathsig_then_existing(monkeypatch): + calls = [] + monkeypatch.setattr(pl, "_is_linux", lambda: True) + monkeypatch.setattr(pl, "_pdeathsig_preexec", lambda: calls.append("death")) + pl.compose_preexec(lambda: calls.append("existing"))() + assert calls == ["death", "existing"] # ordering matters for sandbox hooks + + +def test_compose_preexec_passthrough_off_linux(monkeypatch): + monkeypatch.setattr(pl, "_is_linux", lambda: False) + sentinel = lambda: None # noqa: E731 + assert pl.compose_preexec(sentinel) is sentinel + assert pl.compose_preexec(None) is None + + +# ── Real Linux PDEATHSIG: child dies when the parent dies abnormally ── + + +@pytest.mark.skipif(not IS_LINUX, reason = "PR_SET_PDEATHSIG is Linux-only") +def test_pdeathsig_child_dies_when_parent_sigkilled(tmp_path): + mid = tmp_path / "mid.py" + mid.write_text( + "import sys, subprocess, time\n" + f"sys.path.insert(0, {str(_BACKEND)!r})\n" + "from utils.process_lifetime import child_popen_kwargs\n" + "p = subprocess.Popen(['sleep', '300'], **child_popen_kwargs())\n" + "print(p.pid, flush = True)\n" + "time.sleep(300)\n" + ) + proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True) + try: + sleeper_pid = int(proc.stdout.readline().strip()) + assert _alive(sleeper_pid) + proc.kill() # hard-kill the parent (no graceful shutdown runs) + proc.wait(timeout = 5) + assert _wait_dead(sleeper_pid, 5.0), "child orphaned after parent SIGKILL" + finally: + proc.kill() + + +@pytest.mark.skipif(sys.platform != "win32", reason = "Windows Job Object") +def test_windows_job_kills_child_when_parent_dies(tmp_path): + # Real kill-on-job-close: the parent installs the job and assigns itself, a + # child inherits it automatically, and terminating the parent must reap the + # child (the orphaned-cloudflared.exe scenario). + mid = tmp_path / "mid.py" + mid.write_text( + "import sys, subprocess, time\n" + f"sys.path.insert(0, {str(_BACKEND)!r})\n" + "import utils.process_lifetime as pl\n" + "pl.initialize_parent_lifetime()\n" + "p = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(300)'])\n" + "print(p.pid, int(pl._win_job_handle is not None), flush = True)\n" + "time.sleep(300)\n" + ) + proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True) + try: + first = proc.stdout.readline().split() + child_pid, installed = int(first[0]), first[1] == "1" + assert installed, "Windows Job Object was not installed" + assert _alive(child_pid) + proc.kill() # TerminateProcess the parent -> last job handle closes + proc.wait(timeout = 5) + assert _wait_dead(child_pid, 5.0), "child orphaned after parent killed" + finally: + proc.kill() + + +# ── terminate_all backstop sweep ── + + +@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep") +def test_terminate_all_signals_tracked_and_is_idempotent(): + p = subprocess.Popen(["sleep", "300"]) + pl.adopt_pid(p.pid) + pl.terminate_all() + assert p.wait(timeout = 5) is not None # reap + confirm it died + pl.terminate_all() # registry now empty; must not raise + + +@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep") +def test_terminate_all_escalates_to_sigkill(): + # A child that ignores SIGTERM must still be reaped via SIGKILL after timeout. + p = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(300)", + ] + ) + time.sleep(0.5) # let the handler install + pl.adopt_pid(p.pid) + pl.terminate_all(timeout = 0.3) + assert p.wait(timeout = 5) == -signal.SIGKILL # SIGTERM ignored, SIGKILL wins + + +@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep") +def test_terminate_all_lets_cooperative_child_exit_cleanly(tmp_path): + # A child that handles SIGTERM gets `timeout` to exit cleanly (not -SIGKILL). + marker = tmp_path / "clean.txt" + p = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, sys, time\n" + f"def h(*a): open({str(marker)!r}, 'w').write('clean'); sys.exit(0)\n" + "signal.signal(signal.SIGTERM, h)\n" + "time.sleep(300)\n", + ] + ) + time.sleep(0.5) + pl.adopt_pid(p.pid) + pl.terminate_all(timeout = 3.0) + assert p.wait(timeout = 3) == 0 # exited via its own handler, not SIGKILL + assert marker.read_text() == "clean" + + +def test_forget_pid_unregisters(): + pl.adopt_pid(4242) + assert 4242 in pl._tracked_pids + pl.forget_pid(4242) + assert 4242 not in pl._tracked_pids + + +@pytest.mark.skipif(not IS_POSIX, reason = "POSIX process sweep") +def test_terminate_all_skips_recycled_pid(monkeypatch): + # A tracked pid whose identity changed (recycled) must not be signalled. + p = subprocess.Popen(["sleep", "300"]) + pl.adopt_pid(p.pid) # records the real identity + monkeypatch.setattr(pl, "_pid_identity", lambda _pid: "DIFFERENT") + pl.terminate_all() + assert _alive(p.pid) # left untouched: identity mismatch + p.kill() + p.wait(timeout = 5) + + +@pytest.mark.skipif(not IS_LINUX, reason = "PR_SET_PDEATHSIG is Linux-only") +def test_bind_kills_multiprocessing_child_on_parent_death(tmp_path): + # multiprocessing workers can't take a preexec_fn, so the child binds itself + # via bind_current_process_to_parent_lifetime(). Killing the parent must reap + # it (the gap reviewers found in adopt_pid alone). + mid = tmp_path / "mid_mp.py" + mid.write_text( + "import sys, time, multiprocessing as mp\n" + f"sys.path.insert(0, {str(_BACKEND)!r})\n" + "from utils.process_lifetime import bind_current_process_to_parent_lifetime\n" + "def _child():\n" + " bind_current_process_to_parent_lifetime()\n" + " time.sleep(300)\n" + "if __name__ == '__main__':\n" + " p = mp.get_context('spawn').Process(target = _child, daemon = True)\n" + " p.start()\n" + " print(p.pid, flush = True)\n" + " time.sleep(300)\n" + ) + proc = subprocess.Popen([sys.executable, str(mid)], stdout = subprocess.PIPE, text = True) + try: + child_pid = int(proc.stdout.readline().strip()) + assert _alive(child_pid) + proc.kill() + proc.wait(timeout = 5) + assert _wait_dead(child_pid, 5.0), "mp child orphaned after parent SIGKILL" + finally: + proc.kill() + + +# ── Windows Job Object path (mocked kernel32, runs on Linux CI) ── + + +class _Call: + def __init__(self, name, log, ret): + self.name, self.log, self.ret = name, log, ret + self.restype = self.argtypes = None + + def __call__(self, *a, **k): + self.log.append(self.name) + return self.ret + + +class _FakeKernel32: + def __init__( + self, + log, + create_ret = 4321, + set_ret = 1, + assign_ret = 1, + ): + self.CreateJobObjectW = _Call("create", log, create_ret) + self.SetInformationJobObject = _Call("set", log, set_ret) + self.AssignProcessToJobObject = _Call("assign", log, assign_ret) + self.GetCurrentProcess = _Call("getcur", log, -1) + self.CloseHandle = _Call("close", log, 1) + + +def _patch_windows(monkeypatch, fake): + import ctypes + monkeypatch.setattr(pl, "_is_windows", lambda: True) + monkeypatch.setattr(ctypes, "WinDLL", lambda *a, **k: fake, raising = False) + + +def test_windows_job_install_order(monkeypatch): + log: list[str] = [] + _patch_windows(monkeypatch, _FakeKernel32(log)) + pl._install_windows_job() + assert log.index("create") < log.index("set") < log.index("assign") + assert pl._win_job_handle == 4321 # handle retained + + +def test_windows_job_install_degrades_on_create_failure(monkeypatch): + log: list[str] = [] + _patch_windows(monkeypatch, _FakeKernel32(log, create_ret = 0)) + pl._install_windows_job() # must not raise + assert pl._win_job_handle is None + assert "set" not in log # short-circuited after the failed create + + +def test_windows_job_install_degrades_on_assign_failure(monkeypatch): + log: list[str] = [] + _patch_windows(monkeypatch, _FakeKernel32(log, assign_ret = 0)) + pl._install_windows_job() + assert pl._win_job_handle is None # not retained when assignment fails + assert "close" in log # the orphaned job handle is closed diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 80f42f92d1..a6ba69fffc 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -168,6 +168,13 @@ def _download_child_entry( forms its own process group so the parent can kill the whole transfer, and never logs the token or signed URLs. """ + # Die with Studio on Linux (this mp child gets no parent-set preexec_fn). + try: + from utils.process_lifetime import bind_current_process_to_parent_lifetime + bind_current_process_to_parent_lifetime() + except Exception: + pass + if hasattr(os, "setsid"): try: os.setsid() @@ -277,6 +284,9 @@ def _run_download_attempt( daemon = True, ) proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) stalled = threading.Event() stop_watchdog = start_watchdog( diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 22ffc58b24..965f1a970a 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -43,6 +43,7 @@ from utils.llama_cpp_freshness import ( reset_caches, update_download_size_bytes, ) +from utils.process_lifetime import child_popen_kwargs logger = structlog.get_logger(__name__) @@ -459,6 +460,7 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path stderr = subprocess.STDOUT, text = True, env = env, + **child_popen_kwargs(), ) timed_out = threading.Event() diff --git a/studio/backend/utils/native_path_leases.py b/studio/backend/utils/native_path_leases.py index 7d8514abc8..08671cfe39 100644 --- a/studio/backend/utils/native_path_leases.py +++ b/studio/backend/utils/native_path_leases.py @@ -83,6 +83,15 @@ def child_env_without_native_path_secret(env: Mapping[str, str] | None = None) - def run_without_native_path_secret(target: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: """Run a multiprocessing child target without the native path lease secret.""" + # Runs in the spawned child: bind it to the parent's death (Linux), since + # multiprocessing children cannot be given a preexec_fn by the parent. Shared + # entrypoint for the inference/export/training/data-recipe workers. + try: + from utils.process_lifetime import bind_current_process_to_parent_lifetime + bind_current_process_to_parent_lifetime() + except Exception: + pass + global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET os.environ.pop(LEASE_SECRET_ENV, None) _CACHED_LEASE_SECRET = None diff --git a/studio/backend/utils/process_lifetime.py b/studio/backend/utils/process_lifetime.py new file mode 100644 index 0000000000..3ffd54cc26 --- /dev/null +++ b/studio/backend/utils/process_lifetime.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Bind Studio child processes to the parent's lifetime so none survive an +abnormal parent exit (terminal-window close, Task Manager "End Task", SIGKILL, +crash) -- the cooperative shutdown path only runs on graceful exits. + +Windows: one parent-owned Job Object with JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE. +The parent is assigned to it, children inherit it automatically, and the OS +reaps every process in the job when the parent's last handle closes. Mirrors the +desktop app's job in studio/src-tauri/src/windows_job.rs. + +POSIX: each long-lived child sets prctl(PR_SET_PDEATHSIG) on Linux via a tiny +preexec hook (macOS has no equivalent and relies on the cooperative path + +terminate_all). Linux's signal is per-direct-child only, so multiprocessing +workers are also tracked for terminate_all. + +Best-effort throughout: any failure degrades to today's behavior, never raises. +Stdlib only. +""" + +from __future__ import annotations + +import os +import signal +import sys +import threading +from typing import Callable, Optional + +_PR_SET_PDEATHSIG = 1 +_JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000 +_JobObjectExtendedLimitInformation = 9 + +_lock = threading.Lock() +_initialized = False +_win_job_handle: Optional[int] = None # retained for the interpreter's lifetime +_tracked_pids: "dict[int, Optional[str]]" = {} # pid -> identity, reaped by terminate_all + + +def _is_linux() -> bool: + return sys.platform.startswith("linux") + + +def _is_windows() -> bool: + return sys.platform == "win32" + + +# ── Parent setup ── + + +def initialize_parent_lifetime() -> None: + """Install the parent-death reaper once, as early as possible at startup. + + Windows builds and holds the Job Object; POSIX has nothing to install (the + guarantee is per-child via preexec). Idempotent and never raises. + """ + global _initialized + with _lock: + if _initialized: + return + _initialized = True + if _is_windows(): + _install_windows_job() + + +def _win_signatures(kernel32) -> None: + # Explicit HANDLE-width signatures. Without argtypes, ctypes marshals the + # 64-bit job/process handles as c_int and truncates them on Win64, so the + # job calls silently operate on a bogus handle and assignment fails. + import ctypes + from ctypes import wintypes + + H, BOOL, DWORD = wintypes.HANDLE, wintypes.BOOL, wintypes.DWORD + kernel32.CreateJobObjectW.argtypes = [ctypes.c_void_p, ctypes.c_wchar_p] + kernel32.CreateJobObjectW.restype = H + kernel32.SetInformationJobObject.argtypes = [H, ctypes.c_int, ctypes.c_void_p, DWORD] + kernel32.SetInformationJobObject.restype = BOOL + kernel32.AssignProcessToJobObject.argtypes = [H, H] + kernel32.AssignProcessToJobObject.restype = BOOL + kernel32.GetCurrentProcess.argtypes = [] + kernel32.GetCurrentProcess.restype = H + kernel32.CloseHandle.argtypes = [H] + kernel32.CloseHandle.restype = BOOL + + +def _install_windows_job() -> None: + global _win_job_handle + try: + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + _win_signatures(kernel32) + + class _BASIC(ctypes.Structure): + _fields_ = [ + ("PerProcessUserTimeLimit", ctypes.c_int64), + ("PerJobUserTimeLimit", ctypes.c_int64), + ("LimitFlags", wintypes.DWORD), + ("MinimumWorkingSetSize", ctypes.c_size_t), + ("MaximumWorkingSetSize", ctypes.c_size_t), + ("ActiveProcessLimit", wintypes.DWORD), + ("Affinity", ctypes.c_size_t), + ("PriorityClass", wintypes.DWORD), + ("SchedulingClass", wintypes.DWORD), + ] + + class _IO(ctypes.Structure): + _fields_ = [ + (n, ctypes.c_uint64) + for n in ( + "ReadOperationCount", + "WriteOperationCount", + "OtherOperationCount", + "ReadTransferCount", + "WriteTransferCount", + "OtherTransferCount", + ) + ] + + class _EXT(ctypes.Structure): + _fields_ = [ + ("BasicLimitInformation", _BASIC), + ("IoInfo", _IO), + ("ProcessMemoryLimit", ctypes.c_size_t), + ("JobMemoryLimit", ctypes.c_size_t), + ("PeakProcessMemoryUsed", ctypes.c_size_t), + ("PeakJobMemoryUsed", ctypes.c_size_t), + ] + + job = kernel32.CreateJobObjectW(None, None) + if not job: + return + info = _EXT() + info.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + if not kernel32.SetInformationJobObject( + job, _JobObjectExtendedLimitInformation, ctypes.byref(info), ctypes.sizeof(info) + ): + kernel32.CloseHandle(job) + return + # AssignProcessToJobObject(parent) makes children inherit the job. May + # fail if Studio already runs inside an incompatible host job (pre-Win8); + # degrade to the cooperative path rather than blocking startup. + if not kernel32.AssignProcessToJobObject(job, kernel32.GetCurrentProcess()): + kernel32.CloseHandle(job) + return + _win_job_handle = job # hold the handle so the job is not closed early + except Exception: + pass + + +# ── Child binding ── + + +def _pdeathsig_preexec() -> None: + # Runs in the forked child before exec. prctl is Linux-only; the getppid + # check closes the race where the parent died before this ran. + try: + import ctypes + ctypes.CDLL("libc.so.6", use_errno = True).prctl(_PR_SET_PDEATHSIG, signal.SIGTERM) + if os.getppid() == 1: + os._exit(1) + except Exception: + pass + + +def bind_current_process_to_parent_lifetime() -> None: + """Bind the CURRENT process to its parent's death (Linux). For multiprocessing + children, which cannot take a preexec_fn, so the parent cannot set + PR_SET_PDEATHSIG for them -- the child must do it itself at startup.""" + if _is_linux(): + _pdeathsig_preexec() + + +def compose_preexec(existing: Optional[Callable[[], None]]) -> Optional[Callable[[], None]]: + """Run the PDEATHSIG hook then any caller-supplied preexec (Linux only).""" + if not _is_linux(): + return existing + if existing is None: + return _pdeathsig_preexec + + def _composed() -> None: + _pdeathsig_preexec() + existing() + + return _composed + + +def child_popen_kwargs(preexec_fn: Optional[Callable[[], None]] = None) -> dict: + """Popen kwargs that bind a long-lived child to the parent's lifetime. + + On Linux returns a composed ``preexec_fn`` (PDEATHSIG + any existing one); + empty elsewhere (Windows is covered by the inherited Job Object). Merge via + ``**child_popen_kwargs()`` alongside the caller's existing kwargs. + """ + if _is_linux(): + return {"preexec_fn": compose_preexec(preexec_fn)} + return {} + + +def _pid_identity(pid: int) -> Optional[str]: + # Linux /proc starttime (stat field 22); pins identity so a reused pid is not + # signalled later. None (other platforms / unreadable) disables the check. + if not _is_linux(): + return None + try: + with open(f"/proc/{pid}/stat", encoding = "utf-8") as fh: + stat = fh.read() + return stat[stat.rfind(")") + 2 :].split()[19] # after comm: starttime + except Exception: + return None + + +def forget_pid(pid: Optional[int]) -> None: + """Stop tracking a child the owner has reaped, so terminate_all never + signals a recycled pid.""" + if pid: + _tracked_pids.pop(pid, None) + + +def adopt_pid(pid: Optional[int]) -> None: + """Track a child (e.g. a multiprocessing worker started after the parent job + was set up) and, on Windows, assign it to the job as belt-and-suspenders. + Tolerates a None or already-exited pid.""" + if not pid: + return + _tracked_pids[pid] = _pid_identity(pid) + if _is_windows() and _win_job_handle: + try: + import ctypes + from ctypes import wintypes + + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + _win_signatures(kernel32) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + PROCESS_SET_QUOTA, PROCESS_TERMINATE = 0x0100, 0x0001 + handle = kernel32.OpenProcess(PROCESS_SET_QUOTA | PROCESS_TERMINATE, False, pid) + if handle: + kernel32.AssignProcessToJobObject(_win_job_handle, handle) + kernel32.CloseHandle(handle) + except Exception: + pass + + +def terminate_all(timeout: float = 5.0) -> None: + """Backstop sweep over adopted pids, after per-subsystem cleanup. SIGTERM, + then SIGKILL the survivors after `timeout`. Skips a pid whose identity no + longer matches (recycled). Idempotent and teardown-safe.""" + for pid, identity in list(_tracked_pids.items()): + _tracked_pids.pop(pid, None) + if identity is not None and _pid_identity(pid) != identity: + continue # pid was reused by an unrelated process + try: + if _is_windows(): + os.kill(pid, signal.SIGTERM) + continue + _posix_terminate(pid, timeout) + except Exception: + pass + + +def _posix_terminate(pid: int, timeout: float = 5.0) -> None: + # SIGTERM, give the child up to `timeout` to exit, then SIGKILL. Reaping + # belongs to the child's owner (or init for orphans). Prefer the group + # (covers grandchildren) when pid leads its own group. + import time + + killer = os.kill + try: + if os.getpgid(pid) == pid: + killer = os.killpg + except Exception: + pass + try: + killer(pid, signal.SIGTERM) + except ProcessLookupError: + return + except Exception: + return + deadline = time.monotonic() + max(0.0, timeout) + while time.monotonic() < deadline: + try: + killer(pid, 0) # still alive? + except ProcessLookupError: + return + except Exception: + break + time.sleep(0.05) + try: + killer(pid, signal.SIGKILL) + except Exception: + pass