Coverage for src/lilbee/providers/fleet/child_guard.py: 100%
87 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Bind a spawned llama-swap's lifetime to this process, so a crash cannot orphan it.
3A graceful exit tears the fleet down and a stale engine is reaped on the next
4launch; this closes the window between, when lilbee dies without cleanup (SIGKILL,
5segfault, closed terminal). Each platform binds differently and falls back to the
6next-launch reap when its primitive is unavailable, so a failure here never fails
7a spawn:
9* Linux: ``PR_SET_PDEATHSIG``, set in the preexec of one process-lifetime thread
10 (the signal binds to the forking thread, not the process).
11* Windows: a kill-on-close job object whose handle is held for the process life.
12* Elsewhere (macOS, any POSIX host without prctl): a death pipe (see
13 :func:`_watch_via_death_pipe`).
15``processfamily`` covers only prctl+job (needs pywin32, skips macOS), so the
16syscalls stay custom; the executor is the stdlib's.
17"""
19from __future__ import annotations
21import contextlib
22import ctypes
23import logging
24import os
25import subprocess
26import sys
27import threading
28from concurrent.futures import ThreadPoolExecutor
29from typing import TYPE_CHECKING, Any
31if TYPE_CHECKING:
32 from collections.abc import Callable
34log = logging.getLogger(__name__)
36# prctl(2) PR_SET_PDEATHSIG=1; SIGTERM so llama-swap runs its own shutdown.
37_PR_SET_PDEATHSIG = 1
38_PDEATHSIG = 15
40# Resolved at import so the post-fork child touches an already-loaded handle: a
41# dlopen after fork can deadlock on a lock a sibling thread holds.
42_libc: ctypes.CDLL | None = None
43if sys.platform.startswith("linux"): # pragma: no cover - Linux only
44 try:
45 _libc = ctypes.CDLL("libc.so.6", use_errno=True)
46 except OSError:
47 _libc = None
49# Windows job-object constants (winnt.h).
50_JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
51_JOB_OBJECT_EXTENDED_LIMIT_CLASS = 9
52_PROCESS_SET_QUOTA = 0x0100
53_PROCESS_TERMINATE = 0x0001
55# Live death-pipe write ends, keyed by guarded pid. os.pipe() returns bare fds the
56# GC never closes, so this is the only handle to close one deliberately (on child
57# stop) or leave for the kernel to close on our death (which fires the binding).
58_death_pipe_write_fds: dict[int, int] = {}
61def bind_lifetime_to_parent(parent_pid: int) -> None:
62 """Bind this process's lifetime to *parent_pid*, where the kernel supports it.
64 For a process lilbee spawned that holds resources of its own (an ingest
65 worker owns a GPU fleet): orphaning one leaves the card allocated with
66 nothing left to release it. A no-op off Linux. Exits when the parent is
67 already gone, which the death signal alone would miss.
68 """
69 if _libc is not None:
70 _libc.prctl(_PR_SET_PDEATHSIG, _PDEATHSIG, 0, 0, 0)
71 if os.getppid() != parent_pid:
72 raise SystemExit(1)
75def _make_pdeathsig_preexec(parent_pid: int) -> Callable[[], None] | None:
76 """A preexec binding the child's death to *parent_pid*, or None if libc is absent.
78 Compares against the real pid, not 1, so a lilbee running as pid 1 (a container
79 without an init shim) is not mistaken for a dead parent.
80 """
81 if _libc is None:
82 return None
83 libc = _libc
85 def _set_pdeathsig() -> None:
86 # In the forked child pre-exec: touch only pre-resolved handles, no alloc.
87 libc.prctl(_PR_SET_PDEATHSIG, _PDEATHSIG, 0, 0, 0)
88 if os.getppid() != parent_pid:
89 os._exit(1)
91 return _set_pdeathsig
94class _LifetimeSpawner:
95 """Runs spawns on one process-lifetime thread so PR_SET_PDEATHSIG binds to it.
97 The death signal binds to the forking thread, so a one-worker
98 ``ThreadPoolExecutor`` (reused for every spawn, stopped only at exit or the
99 test-only ``close``) keeps that thread alive for the process.
100 """
102 def __init__(self) -> None:
103 self._executor: ThreadPoolExecutor | None = None
104 self._lock = threading.Lock()
106 def spawn(self, *args: Any, **kwargs: Any) -> subprocess.Popen[Any]:
107 with self._lock:
108 if self._executor is None:
109 self._executor = ThreadPoolExecutor(
110 max_workers=1, thread_name_prefix="fleet-spawner"
111 )
112 executor = self._executor
113 return executor.submit(subprocess.Popen, *args, **kwargs).result()
115 def close(self) -> None:
116 """Stop the spawner thread. Test-only; in production it lives forever."""
117 with self._lock:
118 executor, self._executor = self._executor, None
119 if executor is not None:
120 executor.shutdown(wait=True)
123_spawner = _LifetimeSpawner()
126def _assign_to_kill_on_close_job(pid: int) -> None: # pragma: no cover - Windows only
127 """Put *pid* in a kill-on-close job object.
129 The job handle is deliberately leaked: holding it open is what kills the child
130 when this process ends, and the OS reclaims it then.
131 """
132 # windll is a Windows-only loader absent from other platforms' stubs; the
133 # Any alias keeps the checker quiet without an attr-defined ignore.
134 ct: Any = ctypes
135 kernel32 = ct.windll.kernel32
137 class JobObjectBasicLimitInformation(ctypes.Structure):
138 _fields_ = [
139 ("PerProcessUserTimeLimit", ctypes.c_int64),
140 ("PerJobUserTimeLimit", ctypes.c_int64),
141 ("LimitFlags", ctypes.c_uint32),
142 ("MinimumWorkingSetSize", ctypes.c_size_t),
143 ("MaximumWorkingSetSize", ctypes.c_size_t),
144 ("ActiveProcessLimit", ctypes.c_uint32),
145 ("Affinity", ctypes.c_size_t),
146 ("PriorityClass", ctypes.c_uint32),
147 ("SchedulingClass", ctypes.c_uint32),
148 ]
150 class IoCounters(ctypes.Structure):
151 _fields_ = [(name, ctypes.c_uint64) for name in ("r", "w", "o", "rb", "wb", "ob")]
153 class JobObjectExtendedLimitInformation(ctypes.Structure):
154 _fields_ = [
155 ("BasicLimitInformation", JobObjectBasicLimitInformation),
156 ("IoInfo", IoCounters),
157 ("ProcessMemoryLimit", ctypes.c_size_t),
158 ("JobMemoryLimit", ctypes.c_size_t),
159 ("PeakProcessMemoryUsed", ctypes.c_size_t),
160 ("PeakJobMemoryUsed", ctypes.c_size_t),
161 ]
163 job = kernel32.CreateJobObjectW(None, None)
164 if not job:
165 raise OSError("CreateJobObjectW failed")
166 info = JobObjectExtendedLimitInformation()
167 info.BasicLimitInformation.LimitFlags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
168 if not kernel32.SetInformationJobObject(
169 job, _JOB_OBJECT_EXTENDED_LIMIT_CLASS, ctypes.byref(info), ctypes.sizeof(info)
170 ):
171 raise OSError("SetInformationJobObject failed")
172 handle = kernel32.OpenProcess(_PROCESS_SET_QUOTA | _PROCESS_TERMINATE, False, pid)
173 if not handle:
174 raise OSError("OpenProcess failed")
175 try:
176 if not kernel32.AssignProcessToJobObject(job, handle):
177 raise OSError("AssignProcessToJobObject failed")
178 finally:
179 kernel32.CloseHandle(handle)
182def _watch_via_death_pipe(pid: int) -> None:
183 """Signal *pid* from a detached ``sh`` watcher once a pipe reaches EOF.
185 Portable stand-in for prctl/job objects: we hold the pipe's write end until the
186 child stops or we die, and the watcher (reading the read end as its stdin)
187 signals *pid* at EOF. The write end is O_CLOEXEC, so exec'd children never
188 inherit it; a fork-without-exec child would keep the pipe open past our death,
189 so a bound child's owner must not fork workers (serve is single-process).
191 The read end is passed as the watcher's stdin rather than redirected with
192 ``<&N``: dash mis-dups a multi-digit fd there, silently binding the wrong one.
193 Best-effort; ``kill -0`` skips an already-dead pid, and :func:`release_death_pipe`
194 keeps the pid-recycle window to the child's own stop.
195 """
196 try:
197 read_fd, write_fd = os.pipe()
198 except OSError:
199 log.info("Could not bind the engine to this process; the next launch reaps it.")
200 return
201 script = f"read -r _; kill -0 {pid} 2>/dev/null && kill {pid}"
202 try:
203 _spawner.spawn(
204 ["/bin/sh", "-c", script],
205 stdin=read_fd,
206 start_new_session=True,
207 stdout=subprocess.DEVNULL,
208 stderr=subprocess.DEVNULL,
209 )
210 except OSError:
211 os.close(write_fd)
212 log.info("Could not bind the engine to this process; the next launch reaps it.")
213 else:
214 # A crashed-without-release child can leave a stale entry the OS recycles
215 # this pid into; close it before the overwrite so it never leaks.
216 release_death_pipe(pid)
217 _death_pipe_write_fds[pid] = write_fd
218 finally:
219 os.close(read_fd)
222def release_death_pipe(pid: int) -> None:
223 """Close the death pipe guarding *pid* so its watcher wakes and exits.
225 Called when the guarded child is stopped while this process lives on (fleet
226 reload / model switch), instead of leaving the watcher parked until our death.
227 A no-op for a pid with no death pipe (kernel-bound platforms, or already released).
228 """
229 write_fd = _death_pipe_write_fds.pop(pid, None)
230 if write_fd is not None:
231 with contextlib.suppress(OSError):
232 os.close(write_fd)
235def spawn_bound_child(
236 argv: list[str],
237 *,
238 bind_lifetime: bool = True,
239 death_pipe: bool = True,
240 **popen_kwargs: Any,
241) -> subprocess.Popen[Any]:
242 """Spawn *argv*, by default bound to this process's lifetime.
244 ``bind_lifetime`` is False when the child is meant to outlive this process
245 (``keep_engine_warm``); the next-launch reap is then the only cleanup.
247 ``death_pipe`` is False for a short-lived child: the pipe's watcher lives until
248 we die, so binding a child that outlives neither costs a process and an fd for
249 nothing. Kernel bindings have no such cost and still apply.
251 Pass ``start_new_session=True`` to also put the child in its own group; the
252 binding does not depend on it.
253 """
254 if not bind_lifetime:
255 return _spawner.spawn(argv, **popen_kwargs)
257 preexec = _make_pdeathsig_preexec(os.getpid())
258 if preexec is not None:
259 try:
260 return _spawner.spawn(argv, preexec_fn=preexec, **popen_kwargs)
261 except OSError:
262 log.info("Could not set the death signal on the engine; trying the death pipe.")
264 proc = _spawner.spawn(argv, **popen_kwargs)
265 if sys.platform == "win32":
266 try:
267 _assign_to_kill_on_close_job(proc.pid)
268 except OSError:
269 log.info("Could not bind the engine to this process; the next launch reaps it.")
270 elif death_pipe:
271 _watch_via_death_pipe(proc.pid)
272 return proc