Coverage for src/lilbee/providers/fleet/swap_manager.py: 100%
643 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Supervise the single llama-swap process that fronts every fleet role.
3llama-swap owns each role's llama-server lifecycle; this manages the one proxy
4process and exposes its endpoint and readiness. See docs/architecture.md.
5"""
7from __future__ import annotations
9import contextlib
10import itertools
11import json
12import logging
13import os
14import signal
15import socket
16import subprocess
17import sys
18import threading
19import time
20from collections.abc import Iterable, Iterator
21from dataclasses import dataclass
22from enum import StrEnum
23from functools import lru_cache
24from pathlib import Path
25from typing import TYPE_CHECKING, BinaryIO
27import httpx
28import psutil
30from lilbee.core.health_warnings import HealthWarning
31from lilbee.providers.base import ProviderError, ProviderErrorKind
32from lilbee.providers.fleet.binary import engine_pin, resolve_llama_swap
33from lilbee.providers.fleet.child_guard import release_death_pipe, spawn_bound_child
34from lilbee.providers.fleet.gpu_env import engine_environ
35from lilbee.providers.fleet.groups import SwapGroup
36from lilbee.providers.fleet.launch import role_model_prefix
37from lilbee.providers.fleet.planning import clear_ctx_downshift, log_engine_launch
38from lilbee.providers.fleet.readback import (
39 MEMORY_FLAG,
40 check_launch,
41 check_memory_report,
42 report_missing_log,
43)
44from lilbee.providers.fleet.swap_config import PORT_FLAG, build_swap_config
45from lilbee.runtime.engine_lock import clear_keep_warm, reaper_running
47if TYPE_CHECKING:
48 from lilbee.providers.fleet.launch import InstanceLaunch
49 from lilbee.providers.roles import WorkerRole
51log = logging.getLogger(__name__)
53_HOST = "127.0.0.1"
54# One llama-swap per swap group: the group name lands in the config filename so
55# each group's processes are identified (and stopped) by their own config path,
56# and a placement change can restart one group without touching the others.
57# The writer pid segment is uniqueness, not ownership: the build lock ensures
58# one builder per engine dir, and reaping cleans dead writers' leftovers.
59_CONFIG_FILENAME_TEMPLATE = "llama-swap-{group}.{pid}.json"
60_CONFIG_FILE_GLOB = "llama-swap-*.json"
61# llama-swap's own stdout/stderr (its HTTP access log) is captured to a file in a
62# ``logs/`` dir inside the engine dir, which is the machine slot rather than any
63# one lilbee's data root, so the log sits beside the engine it belongs to instead
64# of beside server.log. Capturing it at all, rather than inheriting the parent's
65# fd, is because a TUI or CLI parent owns the terminal and an inherited fd would bleed
66# llama-swap's request log onto the screen and corrupt the render. Per-model
67# upstream logs are unaffected (those go to llama-swap's /logs API).
68_LOGS_SUBDIR = "logs"
69_LOG_FILENAME_TEMPLATE = "llama-swap-{group}.log"
70# Each writer's state file records its swap's pid/pgid so a later start can
71# stop a dead or unhealthy engine. Health, not ownership, decides sparing.
72_STATE_FILENAME_PREFIX = "llama-swap.state."
73_STATE_FILENAME_SUFFIX = ".json"
74# Also matches the legacy single shared state file ("llama-swap.state.json").
75_STATE_FILE_GLOB = f"{_STATE_FILENAME_PREFIX}*"
76_STATE_KEY_PID = "pid"
77_STATE_KEY_PGID = "pgid"
78_STATE_KEY_CREATED_AT = "created_at"
79_STATE_KEY_NAME = "name"
80_STATE_KEY_MEMBER_PORTS = "member_ports"
81_STATE_KEY_PROXY_PORT = "proxy_port"
82_STATE_KEY_LAUNCHES = "launches"
83_STATE_KEY_ENGINE_PIN = "engine_pin"
84_STATE_KEY_CONFIG_ROOT = "config_root"
85_STATE_KEY_FLEET_VIEW = "fleet_view"
86# Atomic state writes: the dot prefix keeps half-written tmp files out of the
87# reap scan's glob.
88_STATE_TMP_PREFIX = "."
89_STATE_TMP_SUFFIX = ".tmp"
90# Pid reuse guard: a live process at a recorded pid whose create time differs
91# from the recorded one by more than this is a different process.
92_CREATE_TIME_TOLERANCE_S = 1.0
93_LLAMA_SWAP_PROCESS_NAME = "llama-swap"
94_LLAMA_SERVER_PROCESS_NAME = "llama-server"
95_CONFIG_FLAG = "-config"
96_LISTEN_FLAG = "-listen"
97_HEALTH_PATH = "/health"
98_RUNNING_PATH = "/running"
99# Stops every upstream llama-server; llama-swap itself keeps listening.
100_UNLOAD_PATH = "/api/models/unload"
101_HTTP_TIMEOUT_S = 10.0
102# llama-swap's own proxy answers within a second; upstream model loads have their
103# own (longer) budget inside llama-swap, so this only covers the proxy coming up.
104_BOOT_TIMEOUT_S = 30.0
105_BOOT_POLL_S = 0.25
106# Cap on the captured llama-swap output a boot-failure error carries.
107_BOOT_LOG_TAIL_CHARS = 2000
108# Per-group SIGTERM grace before SIGKILL on the manager shutdown/reload path. A
109# hard kill is safe (llama-server holds no persistent state). Note this is NOT
110# the constant the serve handoff waits on: that path goes through stop_engine ->
111# _stop_stale_swap and spends _ORPHAN_STOP_TIMEOUT_S plus the kill/reap waits, so
112# SERVER_LOCK_TIMEOUT budgets only a teardown whose SIGTERMs are honored.
113_STOP_TIMEOUT_S = 2.5
114# Grace for a llama-server that outlived llama-swap before it is force-killed.
115_ORPHAN_STOP_TIMEOUT_S = 5.0
116# Grace for a SIGKILLed process to exit (and release its VRAM) before the next
117# free-memory probe runs.
118_KILL_WAIT_TIMEOUT_S = 5.0
119_PROBE_TIMEOUT_S = 5.0
120# Liveness probes talk to a loopback proxy, so they get their own short budget
121# rather than the module's 10 s general HTTP one. The ladder runs this probe for
122# every group while holding the cross-process build lock, so one wedged port
123# (SYN-accepted but unresponsive) would otherwise stall every other lilbee start
124# for tens of seconds. A local proxy that cannot answer /running this fast is
125# not usable for inference either.
126_LIVENESS_TIMEOUT = httpx.Timeout(connect=0.5, read=2.0, write=2.0, pool=2.0)
127# Run as `python -m`; the frozen build dispatches `-m lilbee.*` to the module's main().
128_REAPER_MODULE = "lilbee.providers.fleet.engine_reaper"
129_REAPER_LOG_FILENAME = "engine-reaper.log"
130# A reaper reports running within one interpreter start; a frozen onefile start
131# on a slow disk is the long case.
132_REAPER_READY_TIMEOUT_S = 10.0
133_REAPER_READY_POLL_S = 0.05
136@lru_cache(maxsize=1)
137def _probe_client() -> httpx.Client:
138 """One shared client for the localhost engine probes.
140 ``httpx.get`` builds a fresh ``Client`` per call, and every ``Client``
141 construction creates an SSL context, which loads the system CA bundle. These
142 probes are plain HTTP to 127.0.0.1, so none of that TLS setup is ever used --
143 and the readiness probe runs on the task bar's timer (up to 10 Hz), which made
144 ``ssl.create_default_context`` 23% of TUI CPU in a py-spy profile. One client
145 builds that at most once and keeps the connection alive between polls.
146 ``trust_env`` is off so a proxy env var cannot redirect a loopback probe.
147 """
148 return httpx.Client(trust_env=False)
151_PROVIDER = "llama-server"
152# /running JSON shape: {"running": [{"model": <id>, "state": "ready", ...}, ...]}.
153_KEY_RUNNING = "running"
154_KEY_MODEL = "model"
155_KEY_STATE = "state"
156_STATE_READY = "ready"
159def _platform_const(module: object, name: str, default: int) -> int:
160 """A platform-conditional stdlib constant (absent on some OSes -> default)."""
161 return getattr(module, name, default)
164_CREATE_NEW_PROCESS_GROUP: int = _platform_const(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
165_DETACHED_PROCESS: int = _platform_const(subprocess, "DETACHED_PROCESS", 0)
166_SIGKILL: int = _platform_const(signal, "SIGKILL", signal.SIGTERM)
169def _atomic_write(path: Path, text: str) -> None:
170 """Write *text* to *path* via a temp file in the same dir, then rename over it.
172 A plain write truncates the destination first, so a process dying mid-write
173 (OOM kill, SIGKILL, disk full) leaves an empty or half-written file behind.
174 For the llama-swap config that means the next spawn hands the engine a file
175 it cannot start from; for the state file it means a sibling's reap scan
176 reads a torn record.
178 The temp name carries the destination's name, and both config and state
179 filenames embed the writing process's pid, so a crash leftover can be told
180 from a live writer's file in flight -- see ``_clean_stale_tmp_files``.
181 """
182 tmp_path = path.with_name(f"{_STATE_TMP_PREFIX}{path.name}{_STATE_TMP_SUFFIX}")
183 tmp_path.write_text(text, encoding="utf-8")
184 os.replace(tmp_path, path)
187def _state_filename(owner_pid: int, group: str) -> str:
188 """The per-owner, per-group state filename for the lilbee process *owner_pid*."""
189 return f"{_STATE_FILENAME_PREFIX}{group}.{owner_pid}{_STATE_FILENAME_SUFFIX}"
192class EngineLifetime(StrEnum):
193 """What ends a spawned engine."""
195 BUILDER = "builder" # the process that spawned it exits or dies
196 USERS = "users" # the dir's reaper, once no user lock and no keep-warm opt-in is left
197 DETACHED = "detached" # nothing; the dir's reaper only unloads its weights
200@dataclass(frozen=True)
201class SwapState:
202 """A running llama-swap's recorded identity and serving contract.
204 Read back from the engine dir's state file, so it describes engines this
205 process did not start. The currency the bind/build ladder is written in:
206 swap_manager records it, provider reads it to decide what a slot is
207 serving, and contract matches it against what this process wants.
208 """
210 pid: int
211 pgid: int | None
212 created_at: float | None = None
213 member_ports: tuple[int, ...] = ()
214 proxy_port: int | None = None
215 launches: tuple[dict, ...] = ()
216 engine_pin: str | None = None
217 origin: PlanOrigin | None = None
220@dataclass(frozen=True)
221class PlanOrigin:
222 """The config root a fleet was planned for, and what its config.toml set for the plan then."""
224 config_root: str
225 fleet_view: str | None
228class SwapManager:
229 """Owns one llama-swap process fronting one role group's servers.
231 The provider runs one manager per role, so restarting a group (a placement
232 or model change) never touches another group's loaded servers.
233 """
235 def __init__(self, data_dir: Path, group: SwapGroup) -> None:
236 self._data_dir = data_dir
237 self._group = group
238 self._config_path = data_dir / _config_filename(os.getpid(), group.value)
239 self._log_path = data_dir / _LOGS_SUBDIR / _LOG_FILENAME_TEMPLATE.format(group=group.value)
240 # Instances whose engine report has already been compared to the estimate,
241 # so the check runs once per start rather than on every readiness poll.
242 self._estimate_checked: set[str] = set()
243 self._launch_by_model: dict[str, InstanceLaunch] = {}
244 # Placement divergences by model id, surfaced on health. Guarded because
245 # role_ready probes off the provider lock while health reads anytime.
246 self._placement_warnings: dict[str, HealthWarning] = {}
247 self._warnings_lock = threading.Lock()
248 self._state_path = data_dir / _state_filename(os.getpid(), group.value)
249 self._proc: subprocess.Popen[bytes] | None = None
250 self._log_file: BinaryIO | None = None
251 # Where this boot's output starts in the append-mode log file.
252 self._log_offset = 0
253 self._port: int | None = None
254 self._member_ports: list[int] = []
255 # Which member port serves which model id, for the direct GET /memory
256 # readback. Only launches this manager spawned are in here: a bound
257 # manager's state record carries the ports but not the mapping, and it
258 # has no launches to check either.
259 self._member_port_by_model: dict[str, int] = {}
260 # The serving contract (per-role model/ctx/slots) persisted in every
261 # state write, so a guest lilbee can bind to this live fleet.
262 self._launches_payload: list[dict] = []
263 self._origin: PlanOrigin | None = None
264 # True when this manager uses an engine another process built: it then
265 # never writes state, never reaps, and never signals engine processes.
266 self._bound = False
268 def start(
269 self,
270 launches: list[InstanceLaunch],
271 *,
272 ttl_seconds: int = 0,
273 lifetime: EngineLifetime = EngineLifetime.BUILDER,
274 origin: PlanOrigin | None = None,
275 ) -> None:
276 """Write the config and spawn llama-swap, waiting for its proxy to answer.
278 The proxy and every member get a freshly allocated free port, which is
279 why llama-swap's own startPort is not used: that assigns a fixed
280 sequential range at config load, so it would collide with a previous
281 instance's server still shutting down (the new llama-server then fails
282 its bind and llama-swap reports it only as "exited prematurely").
284 This narrows that collision rather than removing a race. The ports are
285 picked by binding and closing ephemeral sockets, while llama-swap starts
286 each upstream lazily on its first request, so a member port can sit
287 unbound for as long as it takes that request to arrive and anything else
288 on the box may take it in between. Nothing in llama-swap offers a
289 spawn-time probe to close that window; warming the roles up front
290 shortens it for the roles that are warmed.
292 *lifetime* says what ends the engine. ``USERS`` needs the caller to hold
293 the dir's build lock, and falls back to ``BUILDER`` when no reaper reports
294 running for the dir, so the engine is never left unbound and unwatched.
295 *origin* goes into the state, so a peer can see a fleet planned before a config write.
296 """
297 self._origin = origin
298 # Idempotent safety net; the provider reaps before planning so the GPU
299 # probe already saw the real free memory.
300 self.reap_stale()
301 # Singleton guard: one llama-swap per data_dir for this lilbee. Reap any
302 # llama-swap we already started against this config (a leaked duplicate
303 # from a prior race/reload) before spawning, so they cannot accumulate
304 # and double-book a GPU.
305 _stop_own_fleet(self._config_path, tuple(self._member_ports))
306 ports = _pick_free_ports(1 + len(launches))
307 member_ports = dict(zip([launch.model_id for launch in launches], ports[1:], strict=True))
308 self._member_ports = sorted(member_ports.values())
309 self._member_port_by_model = dict(member_ports)
310 self._launches_payload = [launch.to_state() for launch in launches]
311 self._launch_by_model = {launch.model_id: launch for launch in launches}
312 self._estimate_checked.clear()
313 with self._warnings_lock:
314 self._placement_warnings.clear()
315 self._config_path.parent.mkdir(parents=True, exist_ok=True)
316 self._log_path.parent.mkdir(parents=True, exist_ok=True)
317 _atomic_write(
318 self._config_path,
319 build_swap_config(
320 launches,
321 member_ports,
322 swap=self._group.swaps,
323 ttl_seconds=ttl_seconds,
324 engine_log_dir=self._log_path.parent,
325 ),
326 )
327 self._port = ports[0]
328 # Capture llama-swap's stdout/stderr to a file so its access log never
329 # reaches an inherited terminal (a TUI/CLI parent) and garbles the screen.
330 self._close_log()
331 self._log_path.parent.mkdir(parents=True, exist_ok=True)
332 self._log_file = self._log_path.open("ab")
333 self._log_offset = self._log_path.stat().st_size
334 if lifetime is EngineLifetime.USERS and not ensure_engine_reaper(self._data_dir):
335 log.warning("The engine at %s stops with this process.", self._data_dir)
336 lifetime = EngineLifetime.BUILDER
337 self._proc = spawn_bound_child(
338 [
339 str(resolve_llama_swap()),
340 _CONFIG_FLAG,
341 str(self._config_path),
342 _LISTEN_FLAG,
343 f"{_HOST}:{self._port}",
344 ],
345 bind_lifetime=lifetime is EngineLifetime.BUILDER,
346 env=engine_environ(),
347 stdout=self._log_file,
348 stderr=subprocess.STDOUT,
349 start_new_session=True,
350 creationflags=_CREATE_NEW_PROCESS_GROUP,
351 )
352 self._write_state()
353 self._await_health()
354 for launch in launches:
355 log_engine_launch(launch)
357 def reap_stale(self) -> None:
358 """Kill every dead or unhealthy recorded engine; see :func:`reap_stale`."""
359 reap_stale(self._data_dir)
361 def _process_identity(self) -> tuple[int, int | None, float | None] | None:
362 """(pid, pgid, create time) of the swap this manager runs, or None."""
363 if self._proc is not None:
364 pid = self._proc.pid
365 pgid: int | None = None
366 if sys.platform != "win32":
367 with contextlib.suppress(ProcessLookupError):
368 pgid = os.getpgid(pid)
369 created_at: float | None = None
370 with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied):
371 created_at = psutil.Process(pid).create_time()
372 return pid, pgid, created_at
373 return None
375 def _write_state(self) -> None:
376 """Record the swap's pid/pgid/create time, member ports, and our identity.
378 The write is atomic (tmp file then ``os.replace``) so a sibling's reap
379 scan can never read a torn file and mistake this live record for junk.
380 """
381 identity = self._process_identity()
382 if identity is None:
383 return
384 swap_pid, pgid, created_at = identity
385 state = {
386 _STATE_KEY_PID: swap_pid,
387 _STATE_KEY_PGID: pgid,
388 _STATE_KEY_CREATED_AT: created_at,
389 _STATE_KEY_NAME: _LLAMA_SWAP_PROCESS_NAME,
390 _STATE_KEY_MEMBER_PORTS: self._member_ports,
391 _STATE_KEY_PROXY_PORT: self._port,
392 _STATE_KEY_LAUNCHES: self._launches_payload,
393 _STATE_KEY_ENGINE_PIN: engine_pin(),
394 }
395 if self._origin is not None:
396 state[_STATE_KEY_CONFIG_ROOT] = self._origin.config_root
397 state[_STATE_KEY_FLEET_VIEW] = self._origin.fleet_view
398 _atomic_write(self._state_path, json.dumps(state))
400 def restamp(self, origin: PlanOrigin) -> None:
401 """Record *origin* for the running swap: its plan still matches config.toml."""
402 self._origin = origin
403 self._write_state()
405 @property
406 def log_path(self) -> Path:
407 """The llama-swap log for this group, which records each member's exit."""
408 return self._log_path
410 def endpoint(self) -> str:
411 """Base URL of the llama-swap OpenAI-compatible proxy."""
412 if self._port is None:
413 raise ProviderError(
414 "The local model engine is not running.",
415 provider=_PROVIDER,
416 kind=ProviderErrorKind.SERVER,
417 )
418 return f"http://{_HOST}:{self._port}"
420 def role_ready(self, role: WorkerRole) -> bool:
421 """Whether at least one of *role*'s replica servers is loaded and ready."""
422 prefix = role_model_prefix(role)
423 ready = self._ready_models()
424 self._check_estimates(ready)
425 return any(model.startswith(prefix) for model in ready)
427 def _check_estimates(self, ready: set[str]) -> None:
428 """Compare each newly-ready engine's own report against what it was planned for.
430 The plan is otherwise open-loop, and a wrong estimate only ever surfaces
431 as a failed request much later. Once per instance per start: readiness is
432 polled, and the answer does not change once the engine has loaded.
433 """
434 for model_id in ready - self._estimate_checked:
435 launch = self._launch_by_model.get(model_id)
436 self._estimate_checked.add(model_id)
437 if launch is None:
438 continue
439 # Ready means this role's context loaded, so any reduction taken to
440 # get here has done its job and must not follow the role into the
441 # next plan, a freed machine, or a model the user switched to.
442 clear_ctx_downshift(launch.role)
443 # A launch carrying --memory serves its own report on GET /memory
444 # and was given no trace log to read (swap_config), so the log-side
445 # checks would misfire on it by construction.
446 if MEMORY_FLAG in launch.argv:
447 self._check_memory_estimate(model_id, launch)
448 continue
449 # The engine is ready, so a missing log is not "too early" any more.
450 if report_missing_log(self._log_path.parent, model_id, launch.role):
451 continue
452 if launch.est_vram_bytes <= 0:
453 continue
454 warning = check_launch(
455 self._log_path.parent,
456 model_id,
457 launch.role,
458 launch.model,
459 launch.est_vram_bytes,
460 launch.est_vram_by_device,
461 launch.est_unreported_bytes,
462 launch.est_host_bytes,
463 )
464 if warning is not None:
465 with self._warnings_lock:
466 self._placement_warnings[model_id] = warning
468 def _check_memory_estimate(self, model_id: str, launch: InstanceLaunch) -> None:
469 """Run the estimate check off the engine's own ``GET /memory``.
471 Straight to the member's port rather than through llama-swap: the model
472 is ready, so the server owns its port, and the proxy adds only a routing
473 layer that has nothing to route on for a bare GET. An endpoint that does
474 not answer after the engine took the flag is reported, not swallowed --
475 it is this mode's analog of a ready engine that wrote no log.
476 """
477 port = self._member_port_by_model.get(model_id)
478 if port is None: # bound to another process's engine; nothing was planned here
479 return
480 try:
481 resp = _probe_client().get(f"http://{_HOST}:{port}/memory", timeout=_PROBE_TIMEOUT_S)
482 payload = resp.json() if resp.status_code == httpx.codes.OK else None
483 except (httpx.HTTPError, ValueError):
484 payload = None
485 if payload is None:
486 log.warning(
487 "The %s engine was launched with %s but its /memory endpoint did not "
488 "answer, so its memory use could not be checked against the estimate. "
489 "Placement estimates for this model are unverified.",
490 launch.role.value,
491 MEMORY_FLAG,
492 )
493 return
494 if launch.est_vram_bytes <= 0:
495 return
496 warning = check_memory_report(
497 launch.role,
498 launch.model,
499 launch.est_vram_bytes,
500 launch.est_vram_by_device,
501 payload,
502 launch.est_host_bytes,
503 )
504 if warning is not None:
505 with self._warnings_lock:
506 self._placement_warnings[model_id] = warning
508 def health_warnings(self) -> list[HealthWarning]:
509 """Placement divergences recorded from ready engines in this group."""
510 with self._warnings_lock:
511 return list(self._placement_warnings.values())
513 def is_live(self) -> bool:
514 """Whether the swap process is up and its proxy answers ``/running``."""
515 if self._proc is None or self._proc.poll() is not None:
516 return False
517 if self._port is None:
518 return False
519 return self._proxy_answers()
521 @property
522 def running(self) -> bool:
523 """Whether this manager currently has a spawned llama-swap process."""
524 return self._proc is not None
526 @property
527 def bound(self) -> bool:
528 """Whether this manager rides an engine built by another process."""
529 return self._bound
531 def bind(self, state: SwapState) -> bool:
532 """Use a running engine's proxy without taking any ownership of it.
534 The engine's own state record stays untouched: the binder writes
535 nothing, and shutdown() merely drops the binding.
536 """
537 if state.proxy_port is None:
538 return False
539 self._port = state.proxy_port
540 self._member_ports = list(state.member_ports)
541 if not self._proxy_answers():
542 self._port = None
543 self._member_ports = []
544 return False
545 self._launches_payload = [dict(launch) for launch in state.launches]
546 self._bound = True
547 return True
549 def _proxy_answers(self) -> bool:
550 """Whether the bound proxy port serves llama-swap's running endpoint.
552 Shares state_is_healthy's identity check via _running_endpoint_answers, so
553 bind and reap agree on what "answering" means by construction rather than by
554 two hand-kept-identical copies.
555 """
556 return _running_endpoint_answers(self.endpoint())
558 def shutdown(self) -> None:
559 """Stop every llama-swap this lilbee owns at our config and reap servers.
561 Authoritative teardown keyed on config-path identity, not the single
562 tracked ``Popen``: a warm-up/reset race or a reload can leave several
563 llama-swap processes this lilbee spawned, any of them reparented to init
564 (still holding the engine binary open) -- trusting one handle would leak
565 them. Every llama-swap running against our config is reaped. Unlinks only
566 this owner's state file; another instance's record stays.
567 """
568 if self._bound:
569 # Not ours to stop: drop the binding and leave the engine serving.
570 self._bound = False
571 self._port = None
572 self._member_ports = []
573 self._launches_payload = []
574 return
575 _stop_own_fleet(self._config_path, tuple(self._member_ports))
576 # Nothing is coming back to bind these, so the picker can offer them again.
577 release_reserved_ports([*self._member_ports, *([self._port] if self._port else [])])
578 self._state_path.unlink(missing_ok=True)
579 if self._proc is not None:
580 # Free this engine's death pipe so its watcher exits now, not at our death.
581 release_death_pipe(self._proc.pid)
582 self._proc = None
583 self._port = None
584 self._close_log()
586 def _close_log(self) -> None:
587 """Close the captured llama-swap log handle, if one is open."""
588 if self._log_file is not None:
589 with contextlib.suppress(OSError):
590 self._log_file.close()
591 self._log_file = None
593 def _await_health(self) -> None:
594 """Poll the proxy's /health until it answers, or fail with a clear error."""
595 url = f"{self.endpoint()}{_HEALTH_PATH}"
596 deadline = time.monotonic() + _BOOT_TIMEOUT_S
597 while time.monotonic() < deadline:
598 if self._proc is not None and self._proc.poll() is not None:
599 self._fail("The local model engine exited before it was ready.")
600 with contextlib.suppress(httpx.HTTPError):
601 if _probe_client().get(url, timeout=_PROBE_TIMEOUT_S).status_code == httpx.codes.OK:
602 return
603 time.sleep(_BOOT_POLL_S)
604 self._fail("The local model engine did not start in time.")
606 def _ready_models(self) -> set[str]:
607 """Model ids whose upstream is loaded and ready, per llama-swap's /running.
609 A read-only probe: a concurrent shutdown can clear ``_port`` between the
610 caller's check and ``endpoint()``, raising ProviderError, so that is
611 suppressed too and the probe reports "nothing ready" rather than throwing.
612 """
613 with contextlib.suppress(httpx.HTTPError, ValueError, KeyError, TypeError, ProviderError):
614 payload = (
615 _probe_client()
616 .get(f"{self.endpoint()}{_RUNNING_PATH}", timeout=_PROBE_TIMEOUT_S)
617 .json()
618 )
619 return {
620 entry[_KEY_MODEL]
621 for entry in payload[_KEY_RUNNING]
622 if entry.get(_KEY_STATE) == _STATE_READY
623 }
624 return set()
626 def _boot_log_tail(self) -> str:
627 """The current boot's captured llama-swap output, capped for an error message."""
628 try:
629 with self._log_path.open("rb") as handle:
630 handle.seek(self._log_offset)
631 data = handle.read()
632 except OSError:
633 return ""
634 return data.decode(errors="replace").strip()[-_BOOT_LOG_TAIL_CHARS:]
636 def _fail(self, message: str) -> None:
637 """Tear down and raise a user-facing engine-start error carrying the boot log."""
638 self.shutdown()
639 tail = self._boot_log_tail()
640 if tail:
641 message = f"{message} Engine log ({self._log_path}):\n{tail}"
642 raise ProviderError(message, provider=_PROVIDER, kind=ProviderErrorKind.SERVER)
645# Linux publishes the range here; every other platform is asked via sysctl.
646_PROC_PORT_RANGE = Path("/proc/sys/net/ipv4/ip_local_port_range")
649def _port_range_from(path: Path) -> tuple[int, int] | None:
650 """The two integers in *path*, or ``None`` when it is absent or unreadable."""
651 try:
652 low, high = path.read_text(encoding="utf-8").split()[:2]
653 return int(low), int(high)
654 except (OSError, ValueError):
655 return None
658def _ephemeral_range() -> tuple[int, int] | None:
659 """The port range the kernel hands out for unbound sockets, if it says.
661 ``None`` when neither source answers, which is the signal to fall back to
662 letting the OS choose.
663 """
664 from_proc = _port_range_from(_PROC_PORT_RANGE)
665 if from_proc is not None:
666 return from_proc
667 try: # macOS and the BSDs, which have no procfs entry for this
668 out = subprocess.run(
669 ["/usr/sbin/sysctl", "-n", "net.inet.ip.portrange.first", "net.inet.ip.portrange.last"],
670 capture_output=True,
671 text=True,
672 encoding="utf-8",
673 errors="replace",
674 timeout=5,
675 check=False,
676 )
677 low, high = out.stdout.split()[:2]
678 return int(low), int(high)
679 except (OSError, ValueError, subprocess.SubprocessError):
680 return None
683# Where lilbee looks for engine ports when the kernel's ephemeral range is known.
684# Above the registered-service crowd, below every default ephemeral range.
685_PORT_SEARCH_FLOOR = 20000
686_PORT_WINDOW_SPAN = 8192
687# Block width. Each process searches one block, so concurrent lilbees hold
688# disjoint ranges. A fleet takes one proxy port plus one per member, and embed
689# and vision replicate per GPU, so 64 covers a 30-GPU host; a wider fleet spills
690# into the next block.
691_PORT_BLOCK = 64
693# Ports handed to a child that has not bound them yet. llama-swap binds a member
694# port only on that member's first request, so the probe socket is long closed
695# by then and the port looks free to every later probe. Without this the picker
696# hands the next group exactly what it gave the last one, every time.
697_reserved_ports: set[int] = set()
698_reserved_lock = threading.Lock()
701def release_reserved_ports(ports: Iterable[int]) -> None:
702 """Give *ports* back to the picker, once nothing is expected to bind them."""
703 with _reserved_lock:
704 _reserved_ports.difference_update(ports)
707def _window_span(ceiling: tuple[int, int]) -> int:
708 """How many ports below the ephemeral floor this host leaves to search."""
709 return max(1, min(ceiling[0], _PORT_SEARCH_FLOOR + _PORT_WINDOW_SPAN) - _PORT_SEARCH_FLOOR)
712def _search_start(ceiling: tuple[int, int]) -> int:
713 """First port of the block this process owns.
715 The pid selects a whole block, not an offset: reservation is per-process, and
716 a fleet takes its ports contiguously, so pid-offset starts one apart overlap
717 on all but one port.
718 """
719 blocks = max(1, _window_span(ceiling) // _PORT_BLOCK)
720 return _PORT_SEARCH_FLOOR + (os.getpid() % blocks) * _PORT_BLOCK
723def _pick_free_ports(count: int) -> list[int]:
724 """Bind *count* free localhost ports at once and return them.
726 All sockets stay open until every port is claimed so the OS cannot hand the
727 same port out twice within one allocation.
729 Picked from below the kernel's ephemeral range rather than inside it. The
730 gap between lilbee picking a port and llama-server binding it spans the whole
731 lazy-spawn wait, and a port inside the ephemeral range can be handed to any
732 passing outbound connection during that gap; one below it cannot be handed to
733 anybody, so the only way to lose it is another server binding that exact port
734 on purpose. Falls back to letting the OS choose when the range is unknown.
735 """
736 ceiling = _ephemeral_range()
737 sockets = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(count)]
738 try:
739 for sock in sockets:
740 _bind_below_ephemeral(sock, ceiling)
741 return [int(sock.getsockname()[1]) for sock in sockets]
742 finally:
743 for sock in sockets:
744 sock.close()
747def _bind_below_ephemeral(sock: socket.socket, ceiling: tuple[int, int] | None) -> None:
748 """Bind *sock* to a free, unreserved port under the ephemeral floor.
750 Falls back to letting the OS choose when the range is unknown or the window
751 is used up, which keeps a fleet start working at the cost of returning to the
752 ephemeral range for those ports.
753 """
754 if ceiling is not None and ceiling[0] > _PORT_SEARCH_FLOOR:
755 span = _window_span(ceiling)
756 start = _search_start(ceiling)
757 for offset in range(span):
758 port = _PORT_SEARCH_FLOOR + (start - _PORT_SEARCH_FLOOR + offset) % span
759 with _reserved_lock:
760 if port in _reserved_ports:
761 continue
762 try:
763 sock.bind((_HOST, port))
764 except OSError:
765 continue
766 _reserved_ports.add(port)
767 return
768 sock.bind((_HOST, 0))
771def _live_children(pid: int) -> list[psutil.Process]:
772 """The process's current descendants, or none when it already exited."""
773 try:
774 children: list[psutil.Process] = psutil.Process(pid).children(recursive=True)
775 except psutil.NoSuchProcess:
776 return []
777 return children
780def _reap_survivors(children: list[psutil.Process]) -> None:
781 """Terminate then kill any captured child that is still running."""
782 survivors = [child for child in children if child.is_running()]
783 for child in survivors:
784 with contextlib.suppress(psutil.NoSuchProcess):
785 child.terminate()
786 _, alive = psutil.wait_procs(survivors, timeout=_ORPHAN_STOP_TIMEOUT_S)
787 for child in alive:
788 with contextlib.suppress(psutil.NoSuchProcess):
789 child.kill()
790 _await_killed(alive)
793def _await_killed(procs: list[psutil.Process]) -> None:
794 """Wait for SIGKILLed processes to exit so their VRAM is free before any probe."""
795 if not procs:
796 return
797 _, alive = psutil.wait_procs(procs, timeout=_KILL_WAIT_TIMEOUT_S)
798 for proc in alive:
799 log.warning("Process %s survived SIGKILL; its VRAM may still be held.", proc.pid)
802def _processes_named(needle: str) -> Iterator[psutil.Process]:
803 """Live processes whose executable name contains *needle*.
805 ``name()`` is a cheap field (comm/proc_name); ``cmdline()`` reads the full
806 argument vector and on macOS blocks on entitlement-protected binaries. So the
807 name is the pre-filter and callers pay for ``cmdline()`` only on a match,
808 which keeps a full-process-table scan from stalling on an unrelated process.
809 """
810 for proc in psutil.process_iter(["name"]):
811 # process_iter already skips processes that vanish mid-scan and, per its
812 # ad_value contract, leaves ``name`` as None where it could not be read.
813 name = proc.info["name"] or ""
814 if needle in name:
815 yield proc
818def _swaps_for_config(config_path: Path) -> list[psutil.Process]:
819 """Every live llama-swap (any owner) running against *config_path*.
821 Identity is the ``-config <path>`` argument, which every llama-swap this
822 lilbee starts carries and which survives reparenting to init -- so this finds
823 a leaked duplicate or a swap reparented away from us, neither of which a
824 tracked Popen handle nor a ``children()`` scan would catch.
825 """
826 target = str(config_path)
827 swaps: list[psutil.Process] = []
828 for proc in _processes_named(_LLAMA_SWAP_PROCESS_NAME):
829 try:
830 cmdline = proc.cmdline()
831 except (
832 psutil.NoSuchProcess,
833 psutil.AccessDenied,
834 psutil.ZombieProcess,
835 OSError,
836 SystemError,
837 ):
838 # OSError/SystemError: macOS psutil mishandles entitlement-protected
839 # binaries (sysctl KERN_PROCARGS2), leaking a raw PermissionError or a
840 # C-extension SystemError instead of an AccessDenied.
841 continue
842 # Identity is the -config path; _processes_named already gated on comm.
843 if target in cmdline:
844 swaps.append(proc)
845 return swaps
848def find_live_state(data_dir: Path, group: SwapGroup) -> SwapState | None:
849 """The newest recorded state for *group* at *data_dir* (no liveness check).
851 A record's presence does not prove the engine is up; callers that need that
852 probe it with ``state_is_healthy``. The name reflects that a record is written
853 only for a running engine, not that this function verifies it.
854 """
855 best: SwapState | None = None
856 for state_path in sorted(data_dir.glob(_STATE_FILE_GLOB)):
857 if f".{group.value}." not in f".{state_path.name}":
858 continue
859 state = _load_state(state_path)
860 if state is None:
861 continue
862 if best is None or (state.created_at or 0) > (best.created_at or 0):
863 best = state
864 return best
867def _running_endpoint_answers(base_url: str) -> bool:
868 """Whether *base_url* serves llama-swap's ``/running`` endpoint (identity, not
869 just liveness).
871 Proxy ports are ephemeral: after an engine dies, any unrelated local service
872 that later binds the recorded port and returns a 2xx/3xx to an unknown path
873 would pass a bare status check, so a dead record would look healthy forever and
874 inference clients would bind to a non-engine endpoint. Requiring the ``running``
875 JSON payload shape that only llama-swap produces makes the probe identity-checked.
876 Total: any transport error or non-conforming body reads as "not our engine".
877 """
878 try:
879 resp = _probe_client().get(f"{base_url}{_RUNNING_PATH}", timeout=_LIVENESS_TIMEOUT)
880 except (OSError, httpx.HTTPError):
881 return False
882 if resp.status_code >= httpx.codes.BAD_REQUEST:
883 return False
884 try:
885 return isinstance(resp.json().get(_KEY_RUNNING), list)
886 except (ValueError, AttributeError):
887 return False
890def state_is_healthy(state: SwapState) -> bool:
891 """Whether the engine behind *state* answers on its recorded proxy port."""
892 if state.proxy_port is None:
893 return False
894 return _running_endpoint_answers(f"http://{_HOST}:{state.proxy_port}")
897def engine_record_exists(data_dir: Path) -> bool:
898 """Whether any engine state file is present, without probing proxy health.
900 A filesystem fact, unlike a proxy HTTP probe: it is true for an engine that
901 is live but momentarily unprobeable (fd exhaustion, host thrash), so the
902 ladder can clear a recorded engine before building rather than double-build
903 beside one an HTTP probe failed to see.
904 """
905 return any(data_dir.glob(_STATE_FILE_GLOB))
908def stop_engine(data_dir: Path) -> list[str]:
909 """Stop every engine the dir's state files record, regardless of liveness.
911 The unconditional off switch behind ``lilbee engine stop`` and the
912 last-user-out path: each recorded swap is terminated through its state
913 record (never a Popen handle, so it works on engines this process did
914 not build) and its file removed. A record whose llama-swap is already dead
915 still has its llama-servers (each in its own process group) reaped by
916 recorded port, exactly as reap_stale does -- otherwise the off switch would
917 leave those orphans holding VRAM and delete the ports needed to find them.
918 Stale config files for dead owners are cleaned too, and the persistence
919 opt-in is dropped with the engine it described, so the dir is left as
920 clean as a reap leaves it. Unparseable files are left alone, as in
921 reap_stale: they may be a sibling's in-flight write. Returns the group tokens
922 whose engine was actually alive, so a caller reports only real stops.
923 """
924 _clean_stale_configs(data_dir)
925 # The persistence opt-in describes the engine instance being stopped, so it
926 # dies with it. Cleared here rather than at each call site so no stop path
927 # can leave a mark that makes the next engine sticky-warm.
928 clear_keep_warm(data_dir)
929 stopped: list[str] = []
930 for state_path in sorted(data_dir.glob(_STATE_FILE_GLOB)):
931 state = _load_state(state_path)
932 if state is None:
933 continue
934 if _stop_recorded_engine(state):
935 group = _state_group(state_path.name)
936 if group is not None:
937 stopped.append(group)
938 state_path.unlink(missing_ok=True)
939 return stopped
942def unload_engine_models(data_dir: Path) -> bool:
943 """Unload the models of every engine *data_dir* records; False when any proxy refused.
945 The llama-swap processes stay up, so the next request reloads on demand.
946 """
947 states = (_load_state(path) for path in sorted(data_dir.glob(_STATE_FILE_GLOB)))
948 ports = [state.proxy_port for state in states if state and state.proxy_port is not None]
949 return all([_unload_models(port) for port in ports])
952def engine_answers(data_dir: Path) -> bool:
953 """Whether any engine *data_dir* records answers on its proxy port."""
954 states = (_load_state(path) for path in sorted(data_dir.glob(_STATE_FILE_GLOB)))
955 return any(state_is_healthy(state) for state in states if state)
958def _unload_models(proxy_port: int) -> bool:
959 url = f"http://{_HOST}:{proxy_port}{_UNLOAD_PATH}"
960 try:
961 resp = _probe_client().post(url, timeout=_HTTP_TIMEOUT_S)
962 except (OSError, httpx.HTTPError) as exc:
963 log.warning("Could not unload the models at %s: %s", url, exc)
964 return False
965 if resp.status_code >= httpx.codes.BAD_REQUEST:
966 log.warning("The engine at %s refused the unload: HTTP %d", url, resp.status_code)
967 return False
968 return True
971def ensure_engine_reaper(data_dir: Path) -> bool:
972 """Make sure a live reaper watches *data_dir*; False when none reports running.
974 A reaper reports running by holding the dir's reaper lock. Waiting for that,
975 not only for the launch, also catches a reaper that dies on import.
976 """
977 if reaper_running(data_dir):
978 return True
979 reaper_log = data_dir / _LOGS_SUBDIR / _REAPER_LOG_FILENAME
980 try:
981 reaper_log.parent.mkdir(parents=True, exist_ok=True)
982 log_offset = reaper_log.stat().st_size if reaper_log.exists() else 0
983 # The child keeps its own copy of the handle; the parent closes its copy at once.
984 with reaper_log.open("ab") as reaper_err:
985 reaper = subprocess.Popen( # noqa: S603 - argv is this interpreter and a fixed module
986 [sys.executable, "-m", _REAPER_MODULE, str(data_dir)],
987 stdin=subprocess.DEVNULL,
988 stdout=subprocess.DEVNULL,
989 stderr=reaper_err,
990 start_new_session=True,
991 creationflags=_CREATE_NEW_PROCESS_GROUP | _DETACHED_PROCESS,
992 )
993 except OSError:
994 log.warning("Could not launch the engine reaper for %s.", data_dir)
995 return False
996 deadline = time.monotonic() + _REAPER_READY_TIMEOUT_S
997 while time.monotonic() < deadline and reaper.poll() is None:
998 if reaper_running(data_dir):
999 return True
1000 time.sleep(_REAPER_READY_POLL_S)
1001 # An exited reaper either lost the claim to a live one or failed to start.
1002 if reaper_running(data_dir):
1003 return True
1004 log.warning(
1005 "The engine reaper for %s did not report running. Its log is %s.%s",
1006 data_dir,
1007 reaper_log,
1008 _log_tail_since(reaper_log, log_offset),
1009 )
1010 return False
1013def _log_tail_since(path: Path, offset: int) -> str:
1014 """The text written to *path* after *offset*, capped, as a message suffix ('' when none)."""
1015 try:
1016 with path.open("rb") as handle:
1017 handle.seek(offset)
1018 data = handle.read()
1019 except OSError:
1020 return ""
1021 text = data.decode(errors="replace").strip()[-_BOOT_LOG_TAIL_CHARS:]
1022 return f"\n{text}" if text else ""
1025def reap_stale(data_dir: Path) -> None:
1026 """Kill every dead or unhealthy recorded engine at *data_dir*.
1028 An OOM-killed lilbee leaves llama-swap (and its servers) holding VRAM,
1029 so planning would otherwise see artificially reduced free memory; the
1030 ladder calls this before its GPU probe. Every state file is scanned
1031 (all groups, including legacy names): an engine that is alive AND
1032 answering on its proxy is spared regardless of who started it (a
1033 reload's own healthy groups, or a bindable engine the ladder skipped);
1034 everything else is stopped through its record and its file removed. An
1035 unparseable file is skipped, never deleted: it may be a sibling's
1036 in-flight write. When the swap itself is dead, its servers (each in
1037 its own process group) may still be alive holding VRAM; they are
1038 matched by name plus recorded member port and stopped before the file
1039 is removed.
1041 Module-level (not a method) because it must run before planning decides
1042 which role groups exist, when no per-group manager has been built yet.
1043 """
1044 _clean_stale_tmp_files(data_dir)
1045 _clean_stale_configs(data_dir)
1046 for state_path in sorted(data_dir.glob(_STATE_FILE_GLOB)):
1047 state = _load_state(state_path)
1048 if state is None:
1049 continue
1050 if state_is_healthy(state):
1051 # An answering engine is in use (bind accepts on exactly this
1052 # test); reaping must never disagree with binding.
1053 continue
1054 _stop_recorded_engine(state)
1055 state_path.unlink(missing_ok=True)
1058def _clean_stale_tmp_files(data_dir: Path) -> None:
1059 """Remove crash-leftover state and config tmp files whose writer is dead."""
1060 tmp_glob = f"{_STATE_TMP_PREFIX}*{_STATE_TMP_SUFFIX}"
1061 for tmp_path in data_dir.glob(tmp_glob):
1062 writer_pid = _state_owner_pid(tmp_path.name)
1063 if writer_pid is not None and not psutil.pid_exists(writer_pid):
1064 tmp_path.unlink(missing_ok=True)
1067def _clean_stale_configs(data_dir: Path) -> None:
1068 """Remove per-owner config files whose owner lilbee is gone.
1070 The swaps themselves are reaped from the state files; these leftover config
1071 files are just clutter once their writer pid is dead. A live owner's config
1072 (pid still exists) and a pid-less legacy name are left untouched; skipping on
1073 pid reuse only leaves harmless clutter, never deletes a live owner's config.
1074 """
1075 for config_path in data_dir.glob(_CONFIG_FILE_GLOB):
1076 owner = _config_owner_pid(config_path.name)
1077 if owner is not None and not psutil.pid_exists(owner):
1078 config_path.unlink(missing_ok=True)
1081def _stop_own_fleet(config_path: Path, member_ports: tuple[int, ...]) -> None:
1082 """Stop every llama-swap this lilbee owns at *config_path* and reap upstreams.
1084 Keyed on config-path identity rather than a tracked Popen or the live process
1085 tree: a warm-up/reload race can leave several llama-swap processes this lilbee
1086 started, any of which may be reparented to init, so no single handle or child
1087 scan finds them all. Every llama-swap running against our config is reaped:
1088 the build lock guarantees one builder per engine dir, so no sibling sparing
1089 applies. Each swap runs each llama-server in its own process group, so the
1090 upstreams are swept separately: captured descendants plus any llama-server
1091 still bound to one of our member ports (a respawned upstream the descendant
1092 snapshot missed), then confirmed gone.
1093 """
1094 swaps = list(_swaps_for_config(config_path))
1095 children: list[psutil.Process] = []
1096 for swap in swaps:
1097 children.extend(_live_children(swap.pid))
1098 for swap in swaps:
1099 if sys.platform == "win32":
1100 _hard_stop_proc(swap)
1101 else:
1102 _terminate_proc_group(swap)
1103 _reap_survivors(children + _find_orphan_servers(member_ports))
1106def _terminate_proc_group(proc: psutil.Process) -> None:
1107 """SIGTERM a process's group, escalating to SIGKILL on timeout."""
1108 try:
1109 pgid = os.getpgid(proc.pid)
1110 except (ProcessLookupError, OSError): # pragma: no cover - exited between checks
1111 return
1112 with contextlib.suppress(ProcessLookupError, PermissionError):
1113 os.killpg(pgid, signal.SIGTERM)
1114 try:
1115 proc.wait(timeout=_STOP_TIMEOUT_S)
1116 except psutil.TimeoutExpired:
1117 with contextlib.suppress(ProcessLookupError, PermissionError):
1118 os.killpg(pgid, _SIGKILL)
1119 _await_killed([proc])
1122def _hard_stop_proc(proc: psutil.Process) -> None:
1123 """Terminate a process, escalating to a hard kill on timeout (Windows path)."""
1124 with contextlib.suppress(psutil.NoSuchProcess):
1125 proc.terminate()
1126 try:
1127 proc.wait(timeout=_STOP_TIMEOUT_S)
1128 except psutil.TimeoutExpired:
1129 with contextlib.suppress(psutil.NoSuchProcess):
1130 proc.kill()
1133def _load_state(path: Path) -> SwapState | None:
1134 """Parse a state file into a :class:`SwapState`; ``None`` when absent/corrupt."""
1135 try:
1136 payload = json.loads(path.read_text(encoding="utf-8"))
1137 raw_pgid = payload.get(_STATE_KEY_PGID)
1138 raw_created = payload.get(_STATE_KEY_CREATED_AT)
1139 raw_ports = payload.get(_STATE_KEY_MEMBER_PORTS) or []
1140 raw_proxy = payload.get(_STATE_KEY_PROXY_PORT)
1141 raw_root = payload.get(_STATE_KEY_CONFIG_ROOT)
1142 origin = (
1143 None
1144 if raw_root is None
1145 else PlanOrigin(str(raw_root), payload.get(_STATE_KEY_FLEET_VIEW))
1146 )
1147 return SwapState(
1148 pid=int(payload[_STATE_KEY_PID]),
1149 pgid=int(raw_pgid) if raw_pgid is not None else None,
1150 created_at=float(raw_created) if raw_created is not None else None,
1151 member_ports=tuple(int(port) for port in raw_ports),
1152 proxy_port=int(raw_proxy) if raw_proxy is not None else None,
1153 launches=tuple(payload.get(_STATE_KEY_LAUNCHES) or ()),
1154 engine_pin=payload.get(_STATE_KEY_ENGINE_PIN),
1155 origin=origin,
1156 )
1157 except (OSError, ValueError, KeyError, TypeError):
1158 return None
1161def _is_live_llama_swap(state: SwapState) -> bool:
1162 """True when the recorded pid is alive and is the recorded llama-swap.
1164 A recorded create time that differs from the live process's is pid reuse,
1165 even when the recycled pid runs another instance's llama-swap; a legacy
1166 state file without one falls back to the cmdline match alone.
1167 """
1168 try:
1169 proc = psutil.Process(state.pid)
1170 cmdline = proc.cmdline()
1171 create_time = proc.create_time()
1172 except (psutil.NoSuchProcess, psutil.AccessDenied):
1173 return False
1174 if state.created_at is not None and abs(create_time - state.created_at) > (
1175 _CREATE_TIME_TOLERANCE_S
1176 ):
1177 return False
1178 binary = Path(next(iter(cmdline), "")).name
1179 return _LLAMA_SWAP_PROCESS_NAME in binary
1182def _stop_stale_swap(state: SwapState) -> None:
1183 """TERM-then-KILL a stale llama-swap's group and reap the servers it spawned.
1185 Swept as wide as ``_stop_own_fleet``: a reparented or respawned server is no
1186 longer a descendant, and every caller unlinks the record next, so the member
1187 ports are the last thing that can match it.
1188 """
1189 children = _live_children(state.pid)
1190 try:
1191 proc = psutil.Process(state.pid)
1192 except psutil.NoSuchProcess:
1193 proc = None
1194 if proc is not None:
1195 _signal_stale(state, signal.SIGTERM)
1196 try:
1197 proc.wait(timeout=_ORPHAN_STOP_TIMEOUT_S)
1198 except psutil.TimeoutExpired:
1199 _signal_stale(state, _SIGKILL)
1200 _await_killed([proc])
1201 _reap_survivors(children + _find_orphan_servers(state.member_ports))
1204def _signal_stale(state: SwapState, sig: int) -> None:
1205 """Signal the stale swap's process group, or the pid where groups don't apply."""
1206 if state.pgid is not None and sys.platform != "win32":
1207 with contextlib.suppress(ProcessLookupError, PermissionError):
1208 os.killpg(state.pgid, sig)
1209 return
1210 with contextlib.suppress(psutil.NoSuchProcess):
1211 psutil.Process(state.pid).send_signal(sig)
1214def _stop_recorded_engine(state: SwapState) -> bool:
1215 """Terminate a live llama-swap and its servers, or reap the servers a dead one
1216 orphaned (matched by recorded port, since they run in their own process groups
1217 and outlive the swap). Returns whether anything was actually alive to stop, so
1218 the off switch reports a stale record as "nothing stopped" rather than a false
1219 success.
1220 """
1221 if _is_live_llama_swap(state):
1222 _stop_stale_swap(state)
1223 return True
1224 orphans = _find_orphan_servers(state.member_ports)
1225 _reap_survivors(orphans)
1226 return bool(orphans)
1229def _find_orphan_servers(ports: tuple[int, ...]) -> list[psutil.Process]:
1230 """Live llama-server processes serving one of *ports*.
1232 Both the binary name and the ``--port`` value must match, so an unrelated
1233 process on a recycled port is never killed; a server whose parent is a
1234 live llama-swap belongs to a current run on a reused port and is spared.
1235 """
1236 if not ports:
1237 return []
1238 targets = {str(port) for port in ports}
1239 orphans: list[psutil.Process] = []
1240 for proc in _processes_named(_LLAMA_SERVER_PROCESS_NAME):
1241 try:
1242 cmdline = proc.cmdline()
1243 except (psutil.NoSuchProcess, psutil.AccessDenied):
1244 continue
1245 # Identity is the --port value plus the absence of a live swap parent;
1246 # _processes_named already gated on the executable name (comm).
1247 if _port_argument(cmdline) in targets and not _has_live_swap_parent(proc):
1248 orphans.append(proc)
1249 return orphans
1252def _state_owner_pid(name: str) -> int | None:
1253 """Owner pid embedded in a state or state-tmp filename, ``None`` when absent.
1255 Handles both the group-qualified form (``llama-swap.state.chat.123.json``)
1256 and the legacy pre-group form (``llama-swap.state.123.json``): the pid is
1257 always the last dotted segment of the stem.
1258 """
1259 stem = name.removeprefix(_STATE_TMP_PREFIX).removesuffix(_STATE_TMP_SUFFIX)
1260 stem = stem.removeprefix(_STATE_FILENAME_PREFIX).removesuffix(_STATE_FILENAME_SUFFIX)
1261 try:
1262 return int(stem.rsplit(".", 1)[-1])
1263 except ValueError:
1264 return None
1267def _state_group(name: str) -> str | None:
1268 """Group token from a group-qualified state filename, ``None`` for legacy names.
1270 ``llama-swap.state.chat.123.json`` -> ``chat``; the legacy pre-group form
1271 ``llama-swap.state.123.json`` has no group token.
1272 """
1273 stem = name.removeprefix(_STATE_FILENAME_PREFIX).removesuffix(_STATE_FILENAME_SUFFIX)
1274 head, _, _ = stem.rpartition(".") # drop the trailing pid; group is what remains
1275 return head or None
1278def _config_filename(pid: int, group: str) -> str:
1279 """This owner's config filename for *group* (``llama-swap-<group>.<pid>.json``)."""
1280 return _CONFIG_FILENAME_TEMPLATE.format(group=group, pid=pid)
1283def _config_owner_pid(name: str) -> int | None:
1284 """Owner pid embedded in a config filename, ``None`` for a legacy pid-less name."""
1285 stem = name.removeprefix("llama-swap-").removesuffix(".json")
1286 try:
1287 return int(stem.rsplit(".", 1)[-1])
1288 except ValueError:
1289 return None
1292def _has_live_swap_parent(proc: psutil.Process) -> bool:
1293 """True when *proc*'s parent is a live llama-swap (the server is not orphaned)."""
1294 try:
1295 parent = proc.parent()
1296 if parent is None:
1297 return False
1298 return _LLAMA_SWAP_PROCESS_NAME in parent.name()
1299 except (psutil.NoSuchProcess, psutil.AccessDenied):
1300 return False
1303def _port_argument(cmdline: list[str]) -> str | None:
1304 """The value following the port flag in *cmdline*, or ``None``."""
1305 for flag, value in itertools.pairwise(cmdline):
1306 if flag == PORT_FLAG:
1307 return value
1308 return None