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 <michaelhan2050@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
22e6d64493
commit
8e0d082c92
13 changed files with 684 additions and 16 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
319
studio/backend/tests/test_process_lifetime.py
Normal file
319
studio/backend/tests/test_process_lifetime.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
293
studio/backend/utils/process_lifetime.py
Normal file
293
studio/backend/utils/process_lifetime.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue