diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index f9d7a0cb81..3131eb0ba2 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -91,6 +91,24 @@ def _lifetime_kwargs() -> dict: return {} +def _adopt_pid(pid: int) -> None: + """Record cloudflared so a force quit does not strand it (macOS has no + PDEATHSIG). Best-effort, like _lifetime_kwargs above.""" + try: + from utils.process_lifetime import adopt_pid + adopt_pid(pid) + except Exception: + pass + + +def _forget_pid(pid: int) -> None: + try: + from utils.process_lifetime import forget_pid + forget_pid(pid) + except Exception: + pass + + def _spawn_child(spawn): """Fork on a process-lifetime thread so the PDEATHSIG above means "die with the parent process", not "die when the worker thread that forked me returns".""" @@ -447,6 +465,10 @@ def start(self) -> None: except Exception: _set_studio_tunnel_runtime_active(self, False) raise + # Adopted before the lock drops: a stop() that got in first would + # otherwise reap and forget it while nothing was tracked, and this + # would then record whatever inherited the pid. + _adopt_pid(proc.pid) self._proc = proc threading.Thread( target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True @@ -522,6 +544,7 @@ def stop(self) -> bool: except Exception: pass if _process_exited(proc): + _forget_pid(proc.pid) _set_studio_tunnel_runtime_active(self, False) return True else: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7623f5734e..66ec454009 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -7290,9 +7290,21 @@ def _start_diffusion_server( encoding = "utf-8", errors = "replace", env = utf8_child_env(env), + # Deliberately NOT start_new_session, as with the component + # installer: the desktop stops this backend by signalling its + # process group and force-kills it after five seconds, so a session + # of its own would leave the shim and the visual server holding the + # GPU until the next launch sweeps them. **_windows_hidden_subprocess_kwargs(), **_child_popen_kwargs(), ) + # macOS has no parent-death signal, so the kwargs above are empty there and + # only this record lets the next startup reap a runner holding the GPU. + try: + from utils.process_lifetime import adopt_pid + adopt_pid(self._process.pid) + except Exception as e: + logger.debug(f"Could not track diffusion runner for lifetime sweep: {e}") self._stdout_thread = threading.Thread( target = self._drain_stdout, daemon = True, name = "diffusion-stdout" ) @@ -13284,6 +13296,47 @@ def unload_model(self) -> bool: torch.cuda.empty_cache() return True + @staticmethod + def _leading_process_group(pid): + """The pid's own process group, when it leads one. None otherwise.""" + if not pid or os.name != "posix" or not hasattr(os, "getpgid"): + return None + try: + pgid = os.getpgid(pid) + except OSError: + return None + return pgid if pgid == pid else None + + @staticmethod + def _kill_process_group(pgid): + """Take down what the leader left behind, if anything is still there.""" + if pgid is None or not hasattr(os, "killpg"): + return + try: + os.killpg(pgid, signal.SIGKILL) + except OSError: + pass + + @staticmethod + def _collect_descendants(pid): + """The server's own children, for the kill below. Empty when unreadable.""" + try: + from utils.process_lifetime import collect_descendants + return collect_descendants(pid) + except Exception: + return [] + + @staticmethod + def _terminate_descendants(collected): + """The diffusion shim's visual server, and anything else it started.""" + if not collected: + return + try: + from utils.process_lifetime import terminate_descendants + terminate_descendants(collected, timeout = 5.0) + except Exception as e: + logger.debug(f"Could not terminate server descendants: {e}") + def _kill_process(self): """Terminate the subprocess if running.""" # Stop the watchdog before a deliberate kill so a planned reload/unload @@ -13303,6 +13356,13 @@ def _kill_process(self): terminable = hasattr(self._process, "terminate") if not terminable: logger.debug("no terminable llama-server process to kill; clearing state") + # Both read before the terminate below: getpgid stops answering once the + # wait reaps the leader, and the shim's children are reparented the + # moment it exits. This is the only stop the desktop shutdown waits for, + # so the visual server has to be named while that link still exists. + _pid = getattr(self._process, "pid", None) + _pgid = self._leading_process_group(_pid) + _descendants = self._collect_descendants(_pid) try: if terminable: self._process.terminate() @@ -13323,11 +13383,25 @@ def _kill_process(self): except Exception as e: logger.warning(f"Error killing llama-server process: {e}") finally: + self._kill_process_group(_pgid) + self._terminate_descendants(_descendants) # getattr: teardown must tolerate a partially-built backend (failed # __init__ or a __new__-built instance), as with _llama_log_fh below. if getattr(self, "_stats_logger", None) is not None: self._stats_logger.stop() self._stats_logger = None + # Drop it from the lifetime record only once it is confirmed gone. + # A server that survived a failed kill has to stay recorded, or the + # next startup sweep cannot reap it. The record stores a start-time + # identity, so a recycled pid is never signalled either way. + _killed_pid = getattr(self._process, "pid", None) + _exited = getattr(self._process, "poll", lambda: None)() is not None + if _killed_pid is not None and _exited: + try: + from utils.process_lifetime import forget_pid + forget_pid(_killed_pid) + except Exception: + pass self._process = None self._clear_server_pid() # Clear healthy so a /load during the replacement's warm-up can't @@ -13371,6 +13445,14 @@ def _record_server_pid(cls, pid: int) -> None: since been recycled to a different process (see ``_pid_start_identity``). A bare ``pid`` (no identity) is still accepted on read for compatibility. """ + # Track it generically too: the pidfile holds one server, while the + # process-lifetime record covers every child and is what the startup + # sweep reads where there is no parent-death signal (macOS). + try: + from utils.process_lifetime import adopt_pid + adopt_pid(pid) + except Exception as e: + logger.debug(f"Could not track llama-server for lifetime sweep: {e}") path = cls._server_pidfile_path() if path is None: return diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 159d40668c..541c323aac 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -25,7 +25,7 @@ from pathlib import Path from typing import Callable, Optional -from utils.process_lifetime import child_popen_kwargs +from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid from utils.native_path_leases import child_env_without_native_path_secret from core.inference.sd_cpp_args import ( SdCppGenParams, @@ -413,6 +413,9 @@ def _run( # Bind the child to the parent's lifetime (PR_SET_PDEATHSIG) so a parent crash cannot orphan sd-cli holding VRAM/RAM. **child_popen_kwargs(), ) + # The kwargs above are empty on macOS, so record it too: a crash mid-generation + # would otherwise leave sd-cli holding VRAM with nothing able to find it. + adopt_pid(proc.pid) # Drain stdout on a reader thread so the timeout holds even when the child hangs WITHOUT printing (a plain `for line in proc.stdout` blocks until EOF). Lines, then a None sentinel, go to a queue the main loop polls against a wall-clock deadline. tail: list[str] = [] line_q: "queue.Queue[Optional[str]]" = queue.Queue() @@ -459,6 +462,10 @@ def _drain() -> None: finally: if proc.poll() is None: _terminate(proc) + # Only once it has actually exited: a pid still running has to stay + # recorded, or the next startup has no handle on it. + if proc.poll() is not None: + forget_pid(proc.pid) if ret != 0: raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index e6fc7f0b3b..c63806054f 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -9940,14 +9940,54 @@ def _check_code_safety(code: str) -> str | None: return None +def _adopt_tool_pid(pid: "int | None") -> None: + """Record a tool subprocess for the startup sweep. + + macOS has no parent-death signal, so a force quit mid-call would otherwise + leave a session-leading tool (and whatever it spawned) with nothing able to + find it. Best-effort: a failure here must never break a tool call. + """ + if not pid: + return + try: + from utils.process_lifetime import adopt_pid + adopt_pid(pid) + except Exception: + pass + + +def _forget_tool_pid(proc) -> None: + """Drop the record once the process has actually exited.""" + pid = getattr(proc, "pid", None) + if not pid: + return + try: + if getattr(proc, "poll", lambda: None)() is None: + return + from utils.process_lifetime import forget_pid + forget_pid(pid) + except Exception: + pass + + def _capture_process_group(proc): """Return the setsid process-group id, or ``None`` when unavailable. Captured right after ``Popen`` so a later ``poll()`` / ``wait()`` that reaps - the leader cannot make ``os.getpgid(proc.pid)`` fail first. POSIX-only: - Windows has no process groups (and no ``os.getpgid``), so return ``None`` - there and let the single-pid ``proc.kill()`` fallback handle cleanup. + the leader cannot make ``os.getpgid(proc.pid)`` fail first. + + Windows has no process groups, so capture the wrapper pid instead, tagged + for ``_killpg_captured`` to reach with ``taskkill /T``; returning ``None`` + there left a payload that outlived its wrapper unsignalled. """ + if os.name == "nt": + job = _windows_job_capture(proc) + if job is not None: + return ("windows-job", job) + # No job available, so fall back to the pid, carrying its creation-time + # identity: a posix group id cannot be recycled while a member lives, + # but this bare pid can, and the timeout path may fire much later. + return ("windows-tree", proc.pid, _windows_pid_identity(proc.pid)) if os.name != "posix" or not hasattr(os, "getpgid"): return None try: @@ -9956,10 +9996,120 @@ def _capture_process_group(proc): return None +class _WindowsToolJob: + """A job object holding one tool call's process tree. + + Windows has no process groups, and ``taskkill`` cannot reach a tree whose + root has already exited, which is exactly the case this capture exists for. + The job stays a valid handle on every descendant either way. Created without + kill-on-close, so releasing it never kills a process that outlived the call. + """ + + def __init__(self, handle, kernel32): + self._handle = handle + self._kernel32 = kernel32 + + def terminate(self) -> bool: + if not self._handle: + return False + return bool(self._kernel32.TerminateJobObject(self._handle, 1)) + + def close(self) -> None: + handle, self._handle = self._handle, None + if handle: + try: + self._kernel32.CloseHandle(handle) + except Exception: # noqa: BLE001 - interpreter teardown + pass + + def __del__(self) -> None: + self.close() + + +def _windows_job_capture(proc) -> "_WindowsToolJob | None": + """Put ``proc`` in its own job. ``None`` when that is not possible, leaving + the pid-based fallback.""" + if os.name != "nt": + return None + try: + import ctypes + from ctypes import wintypes + + H, BOOL, UINT = wintypes.HANDLE, wintypes.BOOL, wintypes.UINT + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + # Explicit widths: without them ctypes truncates a 64-bit handle to + # c_int and every call silently works on a bogus one. + kernel32.CreateJobObjectW.argtypes = [ctypes.c_void_p, ctypes.c_wchar_p] + kernel32.CreateJobObjectW.restype = H + kernel32.AssignProcessToJobObject.argtypes = [H, H] + kernel32.AssignProcessToJobObject.restype = BOOL + kernel32.TerminateJobObject.argtypes = [H, UINT] + kernel32.TerminateJobObject.restype = BOOL + kernel32.CloseHandle.argtypes = [H] + kernel32.CloseHandle.restype = BOOL + + job = kernel32.CreateJobObjectW(None, None) + if not job: + return None + # The Popen handle, not a fresh OpenProcess: it already refers to this + # child, so there is no window for the pid to be recycled first. + if not kernel32.AssignProcessToJobObject(job, int(proc._handle)): + kernel32.CloseHandle(job) + return None + return _WindowsToolJob(job, kernel32) + except Exception: # noqa: BLE001 - falls back to the pid-based kill + return None + + +def _windows_pid_identity(pid: int) -> "str | None": + """Process creation time, so a recycled pid is never mistaken for this one.""" + if os.name != "nt": + return None + try: + from utils.process_lifetime import _pid_identity + return _pid_identity(pid) + except Exception: + return None + + +def _windows_taskkill_tree(pid: int, identity: "str | None" = None) -> bool: + """``taskkill /T /F`` a pid and its descendants. True when it succeeded. + + Every tool call runs under a shell wrapper, and Windows has no process + groups, so a bare ``proc.kill()`` reaps the wrapper and orphans the payload + (usually the venv python), which then blocks `unsloth studio update`. + + ``identity`` is the creation time captured at spawn; a mismatch means the pid + now belongs to something else, so nothing is signalled. + """ + if os.name != "nt": + return False + if identity is not None and _windows_pid_identity(pid) != identity: + return False + try: + completed = subprocess.run( + ["taskkill", "/PID", str(pid), "/T", "/F"], + capture_output = True, + timeout = 15, + creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0), + ) + except (OSError, subprocess.SubprocessError): + return False + return completed.returncode in (0, 128) # 128: already gone + + def _kill_process_tree(proc) -> None: """SIGKILL the setsid process group; fall back to single-pid kill.""" if proc.poll() is not None: return + if os.name == "nt": + if _windows_taskkill_tree(proc.pid): + return + try: + proc.kill() + except (ProcessLookupError, PermissionError): + pass + return pgid = None if hasattr(os, "getpgid"): try: @@ -9984,9 +10134,23 @@ def _killpg_captured(pgid) -> None: Once ``proc`` exits, ``os.getpgid(proc.pid)`` fails and ``_kill_process_tree`` short-circuits, so a stdout-holding grandchild that outlived the parent could not otherwise be signaled. The pre-captured setsid group id still targets the - whole tree. No-op with no ``os.killpg`` (Windows) or nothing captured. + whole tree. On Windows the capture is a tagged pid and the equivalent reach + is ``taskkill /T /F``. No-op when nothing was captured. """ - if pgid is None or not hasattr(os, "killpg"): + if pgid is None: + return + if isinstance(pgid, tuple): + if pgid[0] == "windows-job": + pgid[1].terminate() + return + _tag, pid, identity = pgid + # Fail closed: this runs long after the capture, so without a verified + # identity the pid may be someone else's now. The job object still takes + # the whole tree when Studio exits, which is the safe half to keep. + if identity is not None: + _windows_taskkill_tree(pid, identity) + return + if not hasattr(os, "killpg"): return try: os.killpg(pgid, signal.SIGKILL) @@ -10294,6 +10458,7 @@ def _python_exec( # Capture the group before any watcher can reap the leader (see # _capture_process_group); None on Windows. pgid = _capture_process_group(proc) + _adopt_tool_pid(proc.pid) if cancel_event is not None: watcher = threading.Thread( @@ -10351,6 +10516,7 @@ def _python_exec( except Exception as e: return f"Execution error: {e}" finally: + _forget_tool_pid(locals().get("proc")) if tmp_path and os.path.exists(tmp_path): try: os.unlink(tmp_path) @@ -10419,6 +10585,7 @@ def _bash_exec( # Capture the group before any watcher can poll/reap the leader (see # _python_exec); None on Windows. pgid = _capture_process_group(proc) + _adopt_tool_pid(proc.pid) if cancel_event is not None: watcher = threading.Thread( @@ -10450,3 +10617,5 @@ def _bash_exec( except Exception as e: return f"Execution error: {e}" + finally: + _forget_tool_pid(locals().get("proc")) diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index b3ac62e520..f043df2129 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -30,7 +30,7 @@ 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 utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid from . import config @@ -340,6 +340,9 @@ def _spawn_once(self, use_gpu: bool) -> None: **child_popen_kwargs(), ) self._process = proc + # Long-lived, and child_popen_kwargs() is empty on macOS, so the crash + # record is the only thing that can reap it after a force quit. + adopt_pid(proc.pid) self._port = port self._stdout_thread = threading.Thread( target = self._drain_stdout, @@ -425,6 +428,10 @@ def _kill_process(self) -> None: except Exception as e: # noqa: BLE001 logger.warning("error killing llama-server embedder: %s", e) finally: + # Only once it is confirmed gone: a survivor must stay recorded so + # the next startup sweep can reap it. + if proc.poll() is not None: + forget_pid(proc.pid) self._process = None if self._stdout_thread is not None: self._stdout_thread.join(timeout = 2) diff --git a/studio/backend/run.py b/studio/backend/run.py index 0f2ed42fec..b68d3af38a 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1136,6 +1136,79 @@ def _remove_pid_file(): pass +# Windows terminates the process ~5s after a close event, so leave a margin. +_CONSOLE_SHUTDOWN_BUDGET = 4.5 + + +# CTRL_CLOSE / CTRL_LOGOFF / CTRL_SHUTDOWN. Ctrl+C and Ctrl+Break (0 and 1) are +# left out on purpose: Python already delivers those as signals. +_CONSOLE_SHUTDOWN_EVENTS = (2, 5, 6) + + +def _console_event_is_shutdown(event: int) -> bool: + return event in _CONSOLE_SHUTDOWN_EVENTS + + +def _run_console_shutdown(shutdown) -> None: + try: + shutdown() + except Exception as error: + logger.warning("Console-close cleanup failed: %s", error) + + +def _install_windows_console_handler(shutdown) -> bool: + """Run the graceful shutdown when the console window is closed. + + Closing the window raises CTRL_CLOSE_EVENT, which Python never turns into a + signal, so neither a signal handler nor atexit runs and cleanup is skipped. + ``shutdown`` takes no arguments and must not touch signal.signal: Windows + runs this on a thread it creates for the event, and Windows kills the + process about five seconds later, so the work is bounded to fit. No-op off + Windows. + """ + if sys.platform != "win32": + return False + try: + import ctypes + from ctypes import wintypes + + HANDLER = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.DWORD) + + import threading + + def _on_console_event(event: int) -> bool: + if _console_event_is_shutdown(event): + worker = threading.Thread( + target = _run_console_shutdown, args = (shutdown,), daemon = True + ) + worker.start() + worker.join(timeout = _CONSOLE_SHUTDOWN_BUDGET) + return True + # Ctrl+C / Ctrl+Break already arrive as Python signals; pass them + # on rather than shutting down twice. + return False + + callback = HANDLER(_on_console_event) + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.SetConsoleCtrlHandler.argtypes = [HANDLER, wintypes.BOOL] + kernel32.SetConsoleCtrlHandler.restype = wintypes.BOOL + if not kernel32.SetConsoleCtrlHandler(callback, True): + logger.warning( + "Could not install the console-close handler (WinError %s); closing the " + "window will skip subprocess cleanup.", + ctypes.get_last_error(), + ) + return False + # Hold a reference: a collected callback leaves Windows calling into + # freed memory. + globals()["_WINDOWS_CONSOLE_HANDLER"] = callback + logger.info("Console-close handler installed") + return True + except Exception as error: + logger.warning("Could not install the console-close handler: %s", error) + return False + + def _graceful_shutdown(server = None): """Shut down all subprocess backends and the uvicorn server. @@ -1189,8 +1262,9 @@ def _graceful_shutdown(server = None): # 7. Backstop sweep for any adopted child the steps above missed. try: - from utils.process_lifetime import terminate_all + from utils.process_lifetime import clear_breadcrumb, terminate_all terminate_all() + clear_breadcrumb() # nothing left for the next startup to sweep except Exception as e: logger.warning("Error in process-lifetime sweep: %s", e) @@ -1849,9 +1923,18 @@ def run_server( # 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 + from utils.process_lifetime import initialize_parent_lifetime, reap_recorded_children initialize_parent_lifetime() + # macOS has neither PR_SET_PDEATHSIG nor job objects, so a Studio that + # crashed left its sidecars running. Sweep before spawning anything: a + # leftover holds VRAM, a port, and the files an update has to replace. + try: + reaped = reap_recorded_children() + if reaped: + logger.warning("Reaped %d orphan(s) from a previous Studio: %s", len(reaped), reaped) + except Exception as e: + logger.warning("Could not sweep orphans from a previous run: %s", e) # --secure exposes ONLY the Cloudflare link: reject --secure --no-cloudflare, # then force a loopback bind so the raw port is never public (even -H 0.0.0.0). @@ -2488,6 +2571,15 @@ def _signal_handler(signum, frame): if hasattr(signal, "SIGBREAK"): signal.signal(signal.SIGBREAK, _signal_handler) + # NOT _signal_handler: Windows runs this on a thread it creates, and + # signal.signal() off the main thread raises, which would leave the window + # close doing no cleanup at all. + def _console_shutdown(): + _graceful_shutdown(_server) + _shutdown_event.set() + + _install_windows_console_handler(_console_shutdown) + # Keep running until shutdown signal. Event.wait() without a timeout blocks at # the C level on Linux, preventing SIGINT delivery; a short timeout in a loop # lets the interpreter process pending signals. diff --git a/studio/backend/tests/test_child_lifetime_boundary.py b/studio/backend/tests/test_child_lifetime_boundary.py new file mode 100644 index 0000000000..2ab4567cdb --- /dev/null +++ b/studio/backend/tests/test_child_lifetime_boundary.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Where the Windows orphan guarantee holds, and where it stops. + +Children inherit a Job Object with KILL_ON_JOB_CLOSE. The desktop updater clears +that flag before launching the installer, so both sides are pinned here with +real processes. +""" + +import os +import subprocess +import sys +import textwrap +import time +from pathlib import Path + +import pytest + +IS_WINDOWS = sys.platform == "win32" +BACKEND = str(Path(__file__).resolve().parents[1]) + + +def _alive(pid: int) -> bool: + if IS_WINDOWS: + out = subprocess.run( + ["tasklist", "/FI", f"PID eq {pid}", "/NH"], capture_output = True, text = True + ).stdout + return str(pid) in out + try: + os.kill(pid, 0) + except OSError: + return False + return True + + +def _kill(pid: int) -> None: + try: + if IS_WINDOWS: + subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], capture_output = True) + else: + os.kill(pid, 9) + except Exception: + pass + + +PARENT = textwrap.dedent( + """ + import os, subprocess, sys, time + sys.path.insert(0, sys.argv[1]) + if sys.argv[3] == "job": + from utils.process_lifetime import initialize_parent_lifetime + initialize_parent_lifetime() + from core.inference.tools import _get_shell_cmd + payload = ( + "import os,time,pathlib;" + "pathlib.Path(r'%s').write_text(str(os.getpid()));" + "time.sleep(180)" % sys.argv[2] + ) + argv = _get_shell_cmd('"%s" -c "%s"' % (sys.executable, payload)) + kw = {} + if sys.argv[3] == "job": + # The helper Studio's own spawns go through: PR_SET_PDEATHSIG on Linux, + # nothing on macOS, job inheritance on Windows. + from utils.process_lifetime import child_popen_kwargs + kw = child_popen_kwargs() + subprocess.Popen(argv, **kw) + time.sleep(180) + """ +) + + +def _run_case(tmp_path: Path, mode: str) -> bool: + """Start a parent, let it spawn a shell-wrapped payload, hard-kill the + parent, and report whether the payload survived.""" + pidfile = tmp_path / f"{mode}.pid" + script = tmp_path / f"parent_{mode}.py" + script.write_text(PARENT) + parent = subprocess.Popen([sys.executable, str(script), BACKEND, str(pidfile), mode]) + try: + for _ in range(200): + if pidfile.is_file() and pidfile.read_text().strip(): + break + time.sleep(0.1) + payload_pid = int(pidfile.read_text().strip()) + assert _alive(payload_pid) + + # Hard kill, no tree flag: exactly what "End Task" / a crash does. + if IS_WINDOWS: + subprocess.run(["taskkill", "/PID", str(parent.pid), "/F"], capture_output = True) + else: + os.kill(parent.pid, 9) + parent.wait(timeout = 30) + time.sleep(5) + survived = _alive(payload_pid) + print(f"\n[{sys.platform}] mode={mode}: payload {payload_pid} survived = {survived}") + return survived + finally: + try: + _kill(int(pidfile.read_text().strip())) + except Exception: + pass + _kill(parent.pid) + + +@pytest.mark.skipif(not IS_WINDOWS, reason = "job objects are Windows-only") +def test_job_object_reaps_the_whole_tree(tmp_path): + """With the guarantee in force, even a shell grandchild is reaped.""" + assert _run_case(tmp_path, "job") is False + + +@pytest.mark.skipif(not IS_WINDOWS, reason = "job objects are Windows-only") +def test_without_the_job_the_grandchild_is_orphaned(tmp_path): + """This is the state the desktop updater leaves the app in: the job is + still there but KILL_ON_JOB_CLOSE has been cleared and nothing restores + it, so a hard exit leaks every child.""" + assert _run_case(tmp_path, "nojob") is True + + +@pytest.mark.skipif(IS_WINDOWS, reason = "POSIX contrast") +def test_posix_orphan_behaviour(tmp_path): + """Linux reaps via PR_SET_PDEATHSIG on the direct child; macOS has no + equivalent and relies entirely on the cooperative shutdown path.""" + survived = _run_case(tmp_path, "job") + if sys.platform == "darwin": + assert survived, "macOS now reaps orphans -- update this repro" + else: + assert not survived + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-q", "-s"])) diff --git a/studio/backend/tests/test_cloudflare_tunnel.py b/studio/backend/tests/test_cloudflare_tunnel.py index 866e654b8f..d072fab0d5 100644 --- a/studio/backend/tests/test_cloudflare_tunnel.py +++ b/studio/backend/tests/test_cloudflare_tunnel.py @@ -291,6 +291,7 @@ def __init__(self): self.terminated = False self.killed = False self._alive = True + self.pid = 424243 # every real Popen has one; the lifetime record reads it def poll(self): return None if self._alive else 0 diff --git a/studio/backend/tests/test_combined_update.py b/studio/backend/tests/test_combined_update.py index b96d3d030c..53861fcbd9 100644 --- a/studio/backend/tests/test_combined_update.py +++ b/studio/backend/tests/test_combined_update.py @@ -73,9 +73,13 @@ def __init__( ): if on_start is not None: on_start(list(cmd)) + self.pid = 515151 # a real Popen has one; the lifetime record needs it self.returncode = returncode self.stdout = iter(lines or []) + def poll(self): + return self.returncode + def wait(self): return self.returncode diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 0f6a810399..3d625fc5fc 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -45,9 +45,13 @@ def __init__( captured_kwargs.update(kwargs) if on_start is not None: on_start(list(cmd)) + self.pid = 515151 # a real Popen has one; the lifetime record needs it self.returncode = returncode self.stdout = iter(lines or []) + def poll(self): + return self.returncode + def wait(self): return self.returncode diff --git a/studio/backend/tests/test_orphaned_children.py b/studio/backend/tests/test_orphaned_children.py new file mode 100644 index 0000000000..fb79c4e0d8 --- /dev/null +++ b/studio/backend/tests/test_orphaned_children.py @@ -0,0 +1,2431 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""No child outlives the Studio that spawned it. + +The chain this closes: a tool call runs under a shell wrapper, the kill path +reaped only the wrapper on Windows, and the orphaned venv python then made +`unsloth studio update` refuse to run until it was killed by hand. +""" + +import json +import os +import signal +import subprocess +import threading +import sys +import time +from pathlib import Path + +import pytest + +IS_WINDOWS = sys.platform == "win32" + + +def _load_installer_module(): + """install_llama_prebuilt.py, imported from the repo without installing it.""" + import importlib.util + + path = Path(__file__).resolve().parents[2] / "install_llama_prebuilt.py" + name = "install_llama_prebuilt_under_test" + if name in sys.modules: + return sys.modules[name] + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + # Registered before exec: its dataclasses resolve annotations through + # sys.modules, and a module that is not in there fails to build them. + sys.modules[name] = module + try: + spec.loader.exec_module(module) + except Exception: + del sys.modules[name] + raise + return module + + +def _alive(pid: int) -> bool: + if IS_WINDOWS: + out = subprocess.run( + ["tasklist", "/FI", f"PID eq {pid}", "/NH"], + capture_output = True, + text = True, + ).stdout + return str(pid) in out + try: + os.kill(pid, 0) + except OSError: + return False + return True + + +def _kill(pid: int) -> None: + try: + if IS_WINDOWS: + subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], capture_output = True) + else: + os.kill(pid, 9) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# 1. Killing a tool call takes the payload with it +# --------------------------------------------------------------------------- +def test_tool_kill_takes_the_shell_payload_with_it(tmp_path): + from core.inference.tools import _get_shell_cmd, _kill_process_tree + + pidfile = tmp_path / "child.pid" + payload = ( + f"import os,time,pathlib;" + f"pathlib.Path(r'{pidfile}').write_text(str(os.getpid()));" + f"time.sleep(120)" + ) + command = f'"{sys.executable}" -c "{payload}"' + argv = _get_shell_cmd(command) + print(f"\n[{sys.platform}] shell wrapper argv[0] = {argv[0]}") + + kwargs = {} + if not IS_WINDOWS: + kwargs["start_new_session"] = True # what the real spawn does on POSIX + proc = subprocess.Popen(argv, **kwargs) + + for _ in range(100): + if pidfile.is_file() and pidfile.read_text().strip(): + break + time.sleep(0.1) + grandchild = int(pidfile.read_text().strip()) + assert _alive(grandchild) + + try: + _kill_process_tree(proc) + proc.wait(timeout = 10) + time.sleep(2.0) + survived = _alive(grandchild) + print(f"payload pid {grandchild} alive after _kill_process_tree: {survived}") + # Windows reaches the payload via taskkill /T /F, POSIX via killpg. It + # used to survive on Windows, orphaning the venv python that then blocked + # `unsloth studio update`. + assert not survived, "the payload under the shell wrapper was orphaned" + finally: + _kill(grandchild) + + +# --------------------------------------------------------------------------- +# 2. One surviving venv process blocks `unsloth studio update` (Windows) +# --------------------------------------------------------------------------- +@pytest.mark.skipif(not IS_WINDOWS, reason = "the update gate is Windows-only") +def test_update_gate_blocks_on_a_single_orphan(tmp_path): + from unsloth_cli import _studio_runtime_gate + + studio_home = tmp_path / "studio_home" + venv = studio_home / "unsloth_studio" + subprocess.run([sys.executable, "-m", "venv", "--without-pip", str(venv)], check = True) + venv_python = venv / "Scripts" / "python.exe" + assert venv_python.is_file() + + # No orphan: the gate lets the update through. + _studio_runtime_gate.ensure_managed_environment_is_idle(studio_home) + + orphan = subprocess.Popen([str(venv_python), "-c", "import time; time.sleep(120)"]) + try: + time.sleep(2.0) + with pytest.raises(RuntimeError) as excinfo: + _studio_runtime_gate.ensure_managed_environment_is_idle(studio_home) + print(f"\nupdate gate said: {excinfo.value}") + assert "in use by" in str(excinfo.value) + assert str(orphan.pid) in str(excinfo.value) + finally: + _kill(orphan.pid) + + +@pytest.mark.skipif(IS_WINDOWS, reason = "contrast case for POSIX") +def test_update_gate_is_a_noop_on_posix(tmp_path): + """On Linux/macOS nothing checks for a running Studio before an update.""" + from unsloth_cli import _studio_runtime_gate + _studio_runtime_gate.ensure_managed_environment_is_idle(tmp_path / "anything") + + +# --------------------------------------------------------------------------- +# 3. Shutdown paths that are not signals +# --------------------------------------------------------------------------- +def test_console_close_runs_the_graceful_shutdown(): + """Closing the console window is not a signal, so it needs its own handler.""" + import run + + assert hasattr(run, "_install_windows_console_handler") + if not IS_WINDOWS: + assert run._install_windows_console_handler(lambda *a: None) is False + return + + calls = [] + assert run._install_windows_console_handler(lambda *a: calls.append(a)) is True + + +def test_windows_job_status_is_reported(): + """A silent failure meant nobody could tell the guarantee was off.""" + from utils import process_lifetime + + process_lifetime.initialize_parent_lifetime() + in_force, detail = process_lifetime.windows_job_status() + print(f"\n[{sys.platform}] job status: in_force={in_force} detail={detail!r}") + assert detail != "not attempted" + if IS_WINDOWS or sys.platform.startswith("linux"): + assert in_force, detail + else: + # macOS: no kernel mechanism, so it must say so rather than imply cover. + assert not in_force + + +def test_macos_style_orphans_are_recorded_and_reaped(tmp_path, monkeypatch): + """The on-disk child record is the only reaper macOS has after a crash.""" + from utils import process_lifetime as pl + + records = tmp_path / "children" + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(records)) + pl._tracked_pids.clear() + + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(120)"]) + try: + pl.adopt_pid(child.pid) + record = records / f"{os.getpid()}.json" + assert record.is_file(), "the child was not recorded" + + # Pretend a previous Studio wrote this and died. The identity is what + # decides, not the pid: a dead pid is recycled fast on a busy machine + # (macOS especially), and an owner whose start time no longer matches is + # a different process, so its recorded children are orphans. + import json + + payload = json.loads(record.read_text()) + payload["owner_identity"] = "a-previous-studio-that-is-gone" + record.write_text(json.dumps(payload)) + + reaped = pl.reap_recorded_children() + print(f"\nreaped: {reaped}") + assert child.pid in reaped + # Reap our own zombie so the liveness check means something. + child.wait(timeout = 10) + assert not _alive(child.pid) + assert not record.exists(), "the record should be consumed" + finally: + _kill(child.pid) + + +def test_liveness_probe_does_not_kill_what_it_probes(): + """os.kill(pid, 0) is TerminateProcess on Windows, so the probe cannot use it.""" + from utils import process_lifetime as pl + + victim = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + try: + for _ in range(3): + assert pl._pid_alive(victim.pid) is True + time.sleep(0.3) + assert victim.poll() is None, "the liveness probe killed the process" + finally: + _kill(victim.pid) + + dead = subprocess.Popen([sys.executable, "-c", "pass"]) + dead.wait(timeout = 30) + assert pl._pid_alive(dead.pid) is False + + +def test_a_second_studio_does_not_erase_the_first_record(tmp_path, monkeypatch): + """Two Studios can share a home; one record per owner keeps both tracked.""" + from utils import process_lifetime as pl + + records = tmp_path / "children" + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(records)) + pl._tracked_pids.clear() + + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + try: + pl.adopt_pid(child.pid) + # A sibling Studio's record, written under its own pid. + import json + + other = records / "424242.json" + other.write_text( + json.dumps( + { + "owner_pid": 424242, + "owner_identity": "a-studio-that-is-gone", + "children": [], + } + ) + ) + pl.adopt_pid(child.pid) # rewrites ours only + assert other.is_file(), "a sibling's record was erased" + assert (records / f"{os.getpid()}.json").is_file() + finally: + _kill(child.pid) + + +def test_a_live_owner_is_never_reaped(tmp_path, monkeypatch): + """Two Studios at once must not kill each other's children.""" + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + pl._tracked_pids.clear() + + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + try: + pl.adopt_pid(child.pid) + # owner_pid is this process, which is very much alive. + assert pl.reap_recorded_children() == [] + assert _alive(child.pid) + finally: + _kill(child.pid) + + +@pytest.mark.parametrize( + "content", + [ + "", + "not json", + "[]", + "null", + '{"children": "nope"}', + '{"children": [1, 2]}', + '{"children": [{"pid": "x"}]}', + ], +) +def test_a_malformed_record_never_blocks_startup(tmp_path, monkeypatch, content): + """The sweep runs before the server binds, so it must not raise on a record + written by an older build or truncated by a power cut.""" + from utils import process_lifetime as pl + + directory = tmp_path / "children" + directory.mkdir() + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(directory)) + (directory / "700000.json").write_text(content) + assert pl.reap_recorded_children() == [] + + +def test_identity_survives_a_child_renaming_itself(): + """Identity is only ever compared for one pid, and comm is mutable: a worker + calling prctl(PR_SET_NAME) or setproctitle would otherwise read as recycled + and be dropped unsignalled.""" + from utils import process_lifetime as pl + + if not pl._is_linux(): + pytest.skip("Linux only") + + child = subprocess.Popen( + [ + sys.executable, + "-c", + "import ctypes, time\n" + "time.sleep(0.3)\n" + 'ctypes.CDLL("libc.so.6").prctl(15, b"renamed-worker", 0, 0, 0)\n' + "time.sleep(30)\n", + ] + ) + try: + before = pl._pid_identity(child.pid) + with open(f"/proc/{child.pid}/comm", encoding = "utf-8") as fh: + name_before = fh.read().strip() + time.sleep(1.0) + with open(f"/proc/{child.pid}/comm", encoding = "utf-8") as fh: + assert fh.read().strip() != name_before, "the child did not rename itself" + after = pl._pid_identity(child.pid) + assert before == after + assert pl._same_identity(before, after) + finally: + _kill(child.pid) + + +def test_a_record_written_with_a_process_name_still_matches(): + """Records from an earlier build carry starttime:comm.""" + from utils import process_lifetime as pl + + if not pl._is_linux(): + pytest.skip("Linux only") + assert pl._same_identity("196794665:python3", "196794665") is True + assert pl._same_identity("196794665:python3", "196794999") is False + + +def test_a_group_outliving_its_leader_is_still_reaped(tmp_path, monkeypatch): + """The shim can crash while the visual server holds the GPU; the record's + pid is then gone and the group is the only handle left.""" + import json + + from utils import process_lifetime as pl + + if pl._is_windows(): + pytest.skip("POSIX only") + + marker = tmp_path / "grandchild.pid" + leader = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, pathlib\n" + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'])\n" + f"pathlib.Path({str(marker)!r}).write_text(str(child.pid))\n", + ], + start_new_session = True, + ) + try: + leader.wait(timeout = 30) + for _ in range(100): + if marker.is_file(): + break + time.sleep(0.1) + grandchild = int(marker.read_text()) + assert _alive(grandchild) + + record = tmp_path / "900000.json" + record.write_text( + json.dumps( + { + "owner_pid": 900000, + "owner_identity": None, + "children": [ + {"pid": leader.pid, "identity": "gone", "pgid": leader.pid}, + ], + } + ), + encoding = "utf-8", + ) + pl._reap_one_record(record, timeout = 5.0) + time.sleep(0.5) + assert not _alive(grandchild), "the visual server survived its dead leader" + finally: + _kill(leader.pid) + + +def test_a_child_sharing_studios_group_is_never_recorded(): + """Recording that group would make the sweep kill Studio and every sibling.""" + from utils import process_lifetime as pl + + if pl._is_windows(): + pytest.skip("POSIX only") + shared = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + own = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], start_new_session = True + ) + try: + assert pl._own_process_group(shared.pid) is None + assert pl._own_process_group(own.pid) == own.pid + finally: + _kill(shared.pid) + _kill(own.pid) + + +def test_ctrl_c_is_passed_on_by_the_console_handler(): + """Ctrl+C and Ctrl+Break must report "not handled" so Python's own signal + handler still runs. Returning True (or raising inside the ctypes callback, + where the BOOL result is then undefined) would leave Studio unstoppable.""" + import run + + assert run._console_event_is_shutdown(0) is False # CTRL_C_EVENT + assert run._console_event_is_shutdown(1) is False # CTRL_BREAK_EVENT + for close_event in (2, 5, 6): # CLOSE, LOGOFF, SHUTDOWN + assert run._console_event_is_shutdown(close_event) is True + + +def test_a_surviving_child_stays_in_the_record(tmp_path, monkeypatch): + """forget_pid is for confirmed exits; a survivor must remain reapable.""" + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + pl._tracked_pids.clear() + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + try: + pl.adopt_pid(child.pid) + assert child.pid in pl._tracked_pids + # A failed kill leaves poll() as None, so llama_cpp must not forget it. + assert child.poll() is None + finally: + _kill(child.pid) + + +def test_concurrent_adopts_all_survive(tmp_path, monkeypatch): + """Two threads adopting at once must not lose either pid.""" + import json + import threading + + from utils import process_lifetime as pl + + directory = tmp_path / "children" + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(directory)) + pl._tracked_pids.clear() + pids = list(range(900000, 900040)) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: f"id-{pid}") + + def adopt(chunk): + for pid in chunk: + pl.adopt_pid(pid) + + threads = [threading.Thread(target = adopt, args = (pids[i::4],)) for i in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + record = json.loads((directory / f"{os.getpid()}.json").read_text()) + assert sorted(entry["pid"] for entry in record["children"]) == pids + pl._tracked_pids.clear() + + +def test_a_survivor_keeps_its_record_through_a_clean_shutdown(tmp_path, monkeypatch): + """terminate_all cannot confirm every exit, and the record is the only + handle the next startup has on what is left.""" + from utils import process_lifetime as pl + + directory = tmp_path / "children" + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(directory)) + pl._tracked_pids.clear() + + stubborn = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + try: + pl.adopt_pid(stubborn.pid) + # Signalling silently does nothing, as an unkillable child would. + monkeypatch.setattr(pl, "_posix_terminate", lambda pid, timeout = 5.0: None) + monkeypatch.setattr(pl.os, "kill", lambda pid, sig: None) + + survivors = pl.terminate_all() + assert survivors == [stubborn.pid] + + pl.clear_breadcrumb() + record = directory / f"{os.getpid()}.json" + assert record.is_file(), "the only handle on the survivor was deleted" + import json + + assert [e["pid"] for e in json.loads(record.read_text())["children"]] == [stubborn.pid] + finally: + _kill(stubborn.pid) + pl._tracked_pids.clear() + + +def test_a_confirmed_shutdown_still_clears_the_record(tmp_path, monkeypatch): + from utils import process_lifetime as pl + + directory = tmp_path / "children" + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(directory)) + pl._tracked_pids.clear() + + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + pl.adopt_pid(child.pid) + assert pl.terminate_all(timeout = 3.0) == [] + pl.clear_breadcrumb() + assert list(directory.glob("*.json")) == [] + + +def test_a_live_owner_survives_an_unreadable_identity(tmp_path, monkeypatch): + """A `ps` that failed for a moment must not cost a running Studio its + sidecars.""" + import json + + from utils import process_lifetime as pl + + directory = tmp_path / "children" + directory.mkdir() + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(directory)) + + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + try: + (directory / "123456.json").write_text( + json.dumps( + { + "owner_pid": os.getppid(), # alive + "owner_identity": "recorded-when-it-started", + "children": [{"pid": child.pid, "identity": pl._pid_identity(child.pid)}], + } + ) + ) + # The lookup fails only for the owner, as a transient `ps` error would. + real = pl._pid_identity + monkeypatch.setattr( + pl, + "_pid_identity", + lambda pid: None if pid == os.getppid() else real(pid), + ) + assert pl.reap_recorded_children() == [] + assert child.poll() is None, "a live Studio's child was killed" + finally: + _kill(child.pid) + + +def test_an_unverifiable_captured_pid_is_not_taskkilled(monkeypatch): + """The delayed Windows path fails closed: the job object covers the tree.""" + from core.inference import tools + + ran = [] + monkeypatch.setattr( + tools, "_windows_taskkill_tree", lambda pid, identity = None: ran.append(pid) or True + ) + tools._killpg_captured(("windows-tree", 4321, None)) + assert ran == [], "signalled a pid it could not verify" + tools._killpg_captured(("windows-tree", 4321, "created-at-t0")) + assert ran == [4321] + + +def test_the_status_is_not_claimed_when_prctl_is_blocked(monkeypatch): + """Under a seccomp or container policy that rejects prctl, children can + survive a hard parent exit; the diagnostic must not say otherwise.""" + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_is_windows", lambda: False) + monkeypatch.setattr(pl, "_is_linux", lambda: True) + monkeypatch.setattr(pl, "_pdeathsig_available", lambda: False) + monkeypatch.setattr(pl, "_initialized", False) + pl.initialize_parent_lifetime() + in_force, detail = pl.windows_job_status() + assert in_force is False + assert "prctl" in detail + + monkeypatch.setattr(pl, "_pdeathsig_available", lambda: True) + monkeypatch.setattr(pl, "_initialized", False) + pl.initialize_parent_lifetime() + assert pl.windows_job_status()[0] is True + + +def test_the_probe_does_not_arm_anything(monkeypatch): + """PR_GET_PDEATHSIG is read-only, so probing must leave our own setting be.""" + import ctypes + + from utils import process_lifetime as pl + + if not pl._is_linux(): + pytest.skip("Linux only") + libc = ctypes.CDLL("libc.so.6", use_errno = True) + before = ctypes.c_int(0) + libc.prctl(2, ctypes.byref(before), 0, 0, 0) + assert pl._pdeathsig_available() is True + after = ctypes.c_int(0) + libc.prctl(2, ctypes.byref(after), 0, 0, 0) + assert after.value == before.value + + +def test_the_sweep_snapshot_is_taken_under_the_lock(): + """Writes hold _record_lock, so the read has to as well.""" + import ast + import inspect + + from utils import process_lifetime as pl + + tree = ast.parse(inspect.getsource(pl.terminate_all)) + body = ast.dump(tree) + assert "_record_lock" in body + # The snapshot must not be a bare list() over the live dict. + source = inspect.getsource(pl.terminate_all) + assert "for pid, identity in list(_tracked_pids.items())" not in source + + +def test_the_probe_fails_when_only_the_read_is_permitted(monkeypatch): + """seccomp filters prctl on its first argument, so a working GET says + nothing about SET.""" + import ctypes + + from utils import process_lifetime as pl + + class _Libc: + def prctl(self, op, *rest): + return 0 if op == 2 else -1 # GET ok, SET rejected + + monkeypatch.setattr(ctypes, "CDLL", lambda *a, **k: _Libc()) + assert pl._pdeathsig_available() is False + + +def test_sd_cli_is_recorded_while_it_runs(tmp_path, monkeypatch): + """It holds VRAM for the length of a generation and macOS has no + parent-death signal, so a crash mid-run must leave something to reap.""" + from core.inference import sd_cpp_engine + + adopted, forgotten = [], [] + monkeypatch.setattr(sd_cpp_engine, "adopt_pid", adopted.append) + monkeypatch.setattr(sd_cpp_engine, "forget_pid", forgotten.append) + monkeypatch.setattr(sd_cpp_engine, "runtime_env", lambda binary, base = None: dict(os.environ)) + + engine = sd_cpp_engine.SdCppEngine.__new__(sd_cpp_engine.SdCppEngine) + engine.binary = sys.executable + out = tmp_path / "image.png" + engine._run( + [sys.executable, "-c", f"open({str(out)!r}, 'wb').write(b'x')"], + str(out), + timeout = 60, + env = None, + on_log = None, + ) + assert len(adopted) == 1 and adopted[0] > 0 + assert forgotten == adopted, "a finished run must not stay on the record" + + +def test_the_diffusion_runner_is_recorded_too(): + """The llama-server path adopts through _record_server_pid; this one spawns + directly and was invisible to the sweep.""" + import ast + import inspect + + from core.inference import llama_cpp + + source = inspect.getsource(llama_cpp) + tree = ast.parse(source) + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef): + continue + body = ast.dump(node) + if "diffusion-stdout" not in body: + continue + assert "adopt_pid" in body, f"{node.name} spawns a runner it never records" + break + else: + raise AssertionError("could not find the diffusion runner spawn") + + +def test_a_windows_tool_tree_dies_with_its_job(monkeypatch): + """taskkill cannot reach a tree whose root already exited, which is the + case this capture exists for; the job handle still can.""" + from core.inference import tools + + class _Job: + def __init__(self): + self.terminated = 0 + + def terminate(self): + self.terminated += 1 + return True + + taskkilled = [] + monkeypatch.setattr( + tools, "_windows_taskkill_tree", lambda pid, identity = None: taskkilled.append(pid) or True + ) + job = _Job() + tools._killpg_captured(("windows-job", job)) + assert job.terminated == 1 + assert taskkilled == [], "the job is the whole reach; no pid revalidation needed" + + +def test_the_windows_breadcrumb_fallback_takes_the_whole_tree(tmp_path, monkeypatch): + """Killing the leader alone strands its workers, and the record naming them + is deleted straight after.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_is_windows", lambda: True) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: pid == 4242) + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: "same" if pid == 4242 else None) + trees = [] + monkeypatch.setattr(pl, "_windows_terminate_tree", trees.append) + + record = tmp_path / "999.json" + record.write_text( + json.dumps( + { + "owner_pid": 999, + "owner_identity": None, + "children": [{"pid": 4242, "identity": "same"}], + } + ), + encoding = "utf-8", + ) + killed, _deferred = pl._reap_one_record(record, timeout = 1.0) + assert killed == [4242] + assert trees == [4242], "only the leader was signalled" + + +def test_the_diffusion_runner_stays_in_the_backend_group(): + """The desktop stops this backend by signalling its process group and + force-kills it five seconds later, so a session of its own would leave the + shim and the visual server holding the GPU until the next launch.""" + import ast + import inspect + + from core.inference import llama_cpp + + tree = ast.parse(inspect.getsource(llama_cpp)) + for node in ast.walk(tree): + if not isinstance(node, ast.FunctionDef) or "diffusion-stdout" not in ast.dump(node): + continue + for call in ast.walk(node): + if not isinstance(call, ast.Call): + continue + if getattr(call.func, "attr", None) != "Popen": + continue + assert not any( + kw.arg == "start_new_session" for kw in call.keywords + ), "a session of its own is out of reach of the group the desktop stops" + return + raise AssertionError("could not find the diffusion runner spawn") + + +def test_a_group_leader_is_reaped_with_its_children(tmp_path, monkeypatch): + """What start_new_session buys: killpg takes the grandchild too.""" + from utils import process_lifetime as pl + + if pl._is_windows(): + pytest.skip("POSIX only") + + marker = tmp_path / "grandchild.pid" + leader = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, time, pathlib\n" + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'])\n" + f"pathlib.Path({str(marker)!r}).write_text(str(child.pid))\n" + "time.sleep(60)\n", + ], + start_new_session = True, + ) + try: + for _ in range(100): + if marker.is_file(): + break + time.sleep(0.1) + grandchild = int(marker.read_text()) + pl._posix_terminate(leader.pid, timeout = 5.0) + time.sleep(0.5) + assert not _alive(grandchild), "the visual server survived its runner" + finally: + _kill(leader.pid) + + +def test_forgetting_a_leader_keeps_its_live_group(tmp_path, monkeypatch): + """The shim can exit before the visual server it started; dropping the + record then loses the only handle on that group.""" + from utils import process_lifetime as pl + + if pl._is_windows(): + pytest.skip("POSIX only") + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + pl._tracked_pids.clear() + pl._tracked_pgids.clear() + + marker = tmp_path / "grandchild.pid" + leader = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, pathlib\n" + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'])\n" + f"pathlib.Path({str(marker)!r}).write_text(str(child.pid))\n", + ], + start_new_session = True, + ) + try: + pl.adopt_pid(leader.pid) + assert pl._tracked_pgids.get(leader.pid) == leader.pid + leader.wait(timeout = 30) + for _ in range(100): + if marker.is_file(): + break + time.sleep(0.1) + grandchild = int(marker.read_text()) + + pl.forget_pid(leader.pid) + assert leader.pid in pl._tracked_pids, "the group was dropped with its dead leader" + + # And the backstop then takes it. + pl.terminate_all(timeout = 5.0) + time.sleep(0.5) + assert not _alive(grandchild) + finally: + _kill(leader.pid) + + +def test_forgetting_a_leader_with_no_group_left_drops_it(tmp_path, monkeypatch): + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + pl._tracked_pids.clear() + pl._tracked_pgids.clear() + + child = subprocess.Popen([sys.executable, "-c", "pass"], start_new_session = True) + child.wait(timeout = 30) + pl.adopt_pid(child.pid) + pl.forget_pid(child.pid) + assert child.pid not in pl._tracked_pids + assert child.pid not in pl._tracked_pgids + + +def test_the_windows_backstop_takes_the_whole_tree(tmp_path, monkeypatch): + """Same gap the startup sweep had: killing the leader alone strands its + workers, and clear_breadcrumb removes the record straight after.""" + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + monkeypatch.setattr(pl, "_is_windows", lambda: True) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: pid == 4242) + monkeypatch.setattr(pl, "_pid_is_zombie", lambda pid: False) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: "same") + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: "same") + trees = [] + monkeypatch.setattr(pl, "_windows_terminate_tree", trees.append) + killed_singly = [] + monkeypatch.setattr(pl.os, "kill", lambda pid, sig: killed_singly.append(pid)) + + pl._tracked_pids.clear() + pl._tracked_pgids.clear() + pl._tracked_pids[4242] = "same" + pl.terminate_all(timeout = 1.0) + assert trees == [4242] + assert killed_singly == [], "the leader was signalled on its own" + + +def test_the_windows_identity_probe_prototypes_every_handle_call(): + """Without argtypes ctypes marshals a 64-bit HANDLE as c_int, so the handle + is truncated and the probe leaks one per call.""" + import ast + import inspect + + from utils import process_lifetime as pl + + source = inspect.getsource(pl._pid_identity) + tree = ast.parse(source) + prototyped = { + node.targets[0].value.attr + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and isinstance(node.targets[0], ast.Attribute) + and node.targets[0].attr in ("argtypes", "restype") + and isinstance(node.targets[0].value, ast.Attribute) + } + for call in ("OpenProcess", "GetProcessTimes", "CloseHandle"): + assert call in prototyped, f"{call} is called without a signature" + + +def test_a_retained_child_keeps_its_group(tmp_path, monkeypatch): + """terminate_all pops the group before deciding; a pid put back without it + leaves nothing able to reach a descendant once the leader exits.""" + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + pl._tracked_pids.clear() + pl._tracked_pgids.clear() + + # Unverifiable identity, the transient `ps` failure case. + monkeypatch.setattr(pl, "_pid_alive", lambda pid: True) + monkeypatch.setattr(pl, "_pid_is_zombie", lambda pid: False) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: None) + pl._tracked_pids[4242] = "recorded" + pl._tracked_pgids[4242] = 4242 + pl.terminate_all(timeout = 1.0) + assert pl._tracked_pids.get(4242) == "recorded" + assert pl._tracked_pgids.get(4242) == 4242, "the group was dropped with the retained pid" + + # And a confirmed survivor keeps it too. + monkeypatch.setattr(pl, "_pid_identity", lambda pid: "recorded") + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: "recorded") + monkeypatch.setattr(pl, "_posix_terminate", lambda pid, timeout = 5.0: None) + monkeypatch.setattr(pl, "_windows_terminate_tree", lambda pid: None) + assert pl.terminate_all(timeout = 1.0) == [4242] + assert pl._tracked_pgids.get(4242) == 4242 + + +def test_a_worker_record_is_revisited_after_its_owner_dies(tmp_path, monkeypatch): + """A worker's record is skipped while its owner lives, and that owner can be + terminated later in the same sweep by the backend's own record.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + directory = tmp_path / "children" + directory.mkdir(parents = True) + + alive = {9000, 9001} # 9000 is the worker, 9001 its llama-server + monkeypatch.setattr(pl, "_pid_alive", lambda pid: pid in alive) + monkeypatch.setattr(pl, "_pid_is_zombie", lambda pid: False) + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: "same" if pid in alive else None) + monkeypatch.setattr(pl, "_is_windows", lambda: False) + + def terminate(pid, timeout = 5.0): + alive.discard(pid) + + monkeypatch.setattr(pl, "_posix_terminate", terminate) + + # Sorted lexicographically, so the worker's record is read first. + (directory / "9000.json").write_text( + json.dumps( + { + "owner_pid": 9000, + "owner_identity": "same", + "children": [{"pid": 9001, "identity": "same"}], + } + ), + encoding = "utf-8", + ) + (directory / "9999.json").write_text( + json.dumps( + { + "owner_pid": 9999, + "owner_identity": "a-backend-that-crashed", + "children": [{"pid": 9000, "identity": "same"}], + } + ), + encoding = "utf-8", + ) + + killed = pl.reap_recorded_children(timeout = 1.0) + assert sorted(killed) == [9000, 9001], killed + assert list(directory.glob("*.json")) == [] + + +def test_a_group_that_survives_sigkill_keeps_its_record(monkeypatch): + """Reporting it resolved deletes the last handle on something still holding + the GPU.""" + from utils import process_lifetime as pl + + if pl._is_windows(): + pytest.skip("POSIX only") + + def killpg(pgid, sig): + if sig == 0: + return None # always still there + if sig == signal.SIGKILL: + raise PermissionError("not permitted") + + monkeypatch.setattr(pl.os, "killpg", killpg) + assert pl._reap_orphaned_group(4242, 4242, timeout = 0.2) is False + + +def test_a_failed_taskkill_falls_through_to_the_leader(monkeypatch): + """check=False does not raise, so the status is the only signal.""" + import subprocess as sp + + from utils import process_lifetime as pl + + monkeypatch.setattr(pl.os, "kill", lambda pid, sig: fallbacks.append(pid)) + fallbacks = [] + + class _Result: + returncode = 1 + + monkeypatch.setattr(sp, "run", lambda *a, **k: _Result()) + pl._windows_terminate_tree(4242) + assert fallbacks == [4242], "a failed taskkill was treated as success" + + class _Ok: + returncode = 0 + + fallbacks.clear() + monkeypatch.setattr(sp, "run", lambda *a, **k: _Ok()) + pl._windows_terminate_tree(4242) + assert fallbacks == [] + + +def test_a_zombie_owner_does_not_shield_its_children(tmp_path, monkeypatch): + """os.kill(pid, 0) succeeds for a zombie, and the sweep runs once, so its + sidecars would survive until another restart.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + directory = tmp_path / "children" + directory.mkdir(parents = True) + + monkeypatch.setattr(pl, "_is_windows", lambda: False) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: True) + monkeypatch.setattr(pl, "_pid_is_zombie", lambda pid: pid == 9500) + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: "same") + monkeypatch.setattr(pl, "_pid_identity", lambda pid: "same") + terminated = [] + monkeypatch.setattr(pl, "_posix_terminate", lambda pid, timeout = 5.0: terminated.append(pid)) + + (directory / "9500.json").write_text( + json.dumps( + { + "owner_pid": 9500, + "owner_identity": "same", + "children": [{"pid": 9501, "identity": "same"}], + } + ), + encoding = "utf-8", + ) + + pl.reap_recorded_children(timeout = 1.0) + assert terminated == [9501], "a zombie Studio kept its orphans alive" + + +def test_a_group_that_cannot_be_taken_keeps_its_record(tmp_path, monkeypatch): + """_reap_orphaned_group returning False has to leave the breadcrumb, or the + next startup has no handle on the survivors.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + directory = tmp_path / "children" + directory.mkdir(parents = True) + + monkeypatch.setattr(pl, "_is_windows", lambda: False) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: False) + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: None) + monkeypatch.setattr(pl, "_reap_orphaned_group", lambda pgid, pid, timeout: False) + monkeypatch.setattr(pl, "_group_has_members", lambda pgid: True) + + record = directory / "9600.json" + record.write_text( + json.dumps( + { + "owner_pid": 9600, + "owner_identity": "gone", + "children": [{"pid": 9601, "identity": "same", "pgid": 9601}], + } + ), + encoding = "utf-8", + ) + + pl._reap_one_record(record, timeout = 1.0) + assert record.is_file(), "the last handle on a live group was deleted" + + +def test_the_backstop_keeps_a_group_it_could_not_take(monkeypatch): + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_pid_alive", lambda pid: False) + monkeypatch.setattr(pl, "_reap_orphaned_group", lambda pgid, pid, timeout: False) + monkeypatch.setattr(pl, "_group_has_members", lambda pgid: True) + pl._tracked_pids.clear() + pl._tracked_pgids.clear() + pl._tracked_pids[9700] = "same" + pl._tracked_pgids[9700] = 9700 + + assert pl.terminate_all(timeout = 1.0) == [9700] + assert pl._tracked_pids.get(9700) == "same" + assert pl._tracked_pgids.get(9700) == 9700 + + +def test_every_child_starting_path_checks_the_rearm(): + """The backend updater is a child too, and a retry re-enters installUpdate.""" + hook = ( + Path(__file__).resolve().parents[2] / "frontend" / "src" / "hooks" / "use-tauri-update.ts" + ) + source = hook.read_text(encoding = "utf-8") + for block in source.split("async function ")[1:]: + name = block.split("(")[0] + if 'invoke("start_server"' in block or 'invoke("start_backend_update"' in block: + assert "crashCleanupReady()" in block, f"{name} starts a child without the check" + + +def test_a_tool_subprocess_is_recorded_while_it_runs(tmp_path, monkeypatch): + """A force quit mid-call on macOS would otherwise leave a session-leading + tool with nothing able to find it.""" + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + monkeypatch.setenv("UNSLOTH_STUDIO_SANDBOX_HOME", str(tmp_path / "sb")) + + from core.inference import tools + from utils import process_lifetime as pl + + pl._tracked_pids.clear() + pl._tracked_pgids.clear() + tools._workdirs.clear() + + seen = [] + real_adopt = pl.adopt_pid + + def watching_adopt(pid): + seen.append(pid) + return real_adopt(pid) + + monkeypatch.setattr(pl, "adopt_pid", watching_adopt) + tools._python_exec("print('hi')", session_id = "__LOCALID_adopt01") + assert seen, "the tool subprocess was never recorded" + # And it is not left on the record once it has exited. + assert all(pid not in pl._tracked_pids for pid in seen), pl._tracked_pids + + +def test_a_terminated_leader_leaving_a_group_keeps_its_record(tmp_path, monkeypatch): + """The dead-leader branch covered only a leader that was already gone.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + directory = tmp_path / "children" + directory.mkdir(parents = True) + + monkeypatch.setattr(pl, "_is_windows", lambda: False) + # Alive on the way in, gone after the terminate, but its group lives on. + states = iter([True, True, False, False, False]) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: next(states, False)) + monkeypatch.setattr(pl, "_pid_is_zombie", lambda pid: False) + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: "same") + monkeypatch.setattr(pl, "_posix_terminate", lambda pid, timeout = 5.0: None) + monkeypatch.setattr(pl, "_group_has_members", lambda pgid: pgid == 9801) + + record = directory / "9800.json" + record.write_text( + json.dumps( + { + "owner_pid": 9800, + "owner_identity": "gone", + "children": [{"pid": 9801, "identity": "same", "pgid": 9801}], + } + ), + encoding = "utf-8", + ) + + pl._reap_one_record(record, timeout = 1.0) + assert record.is_file(), "the group outlived its leader and lost its record" + + +def test_cloudflared_is_adopted_before_stop_can_reach_it(): + """Publishing _proc first lets a concurrent stop reap and forget it, and the + adoption would then record whatever inherited the pid.""" + import inspect + + from cloudflare_tunnel import CloudflareTunnel + + source = inspect.getsource(CloudflareTunnel) + start = source.index("_adopt_pid(proc.pid)") + publish = source.index("self._proc = proc") + assert start < publish, "_proc is published before the adoption" + # Both inside the same `with self._lock:` block. + lock = source.rindex("with self._lock:", 0, start) + assert source.index("self._proc = proc", lock) < source.index("threading.Thread(", lock) + + +def test_a_group_of_only_zombies_is_not_alive(): + """killpg(pgid, 0) succeeds for a zombie leader, so the plain probe would + keep a finished group's record forever.""" + from utils import process_lifetime as pl + + if pl._is_windows(): + pytest.skip("POSIX only") + + child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], start_new_session = True + ) + try: + assert pl._group_member_pids(child.pid) == [child.pid] + assert pl._group_has_members(child.pid) is True + child.kill() + time.sleep(0.3) + assert pl._pid_is_zombie(child.pid) is True + assert pl._group_has_members(child.pid) is False + finally: + _kill(child.pid) + child.wait(timeout = 10) + + +def test_a_zombie_child_is_not_signalled_or_waited_on(tmp_path, monkeypatch): + """Under a non-reaping PID 1 a zombie answers every probe, so terminating it + burns the full grace period per record and reaps nothing.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + directory = tmp_path / "children" + directory.mkdir(parents = True) + + monkeypatch.setattr(pl, "_is_windows", lambda: False) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: True) + monkeypatch.setattr(pl, "_pid_is_zombie", lambda pid: pid == 9901) + monkeypatch.setattr(pl, "_identity_or_none", lambda pid: None) + monkeypatch.setattr(pl, "_group_has_members", lambda pgid: False) + terminated = [] + monkeypatch.setattr(pl, "_posix_terminate", lambda pid, timeout = 5.0: terminated.append(pid)) + + record = directory / "9900.json" + record.write_text( + json.dumps( + { + "owner_pid": 9900, + "owner_identity": "gone", + "children": [{"pid": 9901, "identity": "same", "pgid": 9901}], + } + ), + encoding = "utf-8", + ) + + started = time.monotonic() + pl._reap_one_record(record, timeout = 5.0) + assert terminated == [], "a zombie was signalled" + assert time.monotonic() - started < 1.0, "startup waited out the grace period" + + +def test_the_component_installer_is_recorded_while_it_runs(): + """It rewrites files in place, so one surviving a crash overlaps the next + launch.""" + import ast + import inspect + + from utils.prebuilt import update_flow + + tree = ast.parse(inspect.getsource(update_flow.stream_installer)) + body = ast.dump(tree) + assert "adopt_pid" in body, "the installer is never recorded" + assert "forget_pid" in body, "the record is never cleared" + + +def test_a_fork_child_does_not_inherit_a_held_record_lock(): + """A fork while another thread was adopting leaves it held here with nobody + to release it, and the first adoption after that blocks forever.""" + from utils import process_lifetime as pl + + held_spawner = pl._spawner_lock + held_record = pl._record_lock + held_record.acquire() + try: + pl._reset_after_fork() + assert pl._record_lock is not held_record, "the child kept the locked record lock" + assert pl._spawner_lock is not held_spawner + assert pl._record_lock.acquire(blocking = False), "the fresh lock is not free" + pl._record_lock.release() + finally: + held_record.release() + + +def test_the_update_hook_asks_the_native_side_after_a_remount(): + """A webview reload rebuilds the hook with its ref back at true while the + native job can still be disarmed.""" + hook = ( + Path(__file__).resolve().parents[2] / "frontend" / "src" / "hooks" / "use-tauri-update.ts" + ).read_text(encoding = "utf-8") + assert "cleanupCheckedRef" in hook + gate = hook[hook.index("async function crashCleanupReady") :] + gate = gate[: gate.index("\n }")] + assert '"desktop_update_cleanup_armed"' in gate, "the gate still trusts the ref alone" + assert gate.index("cleanupCheckedRef") < gate.index( + "if (cleanupRearmedRef.current) return true" + ) + + tauri = Path(__file__).resolve().parents[2] / "src-tauri" / "src" + assert "desktop_update_cleanup_armed" in (tauri / "main.rs").read_text(encoding = "utf-8") + assert "kill_on_close_armed" in (tauri / "windows_job.rs").read_text(encoding = "utf-8") + + +def test_a_transient_identity_failure_is_not_recorded_as_none(monkeypatch): + """None is permanent: an entry with no identity is never signalled, so that + child would survive every later launch.""" + from utils import process_lifetime as pl + + answers = iter([None, None, "started-at"]) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: next(answers, "started-at")) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: True) + assert pl._identity_for_record(os.getpid()) == "started-at" + + +def test_identity_capture_stops_once_the_child_is_gone(monkeypatch): + from utils import process_lifetime as pl + + calls = [] + + def _identity(pid): + calls.append(pid) + return None + + monkeypatch.setattr(pl, "_pid_identity", _identity) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: False) + assert pl._identity_for_record(4242) is None + assert calls == [4242], "kept retrying a process that had already exited" + + +def test_an_unanswerable_cleanup_query_counts_as_disarmed(): + hook = ( + Path(__file__).resolve().parents[2] / "frontend" / "src" / "hooks" / "use-tauri-update.ts" + ).read_text(encoding = "utf-8") + gate = hook[hook.index("async function crashCleanupReady") :] + gate = gate[: gate.index("\n }")] + assert "cleanupRearmedRef.current = !isTauri;" in gate, "a failed query still reads as armed" + + +def test_the_component_installer_stays_in_the_backend_group(): + """The desktop stop path force-kills this backend's group, so a session of + its own would let the installer keep rewriting files after the app reports + the backend stopped.""" + import inspect + + from utils.prebuilt import update_flow + + source = inspect.getsource(update_flow.stream_installer) + assert "start_new_session = " not in source + # Still recorded, which is what the macOS sweep has to work from. + spawn = source.index("subprocess.Popen(") + assert "adopt_pid(proc.pid)" in source[spawn:] + + +def test_the_owner_identity_is_retried_and_then_kept(monkeypatch): + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_owner_identity", None) + answers = iter([None, "started-at"]) + calls = [] + + def _identity(pid): + calls.append(pid) + return next(answers, "started-at") + + monkeypatch.setattr(pl, "_pid_identity", _identity) + monkeypatch.setattr(pl, "_pid_alive", lambda pid: True) + assert pl._own_identity() == "started-at" + assert len(calls) == 2, calls + # Kept, so the record lock never waits on another probe. + assert pl._own_identity() == "started-at" + assert len(calls) == 2, calls + + +def test_a_fork_child_does_not_keep_the_parents_identity(monkeypatch): + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_owner_identity", "the-parents") + pl._reset_after_fork() + assert pl._owner_identity is None, "the child would record its parent's identity" + + +def test_a_fork_child_does_not_claim_the_parents_children(tmp_path, monkeypatch): + """Its record would name processes it never started, and a later startup + would reap them while their real parent is still running.""" + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_tracked_pids", {4321: "started-at"}) + monkeypatch.setattr(pl, "_tracked_pgids", {4321: 4321}) + pl._reset_after_fork() + assert pl._tracked_pids == {} + assert pl._tracked_pgids == {} + + +def test_a_missing_child_identity_is_filled_in_later(tmp_path, monkeypatch): + """None is permanent otherwise: nothing signals an entry it cannot verify, + so that child outlives every shutdown.""" + import json + + from utils import process_lifetime as pl + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path / "children")) + monkeypatch.setattr(pl, "_tracked_pids", {5150: None}) + monkeypatch.setattr(pl, "_tracked_pgids", {}) + monkeypatch.setattr(pl, "_owner_identity", "owner") + monkeypatch.setattr(pl, "_pid_alive", lambda pid: True) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: "found-later") + + pl._write_breadcrumb() + record = pl._breadcrumb_file() + written = json.loads(record.read_text(encoding = "utf-8")) + assert written["children"][0]["identity"] == "found-later" + # Written back, so it costs one probe rather than one per write. + assert pl._tracked_pids[5150] == "found-later" + + +def test_an_exited_child_is_not_probed_again(monkeypatch): + from utils import process_lifetime as pl + + calls = [] + monkeypatch.setattr(pl, "_pid_alive", lambda pid: False) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: calls.append(pid)) + assert pl._refreshed_identity(9999, None) is None + assert calls == [] + + +def test_adoption_installs_the_fork_reset(monkeypatch): + """The Linux spawn path is the only other place that registers it, so on + macOS nothing would have.""" + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_fork_reset_installed", False) + registered = [] + monkeypatch.setattr(os, "register_at_fork", lambda **kw: registered.append(kw)) + monkeypatch.setattr(pl, "_pid_identity", lambda pid: "started-at") + monkeypatch.setattr(pl, "_own_process_group", lambda pid: None) + monkeypatch.setattr(pl, "_write_breadcrumb", lambda: None) + monkeypatch.setattr(pl, "_is_windows", lambda: False) + + pl.adopt_pid(os.getpid()) + assert registered and registered[0]["after_in_child"] is pl._reset_after_fork + pl.forget_pid(os.getpid()) + + +def test_a_group_of_nothing_but_a_zombie_costs_no_grace_period(monkeypatch): + """killpg(pgid, 0) answers for a zombie, and where pid 1 does not reap it + stays that way, so every stale record would wait out the timeout.""" + import time as _time + + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_is_windows", lambda: False) + monkeypatch.setattr(os, "killpg", lambda pgid, sig: None) + monkeypatch.setattr(pl, "_group_has_members", lambda pgid: False) + + started = _time.monotonic() + assert pl._reap_orphaned_group(4711, 4711, timeout = 5.0) is False + assert _time.monotonic() - started < 1.0, "waited out the grace period" + + +def test_an_unanswerable_group_query_is_not_an_empty_group(monkeypatch): + """Reporting empty for a failed ps lets forget_pid drop the only record of + a live descendant.""" + import subprocess as _subprocess + + from utils import process_lifetime as pl + + monkeypatch.setattr(pl, "_is_linux", lambda: False) + monkeypatch.setattr(pl, "_is_windows", lambda: False) + monkeypatch.setattr(sys, "platform", "darwin") + monkeypatch.setattr(os, "killpg", lambda pgid, sig: None) + + class _Failed: + returncode = 1 + stdout = "" + + monkeypatch.setattr(_subprocess, "run", lambda *a, **k: _Failed()) + assert pl._group_member_pids(4711) is None + assert pl._group_has_members(4711) is True, "a failed query read as gone" + + +def test_the_visual_server_goes_down_with_an_ordinary_stop(): + """The desktop shutdown only waits for this path, so what the shim started + has to be taken down here rather than by the next startup sweep.""" + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + source = inspect.getsource(LlamaCppBackend._kill_process) + assert "_collect_descendants" in source + assert "_terminate_descendants" in source + # Named before the wait: the shim's children are reparented once it exits. + assert source.index("_collect_descendants") < source.index("self._process.terminate()") + assert source.index("self._process.terminate()") < source.index("_terminate_descendants") + + +@pytest.mark.skipif(os.name == "nt", reason = "posix termination path") +def test_an_exited_child_does_not_burn_the_shutdown_timeout(): + """A child that exited but has not been waited on answers signal 0 like a + live one, so the wait would spend its whole budget on a dead process. With + several tracked, that is the whole shutdown.""" + from utils.process_lifetime import _pid_is_zombie, _posix_terminate + + proc = subprocess.Popen([sys.executable, "-c", "pass"]) + try: + for _ in range(100): + if _pid_is_zombie(proc.pid): + break + time.sleep(0.05) + assert _pid_is_zombie(proc.pid), "the child never became reapable" + + started = time.monotonic() + _posix_terminate(proc.pid, timeout = 5.0) + elapsed = time.monotonic() - started + assert elapsed < 2.0, f"waited {elapsed:.1f}s on a process that had exited" + # Still ours to reap: the wait must not consume the exit status the + # owner is about to read. + assert proc.wait(timeout = 5) == 0 + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix termination path") +def test_a_running_child_is_still_terminated(): + """The shortcut above must not stop the ordinary kill.""" + from utils.process_lifetime import _posix_terminate + + proc = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"]) + try: + started = time.monotonic() + _posix_terminate(proc.pid, timeout = 5.0) + assert time.monotonic() - started < 5.0 + assert proc.wait(timeout = 5) != 0, "exited on its own, so this proves nothing" + finally: + if proc.poll() is None: + proc.kill() + proc.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_a_group_is_waited_on_while_a_member_is_still_running(): + """The leader going first must not cut the wait short: the group is what + holds the GPU, and a live member is the reason for the SIGKILL.""" + from utils.process_lifetime import _posix_terminate + + leader = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, time;" + "subprocess.Popen([sys.executable, '-c', 'import signal, time;" + "signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(30)']);" + "time.sleep(30)", + ], + start_new_session = True, + ) + try: + time.sleep(1.0) # let the grandchild exist before anything is signalled + started = time.monotonic() + _posix_terminate(leader.pid, timeout = 2.0) + elapsed = time.monotonic() - started + assert elapsed >= 1.5, f"gave up after {elapsed:.1f}s with a member still running" + finally: + try: + os.killpg(leader.pid, signal.SIGKILL) + except Exception: + pass + if leader.poll() is None: + leader.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_a_live_leader_answers_without_scanning_every_process(monkeypatch): + """Enumerating a group reads the state of every process on the machine, and + this runs on each stop; a running leader already settles the question.""" + from utils import process_lifetime as lifetime + + leader = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + start_new_session = True, + ) + try: + for _ in range(50): + if lifetime._pid_alive(leader.pid): + break + time.sleep(0.05) + + def refuse(pgid): + raise AssertionError("scanned every process for a group whose leader is alive") + + monkeypatch.setattr(lifetime, "_group_member_pids", refuse) + assert lifetime._group_has_members(leader.pid) is True + finally: + leader.kill() + leader.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_a_group_outliving_its_leader_is_still_found(): + """The short-circuit must not become the whole answer: the case this record + exists for is a leader that exited while its child holds the GPU.""" + from utils import process_lifetime as lifetime + + leader = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys;" + "subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)']);", + ], + start_new_session = True, + ) + survivor = None + try: + leader.wait(timeout = 30) # the leader goes, its child stays + members = lifetime._group_member_pids(leader.pid) or [] + survivor = next((pid for pid in members if pid != leader.pid), None) + assert survivor, members + assert lifetime._group_has_members(leader.pid) is True + finally: + try: + os.killpg(leader.pid, signal.SIGKILL) + except Exception: + pass + if leader.poll() is None: + leader.kill() + leader.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_forgetting_an_untracked_pid_costs_nothing(): + """Every tool call ends here, most of them with nothing recorded: neither + the group scan nor the record rewrite has anything to do.""" + from utils import process_lifetime as lifetime + + scanned = [] + written = [] + original_scan = lifetime._group_member_pids + original_write = lifetime._write_breadcrumb + lifetime._group_member_pids = lambda pgid: (scanned.append(pgid), original_scan(pgid))[1] + lifetime._write_breadcrumb = lambda: (written.append(1), original_write())[1] + try: + lifetime.forget_pid(999_001) + finally: + lifetime._group_member_pids = original_scan + lifetime._write_breadcrumb = original_write + + assert scanned == [] + assert written == [], "rewrote the record for a pid it never held" + + +def test_a_validation_server_the_installer_started_is_recorded(monkeypatch, tmp_path): + """It is a grandchild: a parent-death signal reaches the installer only, and + the record of the installer pid alone never finds the server holding the GPU.""" + from utils.prebuilt import update_flow + + adopted, forgotten = [], [] + monkeypatch.setattr(update_flow, "adopt_pid", lambda pid: adopted.append(pid)) + monkeypatch.setattr(update_flow, "forget_pid", lambda pid: forgotten.append(pid)) + + class FakeProc: + pid = 4321 + stdout = iter( + [ + "downloading 10% (1/10)\n", + "UNSLOTH_INSTALLER_CHILD started 9911\n", + "validating\n", + "UNSLOTH_INSTALLER_CHILD stopped 9911\n", + "downloading 100% (10/10)\n", + ] + ) + + def wait(self): + return 0 + + def poll(self): + return 0 + + def kill(self): + pass + + monkeypatch.setattr(update_flow.subprocess, "Popen", lambda *a, **k: FakeProc()) + update_flow.stream_installer( + ["x"], + {}, + timeout_seconds = 30, + job = {}, + job_lock = threading.Lock(), + ) + + assert adopted == [4321, 9911], adopted # the installer, then its server + # Dropped on the stop line, before the installer's own record goes. + assert forgotten[0] == 9911, forgotten + + +def test_a_validation_server_left_running_stays_recorded(monkeypatch): + """The installer killed before it could report the stop is exactly the case + the record exists for, so nothing may drop it.""" + from utils.prebuilt import update_flow + + adopted, forgotten = [], [] + monkeypatch.setattr(update_flow, "adopt_pid", lambda pid: adopted.append(pid)) + monkeypatch.setattr(update_flow, "forget_pid", lambda pid: forgotten.append(pid)) + + class FakeProc: + pid = 4321 + stdout = iter(["UNSLOTH_INSTALLER_CHILD started 9912\n"]) + + def wait(self): + return 0 + + def poll(self): + return 0 + + def kill(self): + pass + + monkeypatch.setattr(update_flow.subprocess, "Popen", lambda *a, **k: FakeProc()) + update_flow.stream_installer( + ["x"], + {}, + timeout_seconds = 30, + job = {}, + job_lock = threading.Lock(), + ) + + assert adopted == [4321, 9912], adopted # the installer, then its server + assert 9912 not in forgotten, "dropped the only handle on a running server" + + +def test_the_installer_announces_and_groups_its_validation_server(): + """The other half of the handoff, in the script Studio runs.""" + source = (Path(__file__).resolve().parents[2] / "install_llama_prebuilt.py").read_text( + encoding = "utf-8" + ) + + assert '_announce_child("started", process.pid)' in source + assert '_announce_child("stopped", process.pid)' in source + # Its own group, so a server that starts something of its own is still + # reachable through the one pid that gets recorded. + assert '"start_new_session": True' in source + assert "_terminate_validation_server" in source + assert "killpg" in source + + +def test_a_backend_never_starts_under_a_disarmed_job(): + """The UI gate runs per update action, but a webview remount starts a backend + on its own, so the check belongs where the spawn happens.""" + process_rs = ( + Path(__file__).resolve().parents[2] / "src-tauri" / "src" / "process.rs" + ).read_text(encoding = "utf-8") + + start = process_rs.index("pub fn start_backend(") + spawn = process_rs.index("resolve_backend_binary()", start) + gate = process_rs.index("kill_on_close_armed()", start) + assert gate < spawn, "the armed check runs after the backend is resolved and spawned" + assert "resume_after_update_installer()" in process_rs[start:spawn] + # Fail closed: a re-arm that will not take must stop the start. + assert "Refusing to start the backend with crash cleanup disarmed" in process_rs + + +def test_one_malformed_record_does_not_stop_the_whole_sweep(tmp_path, monkeypatch): + """A record is a file on disk: anything that parses as JSON reaches the + identity comparison, and the sweep runs inside one try in run_server.""" + from utils import process_lifetime as lifetime + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path)) + victim = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(60)"], + start_new_session = True, + ) + # Live, so its identity is readable and the record's is what gets compared + # against it. Unverifiable, so it must survive the sweep. + decoy = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(60)"], + start_new_session = True, + ) + try: + for index, shape in enumerate((123, ["a"], {"b": 1}, None)): + (tmp_path / f"bad_{index}.json").write_text( + json.dumps( + { + "owner_pid": 999_001, + "owner_identity": shape, + "children": [{"pid": decoy.pid, "identity": shape, "pgid": None}], + } + ), + encoding = "utf-8", + ) + # A real orphan behind them, from an owner that is gone. + (tmp_path / "good.json").write_text( + json.dumps( + { + "owner_pid": 999_003, + "owner_identity": "1", + "children": [ + { + "pid": victim.pid, + "identity": lifetime._pid_identity(victim.pid), + "pgid": victim.pid, + } + ], + } + ), + encoding = "utf-8", + ) + + killed = lifetime.reap_recorded_children(timeout = 5.0) + assert victim.pid in killed, killed + victim.wait(timeout = 10) + assert decoy.poll() is None, "signalled a pid it could not verify" + finally: + for proc in (victim, decoy): + if proc.poll() is None: + proc.kill() + proc.wait(timeout = 5) + + +def test_a_malformed_identity_is_never_treated_as_a_match(): + """Unverifiable has to stay unverifiable: an identity that compares equal to + anything would signal whatever holds that pid now.""" + from utils.process_lifetime import _recorded_identity, _same_identity + + for shape in (123, ["1"], {"a": 1}, None, ""): + assert _recorded_identity(shape) is None, shape + assert _same_identity(shape, "1") is False + assert _same_identity("1", shape) is False + assert _recorded_identity("1") == "1" + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_a_group_of_zombies_does_not_hold_up_the_startup_sweep(monkeypatch): + """Where pid 1 does not reap (a container), a member that exits on the + SIGTERM keeps answering killpg(pgid, 0) for the rest of the grace period.""" + from utils import process_lifetime as lifetime + + leader = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + start_new_session = True, + ) + try: + # As a non-reaping pid 1 leaves things: signalled, exited, still listed. + alive = {"value": True} + real_members = lifetime._group_member_pids + monkeypatch.setattr( + lifetime, + "_group_member_pids", + lambda pgid: real_members(pgid) if alive["value"] else [], + ) + monkeypatch.setattr(lifetime, "_pid_is_zombie", lambda pid: not alive["value"]) + + def fake_killpg(pgid, sig): + if sig == signal.SIGTERM: + alive["value"] = False # it exited, and nobody waited on it + return None # the pid is still in the table, so the probe succeeds + + monkeypatch.setattr(lifetime.os, "killpg", fake_killpg) + + started = time.monotonic() + assert lifetime._reap_orphaned_group(leader.pid, leader.pid, 5.0) is True + elapsed = time.monotonic() - started + assert elapsed < 2.0, f"waited {elapsed:.1f}s on a group that had already exited" + finally: + leader.kill() + leader.wait(timeout = 5) + + +def test_rearming_on_backend_start_also_clears_the_cleanup_guard(): + """Re-enabling the job is only half of it: the pre-exit hook has already run + its cleanup, and the guard left set makes the next attempt suspend + kill-on-close without stopping this backend first.""" + process_rs = ( + Path(__file__).resolve().parents[2] / "src-tauri" / "src" / "process.rs" + ).read_text(encoding = "utf-8") + + start = process_rs.index("pub fn start_backend(") + spawn = process_rs.index("resolve_backend_binary()", start) + guard = process_rs[start:spawn] + assert "resume_after_update_installer()" in guard + assert "reset_termination_cleanup()" in guard, "re-armed the job but left the guard set" + assert guard.index("resume_after_update_installer()") < guard.index( + "reset_termination_cleanup()" + ) + + +def test_a_tree_taskkill_could_not_take_keeps_its_record(tmp_path, monkeypatch): + """This fallback runs when the job object was unavailable, so the record is + the only handle on those workers: reading the dead leader as the tree being + gone would strand them for good.""" + from utils import process_lifetime as lifetime + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path)) + monkeypatch.setattr(lifetime, "_is_windows", lambda: True) + monkeypatch.setattr(lifetime, "_identity_or_none", lambda pid: "1") + monkeypatch.setattr(lifetime, "_pid_is_zombie", lambda pid: False) + monkeypatch.setattr(lifetime, "_reap_orphaned_group", lambda pgid, pid, timeout: False) + monkeypatch.setattr(lifetime, "_group_has_members", lambda pgid: False) + + # Alive until it is signalled, gone afterwards: the leader dies, and only + # taskkill's exit status says whether its workers went with it. + state = {"leader": True, "tree": False} + monkeypatch.setattr(lifetime, "_pid_alive", lambda pid: state["leader"]) + + def fake_tree(pid): + state["leader"] = False + return state["tree"] + + monkeypatch.setattr(lifetime, "_windows_terminate_tree", fake_tree) + + def write_record(): + (tmp_path / "previous.json").write_text( + json.dumps( + { + "owner_pid": 999_101, + "owner_identity": "an-owner-that-is-gone", + "children": [{"pid": 999_102, "identity": "1", "pgid": None}], + } + ), + encoding = "utf-8", + ) + + write_record() + assert lifetime.reap_recorded_children(timeout = 1.0) == [999_102] + assert (tmp_path / "previous.json").exists(), "dropped the only handle on a live tree" + + # And once taskkill takes the tree, the record is consumed. + state.update(leader = True, tree = True) + write_record() + assert lifetime.reap_recorded_children(timeout = 1.0) == [999_102] + assert not (tmp_path / "previous.json").exists() + + +def test_a_failed_tree_kill_keeps_the_pid_in_the_shutdown_record(monkeypatch): + """terminate_all reports what is still up so the caller can keep it in the + crash record; a leader-only kill leaves the workers with nothing naming + them.""" + from utils import process_lifetime as lifetime + + monkeypatch.setattr(lifetime, "_is_windows", lambda: True) + monkeypatch.setattr(lifetime, "_pid_identity", lambda pid: "1") + monkeypatch.setattr(lifetime, "_identity_or_none", lambda pid: "1") + monkeypatch.setattr(lifetime, "_pid_is_zombie", lambda pid: False) + monkeypatch.setattr(lifetime, "_group_has_members", lambda pgid: False) + monkeypatch.setattr(lifetime, "_write_breadcrumb", lambda: None) + + state = {"leader": True, "tree": False} + monkeypatch.setattr(lifetime, "_pid_alive", lambda pid: state["leader"]) + + def fake_tree(pid): + state["leader"] = False + return state["tree"] + + monkeypatch.setattr(lifetime, "_windows_terminate_tree", fake_tree) + + def track(): + with lifetime._record_lock: + lifetime._tracked_pids.clear() + lifetime._tracked_pgids.clear() + lifetime._tracked_pids[999_201] = "1" + + track() + assert lifetime.terminate_all(timeout = 1.0) == [999_201] + assert 999_201 in lifetime._tracked_pids, "stopped tracking a tree that may still be up" + + state.update(leader = True, tree = True) + track() + assert lifetime.terminate_all(timeout = 1.0) == [] + assert 999_201 not in lifetime._tracked_pids + + with lifetime._record_lock: + lifetime._tracked_pids.clear() + lifetime._tracked_pgids.clear() + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_the_installer_waits_for_its_validation_group_not_the_leader(tmp_path): + """A child that ignores SIGTERM outlives the leader's exit, and announcing + the stop is what drops the only record of it.""" + installer = _load_installer_module() + + script = tmp_path / "leader.py" + script.write_text( + "import signal, subprocess, sys, time\n" + "signal.signal(signal.SIGTERM, signal.SIG_IGN)\n" + "child = subprocess.Popen([sys.executable, '-c'," + " 'import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN);" + " time.sleep(30)'])\n" + "print(child.pid, flush = True)\n" + "time.sleep(30)\n", + encoding = "utf-8", + ) + proc = subprocess.Popen( + [sys.executable, str(script)], + stdout = subprocess.PIPE, + text = True, + start_new_session = True, + ) + try: + grandchild = int(proc.stdout.readline().strip()) + assert _alive(grandchild) + + gone = installer._terminate_validation_server(proc) + assert gone is True, "reported a stop with the group still up" + time.sleep(0.3) + assert not _alive(grandchild), "the leader went and its child stayed" + finally: + for pid in (proc.pid, locals().get("grandchild")): + if pid: + _kill(pid) + if proc.poll() is None: + proc.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_a_validation_group_that_will_not_die_is_not_announced_as_stopped(monkeypatch): + """False keeps the pid recorded, which is the whole point of announcing it.""" + installer = _load_installer_module() + + proc = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(30)"], + start_new_session = True, + ) + try: + # Signals go nowhere: what an unkillable member looks like from here. + monkeypatch.setattr(installer.os, "killpg", lambda pgid, sig: None) + assert installer._terminate_validation_server(proc, grace = 0.2) is False + finally: + monkeypatch.undo() + proc.kill() + proc.wait(timeout = 5) + + +def test_an_installer_timeout_takes_its_announced_children_with_it(monkeypatch): + """This process keeps running after the error, so no startup sweep is coming + before the retry.""" + from utils.prebuilt import update_flow + + terminated = [] + monkeypatch.setattr(update_flow, "adopt_pid", lambda pid: None) + monkeypatch.setattr(update_flow, "forget_pid", lambda pid: None) + monkeypatch.setattr(update_flow, "terminate_pid", lambda pid: terminated.append(pid)) + + class FakeProc: + pid = 4321 + + def __init__(self): + self.killed = False + self.stdout = self._lines() + + def _lines(self): + yield "UNSLOTH_INSTALLER_CHILD started 9931\n" + while not self.killed: + time.sleep(0.05) + + def wait(self): + return -9 + + def poll(self): + return -9 + + def kill(self): + self.killed = True + + monkeypatch.setattr(update_flow.subprocess, "Popen", lambda *a, **k: FakeProc()) + with pytest.raises(RuntimeError, match = "timed out"): + update_flow.stream_installer( + ["x"], + {}, + timeout_seconds = 1, + job = {}, + job_lock = threading.Lock(), + ) + + assert terminated == [9931], terminated + + +def test_no_process_group_id_is_kept_past_the_group_it_names(): + """A pid is reusable once nothing holds the number as a process group any + more, so a group id cached across an unload eventually names a stranger.""" + source = ( + Path(__file__).resolve().parents[1] / "core" / "inference" / "llama_cpp.py" + ).read_text(encoding = "utf-8") + + assert "_diffusion_pgid" not in source, "a kept group id outlives its group" + + +def test_a_failed_installer_takes_its_announced_children_with_it(monkeypatch): + """A nonzero exit leaves this process running, and its own live record + shields those children from a sweep that will not run anyway.""" + from utils.prebuilt import update_flow + + terminated = [] + monkeypatch.setattr(update_flow, "adopt_pid", lambda pid: None) + monkeypatch.setattr(update_flow, "forget_pid", lambda pid: None) + monkeypatch.setattr(update_flow, "terminate_pid", lambda pid: terminated.append(pid)) + + class FakeProc: + pid = 4321 + stdout = iter( + [ + "UNSLOTH_INSTALLER_CHILD started 9941\n", + "boom\n", + ] + ) + + def wait(self): + return 3 + + def poll(self): + return 3 + + def kill(self): + pass + + monkeypatch.setattr(update_flow.subprocess, "Popen", lambda *a, **k: FakeProc()) + with pytest.raises(update_flow.InstallerExit): + update_flow.stream_installer( + ["x"], + {}, + timeout_seconds = 30, + job = {}, + job_lock = threading.Lock(), + ) + + assert terminated == [9941], terminated + + +def test_a_server_the_installer_reported_stopped_is_not_killed_twice(monkeypatch): + """The ordinary path: it said the group was gone, so there is nothing left + to terminate.""" + from utils.prebuilt import update_flow + + terminated = [] + monkeypatch.setattr(update_flow, "adopt_pid", lambda pid: None) + monkeypatch.setattr(update_flow, "forget_pid", lambda pid: None) + monkeypatch.setattr(update_flow, "terminate_pid", lambda pid: terminated.append(pid)) + + class FakeProc: + pid = 4321 + stdout = iter( + [ + "UNSLOTH_INSTALLER_CHILD started 9942\n", + "UNSLOTH_INSTALLER_CHILD stopped 9942\n", + ] + ) + + def wait(self): + return 0 + + def poll(self): + return 0 + + def kill(self): + pass + + monkeypatch.setattr(update_flow.subprocess, "Popen", lambda *a, **k: FakeProc()) + update_flow.stream_installer( + ["x"], + {}, + timeout_seconds = 30, + job = {}, + job_lock = threading.Lock(), + ) + assert terminated == [] + + +@pytest.mark.skipif(os.name == "nt", reason = "posix process groups") +def test_terminate_pid_takes_the_group_once_its_leader_has_gone(tmp_path, monkeypatch): + """The installer announces a server that has already exited while a child of + its own still holds the GPU. getpgid stops answering for the reaped leader, + so only the recorded group can still reach that child, and this backend + keeps running: no sweep is coming to catch it.""" + from utils import process_lifetime as lifetime + + monkeypatch.setenv("UNSLOTH_STUDIO_CHILD_RECORD", str(tmp_path)) + leader = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys;" + "subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)']);", + ], + start_new_session = True, + ) + survivor = None + try: + lifetime.adopt_pid(leader.pid) # while it still leads a group we can read + assert lifetime._tracked_pgids.get(leader.pid) == leader.pid + leader.wait(timeout = 30) # the leader goes and is reaped; its child stays + members = lifetime._group_member_pids(leader.pid) or [] + survivor = next((pid for pid in members if pid != leader.pid), None) + assert survivor, members + + lifetime.terminate_pid(leader.pid, timeout = 5.0) + for _ in range(100): + if not _alive(survivor): + break + time.sleep(0.05) + assert not _alive(survivor), "left a group member holding the GPU" + finally: + try: + os.killpg(leader.pid, signal.SIGKILL) + except Exception: + pass + if leader.poll() is None: + leader.kill() + leader.wait(timeout = 5) + with lifetime._record_lock: + lifetime._tracked_pids.clear() + lifetime._tracked_pgids.clear() + + +def test_terminate_pid_keeps_a_record_taskkill_could_not_confirm(monkeypatch): + """Windows has no group to fall back on, so a tree kill that failed leaves + the record as the only thing naming those workers.""" + from utils import process_lifetime as lifetime + + monkeypatch.setattr(lifetime, "_is_windows", lambda: True) + monkeypatch.setattr(lifetime, "_group_has_members", lambda pgid: False) + monkeypatch.setattr(lifetime, "_write_breadcrumb", lambda: None) + + state = {"tree": False} + monkeypatch.setattr(lifetime, "_windows_terminate_tree", lambda pid: state["tree"]) + + def track(): + with lifetime._record_lock: + lifetime._tracked_pids.clear() + lifetime._tracked_pgids.clear() + lifetime._tracked_pids[999_301] = "1" + + try: + track() + lifetime.terminate_pid(999_301, timeout = 1.0) + assert 999_301 in lifetime._tracked_pids, "dropped the only handle on a live tree" + + state["tree"] = True # and a confirmed kill still consumes the record + track() + lifetime.terminate_pid(999_301, timeout = 1.0) + assert 999_301 not in lifetime._tracked_pids + finally: + with lifetime._record_lock: + lifetime._tracked_pids.clear() + lifetime._tracked_pgids.clear() + + +def test_announced_children_survive_two_threads_draining_at_once(): + """The watchdog and the reader thread both drain this set, and cancelling a + timer does not stop a callback that already began. A bare + `while announced: announced.pop()` raises KeyError out of whichever thread + loses that race, replacing the installer error the caller should see.""" + import inspect + + from utils.prebuilt import update_flow + + announced = update_flow.AnnouncedChildren() + assert announced.take() is None, "an empty drain must end the loop, not raise" + + pids = list(range(9000, 9500)) + for pid in pids: + announced.add(pid) + + taken: list[int] = [] + failures: list[BaseException] = [] + guard = threading.Lock() + start = threading.Barrier(4) + + def drain(): + try: + start.wait(timeout = 10) + while True: + pid = announced.take() + if pid is None: + return + with guard: + taken.append(pid) + except BaseException as exc: # noqa: BLE001 - reported below + failures.append(exc) + + threads = [threading.Thread(target = drain) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout = 30) + + assert not failures, failures + assert sorted(taken) == pids, "a pid was taken twice or lost" + + # And the installer stream drains through it rather than a bare set. + source = inspect.getsource(update_flow.stream_installer) + assert "AnnouncedChildren()" in source + assert "announced.take()" in source + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason = "PR_SET_PDEATHSIG") +def test_a_validation_server_dies_with_a_killed_installer(tmp_path): + """The server runs in a session of its own, which takes it out of the group + the desktop stop path kills, so it arms the parent-death signal instead: a + SIGKILL mid-validation must not leave it holding the GPU and the staged + files until some later startup sweeps the breadcrumb.""" + studio_dir = Path(__file__).resolve().parents[2] + script = tmp_path / "spawn_validation_server.py" + script.write_text( + "import importlib.util, os, subprocess, sys, time\n" + "studio = sys.argv[1]\n" + "sys.path.insert(0, studio)\n" + "spec = importlib.util.spec_from_file_location(\n" + " 'installer_under_test', os.path.join(studio, 'install_llama_prebuilt.py')\n" + ")\n" + "module = importlib.util.module_from_spec(spec)\n" + "sys.modules['installer_under_test'] = module\n" + "spec.loader.exec_module(module)\n" + "child = subprocess.Popen(\n" + " [sys.executable, '-c', 'import time; time.sleep(120)'],\n" + " **module._validation_server_kwargs(),\n" + ")\n" + "print(child.pid, flush = True)\n" + "time.sleep(120)\n", + encoding = "utf-8", + ) + + installer = subprocess.Popen( + [sys.executable, str(script), str(studio_dir)], + stdout = subprocess.PIPE, + text = True, + ) + server_pid = None + try: + line = installer.stdout.readline().strip() + assert line.isdigit(), f"the installer never started a server: {line!r}" + server_pid = int(line) + assert _alive(server_pid) + + installer.kill() # the crash case: no cooperative shutdown runs + installer.wait(timeout = 10) + for _ in range(200): + if not _alive(server_pid): + break + time.sleep(0.05) + assert not _alive(server_pid), "outlived the installer that started it" + finally: + if installer.poll() is None: + installer.kill() + installer.wait(timeout = 5) + if server_pid is not None and _alive(server_pid): + try: + os.kill(server_pid, signal.SIGKILL) + except OSError: + pass + + +@pytest.mark.skipif(os.name == "nt", reason = "posix termination path") +def test_a_child_in_our_own_group_still_takes_its_children_down(tmp_path): + """killpg is not available for a child that shares this process's group, so + without the walk the diffusion shim dies and the visual server keeps the GPU.""" + from utils import process_lifetime as pl + + marker = tmp_path / "grandchild.pid" + shim = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, time, pathlib\n" + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(60)'])\n" + f"pathlib.Path({str(marker)!r}).write_text(str(child.pid))\n" + "time.sleep(60)\n", + ], + ) # deliberately in this process's group, as the runner now is + grandchild = None + try: + for _ in range(200): + if marker.is_file(): + break + time.sleep(0.05) + grandchild = int(marker.read_text()) + pl._posix_terminate(shim.pid, timeout = 5.0) + for _ in range(100): + if not _alive(grandchild): + break + time.sleep(0.05) + assert not _alive(grandchild), "the visual server survived its runner" + finally: + _kill(shim.pid) + if grandchild is not None: + _kill(grandchild) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix identity check") +def test_terminate_pid_leaves_a_pid_that_is_no_longer_ours_alone(monkeypatch): + """An announced child can exit without the line that clears it, so the pid + reached here may already belong to someone else's tree.""" + from utils import process_lifetime as pl + + stranger = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(60)"]) + try: + monkeypatch.setitem(pl._tracked_pids, stranger.pid, "1") # an impossible start time + monkeypatch.setattr(pl, "_write_breadcrumb", lambda: None) + pl.terminate_pid(stranger.pid, timeout = 1.0) + time.sleep(0.5) + # poll(), not a signal-0 probe: a killed child of this process answers + # that probe as a zombie until it is waited on. + assert stranger.poll() is None, "signalled a pid the record does not match" + finally: + _kill(stranger.pid) + stranger.wait(timeout = 5) + + +@pytest.mark.skipif(os.name == "nt", reason = "posix group handling") +def test_a_validation_server_without_a_death_signal_stays_in_the_group(monkeypatch): + """macOS has no parent-death signal, and the record only exists once the + backend has read the announcement that follows the spawn, so a session of + its own leaves a window in which nothing can reach the server.""" + import importlib + + studio = Path(__file__).resolve().parents[2] + if str(studio) not in sys.path: + sys.path.insert(0, str(studio)) + ilp = importlib.import_module("install_llama_prebuilt") + + monkeypatch.setattr(sys, "platform", "darwin") + assert ilp._validation_server_kwargs() == {}, "no death signal to pair the session with" + + server = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(60)"], + **ilp._validation_server_kwargs(), + ) + signalled: list = [] + monkeypatch.setattr(os, "killpg", lambda pgid, sig: signalled.append((pgid, sig))) + try: + ilp._terminate_validation_server(server, grace = 1.0) + assert signalled == [], f"signalled the group it shares with the installer: {signalled}" + assert server.poll() is not None, "the server was left running" + finally: + monkeypatch.undo() + _kill(server.pid) + server.wait(timeout = 5) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-q", "-s"])) diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 3a332ee19b..18d9723b88 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -33,6 +33,7 @@ def __init__( ): self._alive = alive self.returncode = returncode + self.pid = 424242 # every real Popen has one; the lifetime record reads it self.stdout = iter(()) # drain thread exits immediately def poll(self): diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index a9fca96a0f..5501eec826 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -247,6 +247,7 @@ def __init__( ): type(self).captured_cmd = list(cmd) type(self).captured_env = dict(env or {}) + self.pid = 424242 # a real Popen has one, and the lifetime record needs it self._lines = list(lines) self.returncode = returncode self._out_file = out_file @@ -372,6 +373,7 @@ class _HangingPopen: def __init__(self, cmd, **_kw): self._alive = True + self.pid = 424243 class _Blocking: def __init__(self, owner): diff --git a/studio/backend/utils/prebuilt/update_flow.py b/studio/backend/utils/prebuilt/update_flow.py index 69c1566fc3..31b92a47cf 100644 --- a/studio/backend/utils/prebuilt/update_flow.py +++ b/studio/backend/utils/prebuilt/update_flow.py @@ -25,7 +25,7 @@ import structlog from utils.child_stdio import utf8_child_env -from utils.process_lifetime import child_popen_kwargs +from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid, terminate_pid logger = structlog.get_logger(__name__) @@ -36,6 +36,10 @@ # Matches the installer's download progress lines, e.g. # "Downloading x.zip: 35.0% (12.3 MiB/35.1 MiB) at 8.2 MiB/s". PROGRESS_LINE_RE = re.compile(r"(\d+(?:\.\d+)?)%\s*\(") +# The installer announces each server it starts to validate a build. They are +# grandchildren, so a parent-death signal or a sweep of the installer pid alone +# never reaches them, and one left running holds the GPU and the staged files. +CHILD_PID_LINE_RE = re.compile(r"\AUNSLOTH_INSTALLER_CHILD (started|stopped) (\d+)\Z") # The download dominates the update; extract/validate fill the last slice. DOWNLOAD_PROGRESS_CEILING = 0.95 @@ -282,6 +286,34 @@ def rocm_install_args(asset: Optional[str]) -> list[str]: return ["--has-rocm"] +class AnnouncedChildren: + """The pids the installer reported started, drained one at a time. + + Two threads drain it: the timeout watchdog, and the reader thread in its + `finally` (``Timer.cancel()`` does not stop a callback that has already + begun). A bare ``while pids: pids.pop()`` raises KeyError out of the loser + of that race, replacing the installer error the caller is meant to see, so + emptiness and the take are decided together. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._pids: set[int] = set() + + def add(self, pid: int) -> None: + with self._lock: + self._pids.add(pid) + + def discard(self, pid: int) -> None: + with self._lock: + self._pids.discard(pid) + + def take(self) -> Optional[int]: + """One pid, or None once there are none left.""" + with self._lock: + return self._pids.pop() if self._pids else None + + def stream_installer( cmd: list[str], env: dict[str, str], @@ -310,13 +342,34 @@ def set_progress(fraction: float) -> None: errors = "replace", # Make the Python child emit the UTF-8 we decode above. env = utf8_child_env(env), + # Deliberately NOT start_new_session: the desktop stop path force-kills + # this backend's process group, and a session of its own would take the + # installer out of it, leaving it rewriting files after the app reports + # the backend stopped. **child_popen_kwargs(), ) + # The kwargs above are empty on macOS, so record it: an installer that + # outlives its owner keeps replacing files under the next launch. + adopt_pid(proc.pid) timed_out = threading.Event() + announced = AnnouncedChildren() + + def _stop_announced() -> None: + # This process keeps running after an installer error, so no startup + # sweep is coming and its own record shields these from one anyway: a + # validation server left here holds the GPU and the staged files + # through the retry that follows. + while True: + pid = announced.take() + if pid is None: + return + terminate_pid(pid) + def _kill_on_timeout() -> None: timed_out.set() proc.kill() + _stop_announced() watchdog = threading.Timer(timeout_seconds, _kill_on_timeout) watchdog.daemon = True @@ -328,6 +381,18 @@ def _kill_on_timeout() -> None: tail_lines.append(line) if len(tail_lines) > 80: del tail_lines[0] + child = CHILD_PID_LINE_RE.match(line.strip()) + if child is not None: + # Recorded while it runs and dropped when the installer says it + # stopped; one it never got to report stays for the sweep. + started, child_pid = child.group(1) == "started", int(child.group(2)) + if started: + adopt_pid(child_pid) + announced.add(child_pid) + else: + forget_pid(child_pid) + announced.discard(child_pid) + continue m = PROGRESS_LINE_RE.search(line) if m is None: continue @@ -335,6 +400,11 @@ def _kill_on_timeout() -> None: returncode = proc.wait() finally: watchdog.cancel() + if proc.poll() is not None: + forget_pid(proc.pid) + # Anything it started and never reported as stopped, whether it timed + # out, exited nonzero, or died mid-line. + _stop_announced() if timed_out.is_set(): raise RuntimeError(f"installer timed out after {timeout_seconds}s") if returncode != 0: diff --git a/studio/backend/utils/process_lifetime.py b/studio/backend/utils/process_lifetime.py index 6c48e99394..fa91865421 100644 --- a/studio/backend/utils/process_lifetime.py +++ b/studio/backend/utils/process_lifetime.py @@ -11,10 +11,14 @@ 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 +preexec hook. Linux's signal is per-direct-child only, so multiprocessing workers are also tracked for terminate_all. +macOS has neither mechanism, so tracked children are also recorded on disk and +the next startup sweeps whatever the previous run left behind +(reap_recorded_children). That record is the only reaper macOS has after a +crash, a Force Quit or a closed terminal. + Best-effort throughout: any failure degrades to today's behavior, never raises. Stdlib only. """ @@ -25,6 +29,7 @@ import signal import sys import threading +import time from typing import Callable, Optional _PR_SET_PDEATHSIG = 1 @@ -37,6 +42,56 @@ _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 +# pid -> its own process group, for children started with start_new_session. The +# leader can exit first, and the group is then the only handle on its children. +_tracked_pgids: "dict[int, int]" = {} +# Serialises edit-then-write: two threads adopting at once could otherwise each +# write from its own snapshot, and the older write would drop the newer pid. +_record_lock = threading.Lock() + + +# Whether cleanup-on-abnormal-exit is in force, and why not. A silent failure +# here leaks every child on a crash, so record it and log it. +_win_job_status: "tuple[bool, str]" = (False, "not attempted") + + +def _last_error(ctypes_module) -> int: + # get_last_error is Windows-only; a POSIX probe must not raise here. + getter = getattr(ctypes_module, "get_last_error", None) + try: + return int(getter()) if getter else 0 + except Exception: + return 0 + + +def _record_job_status( + ok: bool, + detail: str, + last_error: int = 0, +) -> None: + global _win_job_status + if last_error: + detail = f"{detail} (WinError {last_error})" + _win_job_status = (ok, detail) + try: + import logging + logger = logging.getLogger(__name__) + if ok: + logger.info("Child-process cleanup on abnormal exit: %s", detail) + else: + logger.warning( + "Child-process cleanup on abnormal exit is NOT guaranteed: %s. Children may " + "survive a crash or a force quit; the startup sweep reaps them on the next " + "launch.", + detail, + ) + except Exception: + pass + + +def windows_job_status() -> "tuple[bool, str]": + """(in_force, detail) for the Windows kill-on-close job.""" + return _win_job_status def _is_linux() -> bool: @@ -63,6 +118,34 @@ def initialize_parent_lifetime() -> None: _initialized = True if _is_windows(): _install_windows_job() + elif _is_linux(): + if _pdeathsig_available(): + _record_job_status(True, "PR_SET_PDEATHSIG per child") + else: + _record_job_status(False, "prctl is unavailable here (seccomp or container policy)") + else: + # macOS: no pdeathsig, no job object; reap_recorded_children covers it. + _record_job_status(False, "no kernel-level parent-death signal on this platform") + + +def _pdeathsig_available() -> bool: + """Whether PR_SET_PDEATHSIG itself works, so the status is not a claim we + cannot keep. + + seccomp can filter prctl on its first argument, so a successful read-only + GET proves nothing about SET. SET is therefore exercised for real, with the + value it already holds: a no-op that still goes through the same filter.""" + _PR_GET_PDEATHSIG, _PR_SET_PDEATHSIG = 2, 1 + try: + import ctypes + + libc = ctypes.CDLL("libc.so.6", use_errno = True) + current = ctypes.c_int(0) + if libc.prctl(_PR_GET_PDEATHSIG, ctypes.byref(current), 0, 0, 0) != 0: + return False + return libc.prctl(_PR_SET_PDEATHSIG, current.value, 0, 0, 0) == 0 + except Exception: + return False def _win_signatures(kernel32) -> None: @@ -132,23 +215,27 @@ class _EXT(ctypes.Structure): job = kernel32.CreateJobObjectW(None, None) if not job: + _record_job_status(False, "CreateJobObjectW failed", _last_error(ctypes)) 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) ): + _record_job_status(False, "SetInformationJobObject failed", _last_error(ctypes)) kernel32.CloseHandle(job) return # AssignProcessToJobObject(parent) makes children inherit the job. May # fail if Unsloth 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()): + _record_job_status(False, "AssignProcessToJobObject failed", _last_error(ctypes)) kernel32.CloseHandle(job) return _win_job_handle = job # hold the handle so the job is not closed early - except Exception: - pass + _record_job_status(True, "kill-on-close job installed") + except Exception as error: + _record_job_status(False, f"{type(error).__name__}: {error}") # ── Child binding ── @@ -223,10 +310,20 @@ def _composed() -> None: def _reset_after_fork() -> None: - """A fork child inherits _spawner_lock in whatever state it was in and a + """A fork child inherits both locks in whatever state they were in and a _spawner whose thread does not exist here. Start clean instead of deadlocking.""" - global _spawner, _spawner_lock + global _spawner, _spawner_lock, _record_lock, _owner_identity _spawner_lock = threading.Lock() + # A different pid here, so the parent's identity is not this process's, and + # the parent's children are not this process's either: adopting anything + # would otherwise write a record claiming them, and a later startup would + # reap them out from under the parent that is still running. + _owner_identity = None + _tracked_pids.clear() + _tracked_pgids.clear() + # A fork while another thread was inside adopt_pid / forget_pid leaves this + # held here with nobody to release it, and the first adoption blocks forever. + _record_lock = threading.Lock() _spawner = None @@ -325,24 +422,496 @@ def child_popen_kwargs(preexec_fn: Optional[Callable[[], None]] = None) -> dict: return {} +def _recorded_identity(value: object) -> "Optional[str]": + """An identity read back from a record, or None when it is not one.""" + return value if isinstance(value, str) and value else None + + +def _same_identity(recorded: str, current: str) -> bool: + """Whether two identities describe the same process. + + Records written before this carried ``starttime:comm`` on Linux, so compare + the start time alone there. Elsewhere the whole string is generated by the + same code and colons are part of the value (Windows FILETIME, macOS lstart). + """ + # A record on disk can hold anything that parses as JSON. Reaching split() + # with a number raises, and run_server catches that around the whole sweep, + # so one bad file would leave every other orphan running. + if not isinstance(recorded, str) or not isinstance(current, str): + return False + if _is_linux(): + return recorded.split(":", 1)[0] == current.split(":", 1)[0] + return recorded == current + + 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 + # Start time pins identity so a reused pid is never signalled later. None + # (unreadable, or a platform with no cheap source) disables the check. + if _is_linux(): + try: + with open(f"/proc/{pid}/stat", encoding = "utf-8") as fh: + stat = fh.read() + # Start time only. comm is mutable (prctl PR_SET_NAME, setproctitle), + # so a child that renames itself would read as a recycled pid and be + # dropped unsignalled -- the orphan this module exists to prevent. + return stat[stat.rfind(")") + 2 :].split()[19] + except Exception: + return None + if _is_windows(): + # Creation time, so a recycled pid is never mistaken for the child that + # was recorded. Same purpose as starttime on Linux. + try: + import ctypes + from ctypes import wintypes + + PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + kernel32.GetProcessTimes.argtypes = [wintypes.HANDLE] + [ + ctypes.POINTER(wintypes.FILETIME) + ] * 4 + kernel32.GetProcessTimes.restype = wintypes.BOOL + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + kernel32.CloseHandle.restype = wintypes.BOOL + handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid) + if not handle: + return None + try: + created = wintypes.FILETIME() + other = [wintypes.FILETIME() for _ in range(3)] + if not kernel32.GetProcessTimes( + handle, ctypes.byref(created), *[ctypes.byref(x) for x in other] + ): + return None + return f"{created.dwHighDateTime}:{created.dwLowDateTime}" + finally: + kernel32.CloseHandle(handle) + except Exception: + return None + if sys.platform == "darwin": + # No /proc; `ps -o lstart` is enough to spot a recycled pid. + try: + import subprocess + + # TZ pinned: lstart is formatted in local time, so a machine that + # changes timezone between adopt and check would otherwise read as a + # different process. + out = subprocess.run( + ["ps", "-o", "lstart=,comm=", "-p", str(pid)], + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 5, + env = {**os.environ, "TZ": "UTC"}, + ) + line = (out.stdout or "").strip() + return line or None + except Exception: + return None + 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: + signals a recycled pid. + + Kept when its process group still has members: the shim can exit before the + visual server it started, and this record is the only handle on that group. + """ + if not pid: + return + with _record_lock: + if pid not in _tracked_pids and pid not in _tracked_pgids: + return # never recorded, so there is nothing to rewrite or check + if _group_has_members(_tracked_pgids.get(pid)): + return _tracked_pids.pop(pid, None) + _tracked_pgids.pop(pid, None) + _write_breadcrumb() + + +def _group_has_members(pgid: object) -> bool: + """Whether the group still holds a process that is actually running. + + ``killpg(pgid, 0)`` alone is not enough: a leader that has exited but has + not been waited on is still a member, so a group whose every member is a + zombie would read as alive and keep its record forever. + """ + if not isinstance(pgid, int) or _is_windows() or not hasattr(os, "killpg"): + return False + try: + os.killpg(pgid, 0) + except Exception: + return False + # The leader is the usual answer, and enumerating a group means reading the + # state of every process on the machine: 62ms on a box with 6000 of them, + # paid on every stop. + if _pid_alive(pgid) and not _pid_is_zombie(pgid): + return True + members = _group_member_pids(pgid) + if members is None: + return True # cannot enumerate, so keep the record rather than lose it + return any(not _pid_is_zombie(pid) for pid in members) + + +def _child_pid_map() -> "Optional[dict[int, list[int]]]": + """Parent pid -> its children, or None when the table cannot be read.""" + if _is_linux(): + try: + table: "dict[int, list[int]]" = {} + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + try: + with open(f"/proc/{entry}/stat", encoding = "utf-8") as fh: + stat = fh.read() + except OSError: + continue + # After the comm field: state, ppid, pgrp, ... + tail = stat[stat.rfind(")") + 2 :].split() + if len(tail) > 1: + table.setdefault(int(tail[1]), []).append(int(entry)) + return table + except Exception: + return None + if sys.platform == "darwin": + try: + import subprocess + + out = subprocess.run( + ["ps", "-A", "-o", "pid=,ppid="], + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 5, + ) + if out.returncode != 0: + return None # a failed query is not an answer, as above + table = {} + for line in (out.stdout or "").splitlines(): + parts = line.split() + if len(parts) != 2: + continue + try: + child, parent = int(parts[0]), int(parts[1]) + except ValueError: + continue + table.setdefault(parent, []).append(child) + return table + except Exception: + return None + return None + + +def collect_descendants(pid: "Optional[int]") -> "list[tuple[int, Optional[str]]]": + """A pid's descendants and their start-time identities. + + Read this BEFORE signalling the parent: its children are reparented the + moment it exits, and nothing then ties them back to it. The identities let + the kill below skip a number that has since moved on to something else. + """ + if not pid or _is_windows(): + return [] + table = _child_pid_map() + if not table: + return [] + found: "list[tuple[int, Optional[str]]]" = [] + seen = {pid} + queue = list(table.get(pid, ())) + while queue: + child = queue.pop(0) + if child in seen: + continue + seen.add(child) + found.append((child, _pid_identity(child))) + queue.extend(table.get(child, ())) + return found + + +def terminate_descendants( + collected: "list[tuple[int, Optional[str]]]", timeout: float = 5.0 +) -> None: + """SIGTERM then SIGKILL what `collect_descendants` found, still alive. + + The POSIX counterpart of the Windows ``taskkill /T``: a child that shares + this process's group cannot be reached with killpg, so its own children are + signalled by pid instead. + """ + if not collected or _is_windows(): + return + live: "list[tuple[int, Optional[str]]]" = [] + for pid, identity in collected: + if not _still_the_same(pid, identity): + continue + try: + os.kill(pid, signal.SIGTERM) + except OSError: + continue + live.append((pid, identity)) + deadline = time.monotonic() + max(0.0, timeout) + while live: + live = [item for item in live if _pid_alive(item[0]) and not _pid_is_zombie(item[0])] + if not live or time.monotonic() >= deadline: + break + time.sleep(0.05) + for pid, identity in live: + if not _still_the_same(pid, identity): + continue + try: + os.kill(pid, signal.SIGKILL) + except OSError: + pass + + +def _still_the_same(pid: int, identity: "Optional[str]") -> bool: + """False only when the pid is provably a different process now.""" + current = _pid_identity(pid) + if identity is None or current is None: + return True + return _same_identity(identity, current) + + +def _group_member_pids(pgid: int) -> "Optional[list[int]]": + """Pids in a process group, or None when they cannot be enumerated.""" + if _is_linux(): + try: + found = [] + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + try: + with open(f"/proc/{entry}/stat", encoding = "utf-8") as fh: + stat = fh.read() + except OSError: + continue + # After the comm field: state, ppid, pgrp, ... + tail = stat[stat.rfind(")") + 2 :].split() + if len(tail) > 2 and tail[2] == str(pgid): + found.append(int(entry)) + return found + except Exception: + return None + if sys.platform == "darwin": + try: + import subprocess + + out = subprocess.run( + ["ps", "-o", "pid=", "-g", str(pgid)], + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 5, + ) + # ps exits nonzero when the group is gone AND when the call itself + # failed, and the two are not the same answer: reporting "empty" + # for a failure lets forget_pid drop the only record of a live + # descendant. Only a clean exit is an answer. + if out.returncode != 0: + return None + return [int(x) for x in (out.stdout or "").split()] + except Exception: + return None + return None + + +# ── Crash-survivable child record ── +# +# macOS has neither PR_SET_PDEATHSIG nor job objects, so a crash leaves every +# sidecar running. Record children as they are adopted; sweep the previous +# run's leftovers at startup. + + +def _breadcrumb_dir(): + from pathlib import Path + + override = os.environ.get("UNSLOTH_STUDIO_CHILD_RECORD") + if override: + return Path(override) + try: + from utils.paths.storage_roots import studio_root + return Path(studio_root()) / "run" / "children" + except Exception: + return None + + +def _breadcrumb_file(): + # One file per owner: two Studios can share a home (different ports), and a + # single shared file would let the second erase the first's children. + directory = _breadcrumb_dir() + return None if directory is None else directory / f"{os.getpid()}.json" + + +_owner_identity: Optional[str] = None + + +def _own_identity() -> "Optional[str]": + """This process's identity, retried like a child's and then kept. + + Recorded as None, any process that later reuses this pid reads as the owner + still running and the children this record names are never reaped. It + cannot change, so it is captured once rather than under the record lock. + """ + global _owner_identity + if _owner_identity is None: + _owner_identity = _identity_for_record(os.getpid()) + return _owner_identity + + +def _refreshed_identity(pid: int, identity: "Optional[str]") -> "Optional[str]": + """Fill in an identity the adoption could not read, while the child lives. + + Recorded as None it is permanent, and neither terminate_all nor the startup + sweep will signal an entry it cannot verify, so that child would outlive + every shutdown. Written back, so this costs one probe per gap. + """ + if identity is not None or not _pid_alive(pid): + return identity + identity = _pid_identity(pid) + if identity is not None: + _tracked_pids[pid] = identity + return identity + + +def _write_breadcrumb() -> None: + path = _breadcrumb_file() + if path is None: + return + try: + import json + + path.parent.mkdir(parents = True, exist_ok = True) + payload = { + "owner_pid": os.getpid(), + "owner_identity": _own_identity(), + "children": [ + { + "pid": pid, + "identity": _refreshed_identity(pid, identity), + "pgid": _tracked_pgids.get(pid), + } + for pid, identity in _tracked_pids.items() + ], + } + tmp = path.with_suffix(".tmp") + tmp.write_text(json.dumps(payload), encoding = "utf-8") + tmp.replace(path) + except Exception: + pass + + +def clear_breadcrumb() -> None: + """Drop our record after a clean shutdown. + + Only when nothing is left: a child that outlived `terminate_all` still needs + the record, or the next startup has no way to find it. + """ + with _record_lock: + if _tracked_pids: + _write_breadcrumb() + return + # Inside the lock: a spawn landing between the check and the unlink + # would otherwise write a record we then delete. + path = _breadcrumb_file() + if path is not None: + _unlink(path) + + +def _identity_or_none(pid) -> "Optional[str]": + return _pid_identity(pid) if isinstance(pid, int) and pid > 0 else None + + +def _pid_alive(pid: int) -> bool: + if _is_windows(): + # NOT os.kill(pid, 0): that is TerminateProcess here, so the probe + # would kill the process it is asking about. + try: + import ctypes + from ctypes import wintypes + + SYNCHRONIZE = 0x0010_0000 + PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 + WAIT_TIMEOUT = 0x102 + ERROR_ACCESS_DENIED = 5 + + kernel32 = ctypes.WinDLL("kernel32", use_last_error = True) + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD] + kernel32.WaitForSingleObject.restype = wintypes.DWORD + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + + handle = kernel32.OpenProcess( + SYNCHRONIZE | PROCESS_QUERY_LIMITED_INFORMATION, False, pid + ) + if not handle: + # ACCESS_DENIED means alive but another user's; else it is gone. + return _last_error(ctypes) == ERROR_ACCESS_DENIED + try: + # Signalled means exited; still waiting means running. + return kernel32.WaitForSingleObject(handle, 0) == WAIT_TIMEOUT + finally: + kernel32.CloseHandle(handle) + except Exception: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + except OSError: + return False + return True + + +def _pid_is_zombie(pid: int) -> bool: + """An exited child nobody waited on. It answers signals like a live process, + so a survivor check has to tell the two apart.""" + if _is_windows(): + return False + if _is_linux(): + try: + with open(f"/proc/{pid}/stat", encoding = "utf-8") as fh: + return fh.read().rsplit(")", 1)[1].split()[0] == "Z" + except Exception: + return False + try: + import subprocess + out = subprocess.run( + ["ps", "-o", "state=", "-p", str(pid)], + capture_output = True, + text = True, + encoding = "utf-8", + errors = "replace", + timeout = 5, + ) + return (out.stdout or "").strip().startswith("Z") + except Exception: + return False + + +def _identity_for_record(pid: int, attempts: int = 3) -> Optional[str]: + """Identity to persist for *pid*, retried while it is still running. + + A `ps` that timed out once would otherwise be recorded as "no identity", + and an entry with no identity is never signalled, so that child survives + every later launch. + """ + import time + + for attempt in range(attempts): + identity = _pid_identity(pid) + if identity is not None: + return identity + if not _pid_alive(pid): + break # already gone: there is nothing left to identify + if attempt + 1 < attempts: + time.sleep(0.05 * (attempt + 1)) + return None def adopt_pid(pid: Optional[int]) -> None: @@ -351,7 +920,17 @@ def adopt_pid(pid: Optional[int]) -> None: Tolerates a None or already-exited pid.""" if not pid: return - _tracked_pids[pid] = _pid_identity(pid) + # Here as well as in the Linux spawn path: this is the first thing that + # writes a record, and without the handler a fork child keeps this + # process's children and later claims them as its own. + _adopt_fork_reset() + identity = _identity_for_record(pid) + pgid = _own_process_group(pid) + with _record_lock: + _tracked_pids[pid] = identity + if pgid is not None: + _tracked_pgids[pid] = pgid + _write_breadcrumb() if _is_windows() and _win_job_handle: try: import ctypes @@ -370,35 +949,387 @@ def adopt_pid(pid: Optional[int]) -> None: pass -def terminate_all(timeout: float = 5.0) -> None: +def terminate_all(timeout: float = 5.0) -> "list[int]": """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 + then SIGKILL the survivors after `timeout`. Idempotent and teardown-safe. + + Signals only a pid whose recorded start-time identity still matches. A pid + that cannot be verified is left alone and kept in the record, so the startup + sweep can retry it rather than this process signalling a stranger. + + Returns the pids still alive afterwards, so the caller can keep them in the + crash record rather than dropping the only handle on them.""" + survivors: "list[int]" = [] + # Snapshot under the same lock the writes take: a request thread can still + # reach adopt_pid while this runs. + with _record_lock: + tracked = list(_tracked_pids.items()) + for pid, identity in tracked: + with _record_lock: + _tracked_pids.pop(pid, None) + pgid = _tracked_pgids.pop(pid, None) + current = _pid_identity(pid) + if not _pid_alive(pid): + # Same as the startup sweep: a leader that exited first can still + # have a group holding the GPU behind it. + if not _reap_orphaned_group(pgid, pid, timeout) and _group_has_members(pgid): + survivors.append(pid) + with _record_lock: + _tracked_pids[pid] = identity + if pgid is not None: + _tracked_pgids[pid] = pgid + continue + if current is not None and identity is not None and not _same_identity(identity, current): + continue # definitely recycled by something else; drop it + if identity is None or current is None: + # Cannot prove this is still our child, so do not signal it. Keep it + # recorded while it is alive: the startup sweep runs the same test. + if _pid_alive(pid) and not _pid_is_zombie(pid): + with _record_lock: + _tracked_pids[pid] = identity + # The group goes back with it: keeping the pid but dropping + # its group leaves nothing able to reach a descendant once + # the leader exits. + if pgid is not None: + _tracked_pgids[pid] = pgid + continue + tree_stands = False try: if _is_windows(): - os.kill(pid, signal.SIGTERM) - continue + # The tree: a leader killed alone strands its workers, and the + # record naming them is cleared right after. + tree_stands = not _windows_terminate_tree(pid) + else: + _posix_terminate(pid, timeout) + except Exception: + pass + current_now = _identity_or_none(pid) + still_ours = ( + _pid_alive(pid) + and not _pid_is_zombie(pid) + and current_now is not None + and _same_identity(identity, current_now) + ) + # Or the leader went and its group did not: same loss of the only handle, + # which on Windows is a tree taskkill could not take. + if still_ours or tree_stands or _group_has_members(pgid): + survivors.append(pid) + with _record_lock: + _tracked_pids[pid] = identity + if pgid is not None: + _tracked_pgids[pid] = pgid + return survivors + + +def terminate_pid(pid: "Optional[int]", timeout: float = 5.0) -> None: + """Stop one tracked child now, tree and all, and drop its record. + + For an owner that has to give up on a child before its own shutdown, and + cannot leave it for a sweep that will not run while this process lives. + """ + if not pid: + return + with _record_lock: + identity = _tracked_pids.get(pid) + pgid = _tracked_pgids.get(pid) + # Same test terminate_all runs. An announced child can exit without the line + # that clears it, and its pid is free the moment the group behind it empties, + # so signalling on the number alone can take a stranger's tree down. + current = _pid_identity(pid) + if identity is not None and current is not None and not _same_identity(identity, current): + # A pid is only reusable once nothing holds the number as a process + # group either, so there is no group of ours left to reap here. + forget_pid(pid) + return + if _pid_alive(pid) and (identity is None or current is None): + # Cannot prove this is still our child. Leave it alone and keep the + # record: the startup sweep repeats the test with a fresh reading. + return + tree_stands = False + try: + if _is_windows(): + # False is "only the leader was signalled": nothing else names those + # workers, so the record has to outlive this call. + tree_stands = not _windows_terminate_tree(pid) + else: _posix_terminate(pid, timeout) + # A leader that exited first takes getpgid with it, so _posix_terminate + # signals the dead pid alone; the recorded group is the only handle + # left on the session it started. + if _group_has_members(pgid): + _reap_orphaned_group(pgid, pid, timeout) + except Exception: # noqa: BLE001 - best effort, like the rest of this + pass + if tree_stands: + return + forget_pid(pid) # keeps the record if its group is still up + + +def reap_recorded_children(timeout: float = 5.0) -> "list[int]": + """Kill children recorded by a previous Studio that is no longer running. + + Runs once at startup, before anything new spawns. Every record in the + directory is considered, so a Studio that crashed while a sibling was + running is still cleaned up. A child is only signalled when its recorded + start-time identity still matches, so a recycled pid is never touched. + """ + directory = _breadcrumb_dir() + if directory is None or not directory.is_dir(): + return [] + killed: "list[int]" = [] + # A record is skipped while its owner is alive, and that owner may be + # terminated later in the same sweep by the record of the backend that + # spawned it. Only those deferred records are revisited, so nothing is + # signalled twice. + pending = sorted(directory.glob("*.json")) + for _pass in range(4): + deferred: "list" = [] + found: "list[int]" = [] + for path in pending: + reaped, owner_alive = _reap_one_record(path, timeout) + found.extend(reaped) + if owner_alive: + deferred.append(path) + killed.extend(found) + if not deferred or not found: + break + pending = deferred + return killed + + +def _reap_one_record(path, timeout: float) -> "tuple[list[int], bool]": + """``(pids signalled, whether it was left alone because its owner is alive)``.""" + import json + + killed: "list[int]" = [] + try: + record = json.loads(path.read_text(encoding = "utf-8")) + except Exception: + _unlink(path) + return killed, False + if not isinstance(record, dict): + _unlink(path) + return killed, False + + owner_pid = record.get("owner_pid") + # Anything that is not the string this wrote is no identity at all, and is + # treated like a missing one rather than trusted or crashed on. + owner_identity = _recorded_identity(record.get("owner_identity")) + # Identity decides, not the pid: pids recycle, and this process may have + # inherited the pid of the Studio that wrote the record. + # Inconclusive counts as a match: a `ps` that failed for a moment must not + # make a live Studio look gone and cost it its running sidecars. + current_owner = _identity_or_none(owner_pid) + owner_matches = ( + owner_identity is None + or current_owner is None + or _same_identity(owner_identity, current_owner) + ) + if owner_pid == os.getpid() and owner_matches: + return killed, False # our own record + if ( + isinstance(owner_pid, int) + and owner_matches + and _pid_alive(owner_pid) + # os.kill(pid, 0) succeeds for a zombie, and a Studio nobody has waited + # on yet is still a dead one whose sidecars are orphans. + and not _pid_is_zombie(owner_pid) + ): + return killed, True # that Studio is still running; its children are its own + + unresolved = False + children = record.get("children") + for entry in children if isinstance(children, list) else []: + pid = entry.get("pid") if isinstance(entry, dict) else None + if not isinstance(pid, int): + continue + # A zombie is a dead leader nobody has waited on: it holds nothing, and + # signalling it would burn the whole grace period answering probes. + if not _pid_alive(pid) or _pid_is_zombie(pid): + # The leader can exit first and leave the group running (the shim + # crashing while its visual server holds the GPU). The group is then + # the only handle left on those children. + pgid = entry.get("pgid") + if _reap_orphaned_group(pgid, pid, timeout): + killed.append(pid) + elif _group_has_members(pgid): + # Still there, so this record is the only handle on it. + unresolved = True + continue + identity = _recorded_identity(entry.get("identity")) + current = _identity_or_none(pid) + if identity is None or current is None: + # Unverifiable: never signal a pid that might now be something else. + unresolved = True + continue + if not _same_identity(identity, current): + continue # recycled by an unrelated process + tree_stands = False + if _is_windows(): + # The tree, not the leader: this fallback runs when the Job Object + # is unavailable, and killing a leader alone strands its workers + # while the record that named them is deleted. + tree_stands = not _windows_terminate_tree(pid) + else: + _posix_terminate(pid, timeout = timeout) + killed.append(pid) + if tree_stands: + # Workers may still be running with nothing else naming them, so + # this record has to survive for the next launch to retry. + unresolved = True + # These belong to a process that is gone, so they cannot be our zombies; + # one still running means the kill did not take, and the record is the + # only handle the next launch would have on it. + if _pid_alive(pid) and not _pid_is_zombie(pid): + unresolved = True + elif _group_has_members(entry.get("pgid")): + # The leader died but its group did not, which is the same loss of + # the only handle, just one step later. + unresolved = True + + if not unresolved: + _unlink(path) + if killed: + try: + import logging + logging.getLogger(__name__).warning( + "Reaped %d orphaned child process(es) left by a previous Studio: %s", + len(killed), + killed, + ) except Exception: pass + return killed, False + + +def _own_process_group(pid: int) -> Optional[int]: + """The pid's process group, but only when it leads one (start_new_session). + + A child sharing Studio's group must never be recorded: killing that group + would take Studio and every sibling with it. + """ + if _is_windows() or not hasattr(os, "getpgid"): + return None + try: + pgid = os.getpgid(pid) + except Exception: + return None + return pgid if pgid == pid else None + + +def _reap_orphaned_group(pgid: object, pid: int, timeout: float) -> bool: + """Signal a recorded group whose leader is already gone. True when it was. + + Safe without an identity check: the group id is the dead leader's pid, and + the kernel holds that number for as long as any task still references it as + a process group, so it cannot have been handed to an unrelated group while + members remain. + """ + if not isinstance(pgid, int) or pgid != pid or _is_windows() or not hasattr(os, "killpg"): + return False + try: + os.killpg(pgid, 0) # the group is still there + except Exception: + return False + # A zombie answers that probe, and where pid 1 does not reap it stays that + # way, so without this every stale record costs the whole grace period. + if not _group_has_members(pgid): + return False + import time + + try: + os.killpg(pgid, signal.SIGTERM) + except Exception: + return False + deadline = time.monotonic() + max(0.0, timeout) + next_state_check = 0.0 + while time.monotonic() < deadline: + try: + os.killpg(pgid, 0) + except Exception: + return True + now = time.monotonic() + if now >= next_state_check: + # Where pid 1 does not reap (a container), a member that exits on + # the SIGTERM stays a zombie and keeps answering that probe, so the + # wait would run its full length for a group that is already gone. + next_state_check = now + 0.5 + if not _group_has_members(pgid): + return True + time.sleep(0.1) + try: + os.killpg(pgid, signal.SIGKILL) + except Exception: + pass + # Only gone counts: a SIGKILL that could not be delivered leaves this record + # as the last handle on whatever is still holding the GPU. + try: + os.killpg(pgid, 0) + return False + except Exception: + return True + + +def _windows_terminate_tree(pid: int) -> bool: + """``taskkill /T /F``. True when the whole tree is gone. + + False means only the leader was signalled and its workers may still be + running. This is the fallback for a root with no job object, so the record + naming it is the only handle left on them: the caller keeps it rather than + reading the dead leader as the tree being gone. Caller has already verified + the pid's identity. + """ + import subprocess + + try: + completed = subprocess.run( + ["taskkill", "/PID", str(pid), "/T", "/F"], + capture_output = True, + timeout = 15, + creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0), + ) + # check = False does not raise, so the status is the only signal that + # the tree is still standing. 128 is "already gone". + if completed.returncode in (0, 128): + return True + except (OSError, subprocess.SubprocessError): + pass + try: + os.kill(pid, signal.SIGTERM) + except OSError: + pass + return False + + +def _unlink(path) -> None: + try: + path.unlink(missing_ok = True) + 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 + group_leader = False try: - if os.getpgid(pid) == pid: - killer = os.killpg + group_leader = os.getpgid(pid) == pid except Exception: pass + # A child sharing this process's group cannot be reached with killpg, so its + # own children have to be named individually -- and named now, while the + # parent that links them is still alive. Always run, so a return out of the + # signalling below still takes the tree. + descendants = [] if group_leader else collect_descendants(pid) + try: + _posix_terminate_one(pid, group_leader, timeout) + finally: + terminate_descendants(descendants, timeout) + + +def _posix_terminate_one(pid: int, group_leader: bool, timeout: float) -> None: + killer = os.killpg if group_leader else os.kill try: killer(pid, signal.SIGTERM) except ProcessLookupError: @@ -406,6 +1337,7 @@ def _posix_terminate(pid: int, timeout: float = 5.0) -> None: except Exception: return deadline = time.monotonic() + max(0.0, timeout) + next_state_check = 0.0 while time.monotonic() < deadline: try: killer(pid, 0) # still alive? @@ -413,6 +1345,16 @@ def _posix_terminate(pid: int, timeout: float = 5.0) -> None: return except Exception: break + now = time.monotonic() + if now >= next_state_check: + # An exited child nobody has waited on answers signal 0 exactly like + # a live one, so without this the whole timeout is spent on a process + # that is already gone. Reading the state costs a fork off Linux, + # hence twice a second rather than at the poll rate. + next_state_check = now + 0.5 + gone = (not _group_has_members(pid)) if group_leader else _pid_is_zombie(pid) + if gone: + return time.sleep(0.05) try: killer(pid, signal.SIGKILL) diff --git a/studio/frontend/src/hooks/use-tauri-update.ts b/studio/frontend/src/hooks/use-tauri-update.ts index ec00dbc36d..0c95b503e5 100644 --- a/studio/frontend/src/hooks/use-tauri-update.ts +++ b/studio/frontend/src/hooks/use-tauri-update.ts @@ -117,6 +117,25 @@ export function useTauriUpdate(isExternalServer = false) { const startupScheduledRef = useRef(false); const checkingRef = useRef(false); const updatingRef = useRef(false); + // Windows kill-on-close: false once a re-arm has failed, and every path that + // starts a backend has to check it or the orphan risk comes straight back. + const cleanupRearmedRef = useRef(true); + // A webview reload rebuilds this hook with the ref back at its initial value + // while the native job may still be disarmed, so the first gate of a mount + // asks the native side instead of trusting it. + const cleanupCheckedRef = useRef(false); + + async function resumeCleanup(): Promise { + try { + const { invoke } = await import("@tauri-apps/api/core"); + await invoke("resume_desktop_update_cleanup"); + cleanupRearmedRef.current = true; + } catch (e) { + console.error("Could not re-arm crash cleanup after a failed update:", e); + cleanupRearmedRef.current = false; + } + return cleanupRearmedRef.current; + } function replaceInfo(nextInfo: UpdateInfo | null) { infoRef.current = nextInfo; @@ -291,8 +310,10 @@ export function useTauriUpdate(isExternalServer = false) { updatingRef.current = true; const cleanups: (() => void)[] = []; - try { + // A retry re-enters here, and start_backend_update spawns an + // environment-mutating child of its own. + if (!(await crashCleanupReady())) return; const { policy } = await resolveUpdatePolicy(); if (policy.mode === "manual_linux_package") { const version = info?.version ?? updateRef.current?.version; @@ -381,20 +402,32 @@ export function useTauriUpdate(isExternalServer = false) { break; } }); + } catch (installError) { + // Failed or cancelled: we keep running, so the cleanup the pre-exit hook + // stood down has to come back. + await resumeCleanup(); + throw installError; } finally { - // Success, failure and cancel alike: the installer is no longer running. publishShellUpdateActive(false); } - // relaunch() re-execs with the original argv, so flag the inherited --hidden as not a login - // start. It only fails when there is a --hidden to suppress, so let it stop the restart. - await invoke("mark_in_app_relaunch"); - const { relaunch } = await import("@tauri-apps/plugin-process"); + // Deliberately NOT re-arming kill-on-close before the restart: relaunch() + // starts the replacement as a child, so it inherits this job, and re-arming + // would make this process kill it on the way out. + // The whole handoff is inside the recovery scope: anything that throws here + // leaves this process running with cleanup still stood down. try { + // relaunch() re-execs with the original argv, so flag the inherited --hidden as not a + // login start. It only fails when there is a --hidden to suppress, so let it stop the + // restart. + await invoke("mark_in_app_relaunch"); + const { relaunch } = await import("@tauri-apps/plugin-process"); await relaunch(); } catch (relaunchError) { // No replacement process, so the marker would outlive it and unhide a later login start. await invoke("clear_in_app_relaunch").catch(() => {}); + // Still this process, so the cleanup has to come back after all. + await resumeCleanup(); throw relaunchError; } } catch (e) { @@ -403,6 +436,13 @@ export function useTauriUpdate(isExternalServer = false) { // Shell update failed, so restart the backend on the updated code. if (phaseRef.current === "shell_download" || phaseRef.current === "shell_install") { + // A backend started under a job that still has kill-on-close disabled is + // the orphan this PR exists to prevent, so retry the re-arm and stop here + // if it will not take. + if (!(await crashCleanupReady())) { + retainFailure(msg, phaseRef.current ?? "shell_install"); + return; + } try { const { invoke } = await import("@tauri-apps/api/core"); await invoke("start_server", { port: 8888 }); @@ -436,8 +476,33 @@ export function useTauriUpdate(isExternalServer = false) { await installUpdate(); } + /** Every path that starts a child has to clear this first. */ + async function crashCleanupReady(): Promise { + if (!cleanupCheckedRef.current) { + cleanupCheckedRef.current = true; + try { + const { invoke } = await import("@tauri-apps/api/core"); + cleanupRearmedRef.current = await invoke("desktop_update_cleanup_armed"); + } catch { + // On the desktop this is the one answer we cannot assume: fail closed + // and let the gate below re-arm. In the browser there is no job at all. + cleanupRearmedRef.current = !isTauri; + } + } + if (cleanupRearmedRef.current) return true; + if (await resumeCleanup()) return true; + setError( + "Crash cleanup could not be re-armed. Restart Unsloth before continuing.", + ); + setStatus("error"); + return false; + } + async function skipAndRestart() { const skippedError = error; + // Same gate as the recovery path: this is offered on every error, so it is + // the other way a user could start a backend under a disarmed job. + if (!(await crashCleanupReady())) return; try { const { invoke } = await import("@tauri-apps/api/core"); await invoke("start_server", { port: 8888 }); diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index a634a5828d..b1bbcfae70 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -18,6 +18,7 @@ import random import re import shutil +import signal import site import socket import ssl @@ -5299,8 +5300,13 @@ def validate_server( stderr = subprocess.STDOUT, text = True, env = binary_env(server_path, install_dir, host, runtime_line = runtime_line), + **_validation_server_kwargs(), **windows_hidden_subprocess_kwargs(), ) + # For the caller that spawned this script: a validation server is + # its grandchild, so a crash here leaves it holding the GPU and + # the staged files with nothing recording where it is. + _announce_child("started", process.pid) deadline = time.time() + 60 startup_started = time.time() response_body = "" @@ -5360,13 +5366,11 @@ def validate_server( + ("\n" + response_body if response_body else "") ) finally: - if process is not None and process.poll() is None: - process.terminate() - try: - process.wait(timeout = 5) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout = 5) + if process is not None: + # Only once nothing is left in its group: announcing the stop + # drops the record, which is the only handle on a survivor. + if _terminate_validation_server(process): + _announce_child("stopped", process.pid) try: log_path.unlink(missing_ok = True) except Exception: @@ -5376,6 +5380,134 @@ def validate_server( raise PrebuiltFallback("llama-server validation failed unexpectedly") +_PR_SET_PDEATHSIG = 1 + + +def _validation_server_leads_group() -> bool: + """Whether the validation server is started as its own group leader. + + Only where the parent-death signal can be armed with it, so the two always + agree: the kill path may only killpg a group the server actually leads. + """ + return os.name == "posix" and sys.platform.startswith("linux") + + +def _validation_server_kwargs() -> "dict[str, Any]": + """Popen kwargs tying a validation server to this installer's lifetime. + + Its own group keeps whatever the server starts reachable through the leader + alone, but it also takes the server out of the group Studio force-kills, so + the parent-death signal is armed alongside it: an installer that is killed + mid-validation must not leave a server holding the GPU and the staged files + until some later startup sweeps the breadcrumb. + + macOS has no such signal, and the record only exists once the backend has + read the announcement that follows the spawn, so a session of its own there + would leave a window in which nothing can reach the server at all. It stays + in the inherited group instead, and the termination walk covers whatever it + starts. + """ + if not _validation_server_leads_group(): + return {} + kwargs: "dict[str, Any]" = {"start_new_session": True} + owner_pid = os.getpid() # read pre-fork, so the child can tell reparenting apart + + def _arm_parent_death() -> None: + # Post-fork, pre-exec: fork clears the setting, and subprocess runs + # preexec_fn after setsid, neither of which undoes it (a plain execve + # preserves it too). getppid closes the race where the installer was + # already gone by the time this ran. + try: + import ctypes + ctypes.CDLL("libc.so.6", use_errno = True).prctl(_PR_SET_PDEATHSIG, signal.SIGTERM) + if os.getppid() != owner_pid: + os._exit(1) + except Exception: + pass + + kwargs["preexec_fn"] = _arm_parent_death + return kwargs + + +def _announce_child(state: str, pid: int) -> None: + """Tell whoever runs this script about a server it started. + + Studio adopts the pid so its own sweep can reach it; run by hand the line is + just noise on stdout. + """ + print(f"UNSLOTH_INSTALLER_CHILD {state} {pid}", flush = True) + + +def _group_is_running(pgid: "int | None") -> bool: + """Whether anything is left in that process group.""" + if pgid is None or os.name != "posix" or not hasattr(os, "killpg"): + return False + try: + os.killpg(pgid, 0) + except OSError: + return False + return True + + +def _wait_for_group(pgid: "int | None", grace: float) -> None: + """Give a signalled group its moment to go.""" + deadline = time.time() + grace + while _group_is_running(pgid) and time.time() < deadline: + time.sleep(0.05) + + +def _terminate_validation_server(process: "subprocess.Popen", grace: float = 5.0) -> bool: + """Stop the validation server and anything it started. True when it is gone. + + Where it leads its own group, signalling the leader alone would leave a + child of its own behind, and a child that ignores SIGTERM outlives the + leader's exit. False means something is still in that group, and the caller + keeps the pid announced so a later sweep can still reach it. + + Where it does not lead one it shares this installer's group, and killpg + would take the installer and Studio with it, so only the server itself is + signalled. + """ + pgid = None + if _validation_server_leads_group() and hasattr(os, "getpgid"): + try: + pgid = os.getpgid(process.pid) + except OSError: + # Already reaped, so its own pid is the group id: it was started in + # a session of its own, and the kernel holds that number for as long + # as any member of the group is still there. + pgid = process.pid + if pgid != process.pid: + pgid = None # not the leader after all; never signal a shared group + if pgid is not None: + try: + os.killpg(pgid, signal.SIGTERM) + except OSError: + process.terminate() + else: + process.terminate() + try: + process.wait(timeout = grace) + except subprocess.TimeoutExpired: + pass + # The leader exiting is not the answer: a member that ignored the SIGTERM + # is still holding the GPU and the staged files. + _wait_for_group(pgid, grace) + if pgid is not None and _group_is_running(pgid): + try: + os.killpg(pgid, signal.SIGKILL) + except OSError: + pass + _wait_for_group(pgid, grace) + elif pgid is None and process.poll() is None: + process.kill() + try: + process.wait(timeout = grace) + except subprocess.TimeoutExpired: + pass + return not _group_is_running(pgid) and process.poll() is not None + + def collect_system_report(host: HostInfo, choice: AssetChoice | None, install_dir: Path) -> str: lines = [ f"platform={host.system} machine={host.machine}", diff --git a/studio/src-tauri/src/desktop_updater.rs b/studio/src-tauri/src/desktop_updater.rs index 344dddc4c2..16d542a16d 100644 --- a/studio/src-tauri/src/desktop_updater.rs +++ b/studio/src-tauri/src/desktop_updater.rs @@ -13,6 +13,36 @@ pub(crate) struct DesktopUpdateMetadata { raw_json: serde_json::Value, } +/// Re-arm crash cleanup when the installer never took over. +/// +/// `on_before_exit` clears kill-on-close assuming the installer is about to +/// replace this process; a failed or cancelled install leaves the app running +/// with no reaper for its children. +#[tauri::command] +pub(crate) async fn resume_desktop_update_cleanup() -> Result<(), String> { + // Both halves of what the pre-exit hook consumed: the exit guard it spent, so a + // retry actually reaps the backend again, and kill-on-close itself. + crate::reset_termination_cleanup(); + #[cfg(windows)] + { + crate::windows_job::resume_after_update_installer().map_err(|error| error.to_string())?; + } + Ok(()) +} + +/// Whether crash cleanup is armed right now, for a UI that has just remounted. +#[tauri::command] +pub(crate) async fn desktop_update_cleanup_armed() -> Result { + #[cfg(windows)] + { + crate::windows_job::kill_on_close_armed().map_err(|error| error.to_string()) + } + #[cfg(not(windows))] + { + Ok(true) + } +} + #[tauri::command] pub(crate) async fn check_desktop_update( webview: tauri::Webview, diff --git a/studio/src-tauri/src/main.rs b/studio/src-tauri/src/main.rs index 8f36235dbf..d838b9d46b 100644 --- a/studio/src-tauri/src/main.rs +++ b/studio/src-tauri/src/main.rs @@ -31,7 +31,8 @@ use std::ffi::OsStr; use std::fs; use std::io::Write; use std::path::{Path, PathBuf}; -use std::sync::{Mutex, MutexGuard, Once, OnceLock}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Mutex, MutexGuard, OnceLock}; use tauri::menu::{MenuBuilder, MenuItemBuilder}; use tauri::tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent}; use tauri::{Emitter, Manager}; @@ -40,7 +41,7 @@ use tauri_plugin_window_state::{AppHandleExt, StateFlags}; /// Serializes the exit paths that reap the backend: `request_quit` (tray "Quit" and, /// outside macOS, the close button), the Unix signal listener, and `RunEvent::Exit`. /// Exactly one runs cleanup; the others block, so the process never exits mid-reap. -static TERMINATION_CLEANUP: Once = Once::new(); +static TERMINATION_CLEANUP: Mutex = Mutex::new(false); const IN_APP_RELAUNCH_MARKER_FILE: &str = "in-app-relaunch-v1"; @@ -786,6 +787,16 @@ fn confirm_quit_during_update(app: &tauri::AppHandle) -> bool { .blocking_show() } +/// Let a later exit run cleanup again. A cancelled or failed installer leaves the +/// app running with a spent guard, so the retry's pre-exit hook would reap nothing +/// and then suspend the job with the backend still holding the environment open. +fn reset_termination_cleanup() { + match TERMINATION_CLEANUP.lock() { + Ok(mut done) => *done = false, + Err(poisoned) => *poisoned.into_inner() = false, + } +} + /// Native AppKit termination needs to decide synchronously whether it can terminate now or must /// return NSTerminateLater while the shared confirmation sequence runs. #[cfg(target_os = "macos")] @@ -802,14 +813,21 @@ fn quit_requires_confirmation(app: &tauri::AppHandle) -> bool { } fn cleanup_child_processes(app: &tauri::AppHandle) { - // `call_once_force` rather than `call_once`: a panicking cleanup poisons the `Once` - // instead of deadlocking the exit paths waiting on it, and the next caller retries - // it rather than exiting on a backend that was never reaped. - TERMINATION_CLEANUP.call_once_force(|state| { - if state.is_poisoned() { + // A resettable guard rather than a `Once`, so `reset_termination_cleanup` can + // re-arm it. The flag is set after the body, so a panicking cleanup poisons the + // lock and the next caller retries it rather than exiting on a backend that was + // never reaped. + let mut done = match TERMINATION_CLEANUP.lock() { + Ok(guard) => guard, + Err(poisoned) => { warn!("Previous termination cleanup panicked, retrying it"); + poisoned.into_inner() } - + }; + if *done { + return; + } + { let diagnostics_state = app .try_state::() .map(|state| state.inner().clone()); @@ -831,7 +849,8 @@ fn cleanup_child_processes(app: &tauri::AppHandle) { .expect("ShutdownFlag must be managed"); let _ = process::stop_backend(&backend_state, &shutdown, diagnostics_state.as_ref()); } - }); + } + *done = true; } /// The backend is spawned as its own process-group leader, so nothing reaps it when the @@ -1523,6 +1542,8 @@ fn main() { desktop_update_policy::check_desktop_manual_update, desktop_update_policy::desktop_update_policy, desktop_updater::check_desktop_update, + desktop_updater::desktop_update_cleanup_armed, + desktop_updater::resume_desktop_update_cleanup, diagnostics::collect_support_diagnostics, native_clipboard::read_native_clipboard_files, native_clipboard::read_native_clipboard_png, diff --git a/studio/src-tauri/src/process.rs b/studio/src-tauri/src/process.rs index ff06b6260f..ad3f618ce5 100644 --- a/studio/src-tauri/src/process.rs +++ b/studio/src-tauri/src/process.rs @@ -1109,6 +1109,20 @@ pub fn start_backend( #[cfg(windows)] let _runtime_launch_guard = acquire_studio_runtime_launch_guard()?; + // A backend started while the job is disarmed is the orphan this guards + // against. The UI gate is per update action, and a webview remount starts + // one on its own, so the check belongs on the path that actually spawns. + #[cfg(windows)] + if !crate::windows_job::kill_on_close_armed().unwrap_or(false) { + crate::windows_job::resume_after_update_installer().map_err(|error| { + format!("Refusing to start the backend with crash cleanup disarmed: {error}") + })?; + // The same pair the UI's resume does: the pre-exit hook has already run + // its cleanup, and leaving that guard set means the next attempt's hook + // suspends kill-on-close without stopping this backend first. + crate::reset_termination_cleanup(); + } + let bin = match resolve_backend_binary() { Ok(bin) => bin, Err(msg) => { diff --git a/studio/src-tauri/src/windows_job.rs b/studio/src-tauri/src/windows_job.rs index cc6ec37e16..151b512954 100644 --- a/studio/src-tauri/src/windows_job.rs +++ b/studio/src-tauri/src/windows_job.rs @@ -97,8 +97,10 @@ unsafe fn set_kill_on_close(job: &OwnedHandle, enabled: bool) -> std::io::Result /// Keep the updater-launched installer alive after Tauri exits this process. /// -/// The updater calls this from its native pre-exit hook after extraction and -/// immediately before it launches the installer. +/// Called from the updater's pre-exit hook, just before the installer launches. +/// `cleanup_child_processes` has already reaped the backend tree by then, which +/// is what has to be gone: a child that outlives the app keeps +/// `%STUDIO_HOME%\unsloth_studio` open and the next update refuses to run. #[cfg(windows)] pub fn suspend_for_update_installer() -> std::io::Result<()> { let Some(job) = APP_JOB.get().and_then(Option::as_ref) else { @@ -110,6 +112,34 @@ pub fn suspend_for_update_installer() -> std::io::Result<()> { Ok(()) } +/// Whether kill-on-close is currently in force. +/// +/// The webview can reload after a suspension, which resets whatever the UI +/// remembered while the job stays disarmed, so the UI asks rather than assumes. +/// No job means nothing to disarm, and the caller has nothing to re-arm. +#[cfg(windows)] +pub fn kill_on_close_armed() -> std::io::Result { + let Some(job) = APP_JOB.get().and_then(Option::as_ref) else { + return Ok(true); + }; + + let limits = unsafe { query_job_limits(job) }?; + Ok(limits.BasicLimitInformation.LimitFlags & JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE != 0) +} + +/// Re-arm crash cleanup after a suspension that did not end in an exit (a failed +/// or abandoned install). Without it the app runs on with no reaper at all. +#[cfg(windows)] +pub fn resume_after_update_installer() -> std::io::Result<()> { + let Some(job) = APP_JOB.get().and_then(Option::as_ref) else { + return Ok(()); + }; + + unsafe { set_kill_on_close(job, true) }?; + info!("Windows job cleanup re-armed after the updater did not exit"); + Ok(()) +} + #[cfg(all(test, windows))] mod tests { use super::*;