Coverage for src/lilbee/providers/fleet/swap_manager.py: 100%

1"""Supervise the single llama-swap process that fronts every fleet role. 

2 

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""" 

6 

7from __future__ import annotations 

8 

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 functools import lru_cache 

23from pathlib import Path 

24from typing import TYPE_CHECKING, BinaryIO 

25 

26import httpx 

27import psutil 

28 

29from lilbee.providers.base import ProviderError, ProviderErrorKind 

30from lilbee.providers.fleet.binary import engine_pin, resolve_llama_swap 

31from lilbee.providers.fleet.child_guard import release_death_pipe, spawn_bound_child 

32from lilbee.providers.fleet.groups import SwapGroup 

33from lilbee.providers.fleet.launch import role_model_prefix 

34from lilbee.providers.fleet.planning import clear_ctx_downshift 

35from lilbee.providers.fleet.readback import check_launch, report_missing_log 

36from lilbee.providers.fleet.swap_config import PORT_FLAG, build_swap_config 

37from lilbee.runtime.engine_lock import clear_keep_warm 

38 

39if TYPE_CHECKING: 

40 from lilbee.providers.fleet.launch import InstanceLaunch 

41 from lilbee.providers.roles import WorkerRole 

42 

43log = logging.getLogger(__name__) 

44 

45_HOST = "127.0.0.1" 

46# One llama-swap per swap group: the group name lands in the config filename so 

47# each group's processes are identified (and stopped) by their own config path, 

48# and a placement change can restart one group without touching the others. 

49# The writer pid segment is uniqueness, not ownership: the build lock ensures 

50# one builder per engine dir, and reaping cleans dead writers' leftovers. 

51_CONFIG_FILENAME_TEMPLATE = "llama-swap-{group}.{pid}.json" 

52_CONFIG_FILE_GLOB = "llama-swap-*.json" 

53# llama-swap's own stdout/stderr (its HTTP access log) is captured to a file in a 

54# ``logs/`` dir inside the engine dir, which is the machine slot rather than any 

55# one lilbee's data root, so the log sits beside the engine it belongs to instead 

56# of beside server.log. Capturing it at all, rather than inheriting the parent's 

57# fd, is because a TUI or CLI parent owns the terminal and an inherited fd would bleed 

58# llama-swap's request log onto the screen and corrupt the render. Per-model 

59# upstream logs are unaffected (those go to llama-swap's /logs API). 

60_LOGS_SUBDIR = "logs" 

61_LOG_FILENAME_TEMPLATE = "llama-swap-{group}.log" 

62# Each writer's state file records its swap's pid/pgid so a later start can 

63# stop a dead or unhealthy engine. Health, not ownership, decides sparing. 

64_STATE_FILENAME_PREFIX = "llama-swap.state." 

65_STATE_FILENAME_SUFFIX = ".json" 

66# Also matches the legacy single shared state file ("llama-swap.state.json"). 

67_STATE_FILE_GLOB = f"{_STATE_FILENAME_PREFIX}*" 

68_STATE_KEY_PID = "pid" 

69_STATE_KEY_PGID = "pgid" 

70_STATE_KEY_CREATED_AT = "created_at" 

71_STATE_KEY_NAME = "name" 

72_STATE_KEY_MEMBER_PORTS = "member_ports" 

73_STATE_KEY_PROXY_PORT = "proxy_port" 

74_STATE_KEY_LAUNCHES = "launches" 

75_STATE_KEY_ENGINE_PIN = "engine_pin" 

76# Atomic state writes: the dot prefix keeps half-written tmp files out of the 

77# reap scan's glob. 

78_STATE_TMP_PREFIX = "." 

79_STATE_TMP_SUFFIX = ".tmp" 

80# Pid reuse guard: a live process at a recorded pid whose create time differs 

81# from the recorded one by more than this is a different process. 

82_CREATE_TIME_TOLERANCE_S = 1.0 

83_LLAMA_SWAP_PROCESS_NAME = "llama-swap" 

84_LLAMA_SERVER_PROCESS_NAME = "llama-server" 

85_CONFIG_FLAG = "-config" 

86_LISTEN_FLAG = "-listen" 

87_HEALTH_PATH = "/health" 

88_RUNNING_PATH = "/running" 

89_HTTP_TIMEOUT_S = 10.0 

90# llama-swap's own proxy answers within a second; upstream model loads have their 

91# own (longer) budget inside llama-swap, so this only covers the proxy coming up. 

92_BOOT_TIMEOUT_S = 30.0 

93_BOOT_POLL_S = 0.25 

94# Cap on the captured llama-swap output a boot-failure error carries. 

95_BOOT_LOG_TAIL_CHARS = 2000 

96# Per-group SIGTERM grace before SIGKILL on the manager shutdown/reload path. A 

97# hard kill is safe (llama-server holds no persistent state). Note this is NOT 

98# the constant the serve handoff waits on: that path goes through stop_engine -> 

99# _stop_stale_swap and spends _ORPHAN_STOP_TIMEOUT_S plus the kill/reap waits, so 

100# SERVER_LOCK_TIMEOUT budgets only a teardown whose SIGTERMs are honored. 

101_STOP_TIMEOUT_S = 2.5 

102# Grace for a llama-server that outlived llama-swap before it is force-killed. 

103_ORPHAN_STOP_TIMEOUT_S = 5.0 

104# Grace for a SIGKILLed process to exit (and release its VRAM) before the next 

105# free-memory probe runs. 

106_KILL_WAIT_TIMEOUT_S = 5.0 

107_PROBE_TIMEOUT_S = 5.0 

108# Liveness probes talk to a loopback proxy, so they get their own short budget 

109# rather than the module's 10 s general HTTP one. The ladder runs this probe for 

110# every group while holding the cross-process build lock, so one wedged port 

111# (SYN-accepted but unresponsive) would otherwise stall every other lilbee start 

112# for tens of seconds. A local proxy that cannot answer /running this fast is 

113# not usable for inference either. 

114_LIVENESS_TIMEOUT = httpx.Timeout(connect=0.5, read=2.0, write=2.0, pool=2.0) 

115 

116 

117@lru_cache(maxsize=1) 

118def _probe_client() -> httpx.Client: 

119 """One shared client for the localhost engine probes. 

120 

121 ``httpx.get`` builds a fresh ``Client`` per call, and every ``Client`` 

122 construction creates an SSL context, which loads the system CA bundle. These 

123 probes are plain HTTP to 127.0.0.1, so none of that TLS setup is ever used -- 

124 and the readiness probe runs on the task bar's timer (up to 10 Hz), which made 

125 ``ssl.create_default_context`` 23% of TUI CPU in a py-spy profile. One client 

126 builds that at most once and keeps the connection alive between polls. 

127 ``trust_env`` is off so a proxy env var cannot redirect a loopback probe. 

128 """ 

129 return httpx.Client(trust_env=False) 

130 

131 

132_PROVIDER = "llama-server" 

133# /running JSON shape: {"running": [{"model": <id>, "state": "ready", ...}, ...]}. 

134_KEY_RUNNING = "running" 

135_KEY_MODEL = "model" 

136_KEY_STATE = "state" 

137_STATE_READY = "ready" 

138 

139 

140def _platform_const(module: object, name: str, default: int) -> int: 

141 """A platform-conditional stdlib constant (absent on some OSes -> default).""" 

142 return getattr(module, name, default) 

143 

144 

145_CREATE_NEW_PROCESS_GROUP: int = _platform_const(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) 

146_SIGKILL: int = _platform_const(signal, "SIGKILL", signal.SIGTERM) 

147 

148 

149def _atomic_write(path: Path, text: str) -> None: 

150 """Write *text* to *path* via a temp file in the same dir, then rename over it. 

151 

152 A plain write truncates the destination first, so a process dying mid-write 

153 (OOM kill, SIGKILL, disk full) leaves an empty or half-written file behind. 

154 For the llama-swap config that means the next spawn hands the engine a file 

155 it cannot start from; for the state file it means a sibling's reap scan 

156 reads a torn record. 

157 

158 The temp name carries the destination's name, and both config and state 

159 filenames embed the writing process's pid, so a crash leftover can be told 

160 from a live writer's file in flight -- see ``_clean_stale_tmp_files``. 

161 """ 

162 tmp_path = path.with_name(f"{_STATE_TMP_PREFIX}{path.name}{_STATE_TMP_SUFFIX}") 

163 tmp_path.write_text(text, encoding="utf-8") 

164 os.replace(tmp_path, path) 

165 

166 

167def _state_filename(owner_pid: int, group: str) -> str: 

168 """The per-owner, per-group state filename for the lilbee process *owner_pid*.""" 

169 return f"{_STATE_FILENAME_PREFIX}{group}.{owner_pid}{_STATE_FILENAME_SUFFIX}" 

170 

171 

172@dataclass(frozen=True) 

173class SwapState: 

174 """A running llama-swap's recorded identity and serving contract. 

175 

176 Read back from the engine dir's state file, so it describes engines this 

177 process did not start. The currency the bind/build ladder is written in: 

178 swap_manager records it, provider reads it to decide what a slot is 

179 serving, and contract matches it against what this process wants. 

180 """ 

181 

182 pid: int 

183 pgid: int | None 

184 created_at: float | None = None 

185 member_ports: tuple[int, ...] = () 

186 proxy_port: int | None = None 

187 launches: tuple[dict, ...] = () 

188 engine_pin: str | None = None 

189 

190 

191class SwapManager: 

192 """Owns one llama-swap process fronting one role group's servers. 

193 

194 The provider runs one manager per role, so restarting a group (a placement 

195 or model change) never touches another group's loaded servers. 

196 """ 

197 

198 def __init__(self, data_dir: Path, group: SwapGroup) -> None: 

199 self._data_dir = data_dir 

200 self._group = group 

201 self._config_path = data_dir / _config_filename(os.getpid(), group.value) 

202 self._log_path = data_dir / _LOGS_SUBDIR / _LOG_FILENAME_TEMPLATE.format(group=group.value) 

203 # Instances whose engine report has already been compared to the estimate, 

204 # so the check runs once per start rather than on every readiness poll. 

205 self._estimate_checked: set[str] = set() 

206 self._launch_by_model: dict[str, InstanceLaunch] = {} 

207 self._state_path = data_dir / _state_filename(os.getpid(), group.value) 

208 self._proc: subprocess.Popen[bytes] | None = None 

209 self._log_file: BinaryIO | None = None 

210 # Where this boot's output starts in the append-mode log file. 

211 self._log_offset = 0 

212 self._port: int | None = None 

213 self._member_ports: list[int] = [] 

214 # The serving contract (per-role model/ctx/slots) persisted in every 

215 # state write, so a guest lilbee can bind to this live fleet. 

216 self._launches_payload: list[dict] = [] 

217 # True when this manager uses an engine another process built: it then 

218 # never writes state, never reaps, and never signals engine processes. 

219 self._bound = False 

220 

221 def start( 

222 self, launches: list[InstanceLaunch], *, ttl_seconds: int = 0, bind_lifetime: bool = True 

223 ) -> None: 

224 """Write the config and spawn llama-swap, waiting for its proxy to answer. 

225 

226 The proxy and every member get a freshly allocated free port, which is 

227 why llama-swap's own startPort is not used: that assigns a fixed 

228 sequential range at config load, so it would collide with a previous 

229 instance's server still shutting down (the new llama-server then fails 

230 its bind and llama-swap reports it only as "exited prematurely"). 

231 

232 This narrows that collision rather than removing a race. The ports are 

233 picked by binding and closing ephemeral sockets, while llama-swap starts 

234 each upstream lazily on its first request, so a member port can sit 

235 unbound for as long as it takes that request to arrive and anything else 

236 on the box may take it in between. Nothing in llama-swap offers a 

237 spawn-time probe to close that window; warming the roles up front 

238 shortens it for the roles that are warmed. 

239 

240 ``bind_lifetime`` binds the engine to this process so a crash cannot orphan 

241 it; it is False for a keep-warm fleet that is meant to outlive lilbee. 

242 """ 

243 # Idempotent safety net; the provider reaps before planning so the GPU 

244 # probe already saw the real free memory. 

245 self.reap_stale() 

246 # Singleton guard: one llama-swap per data_dir for this lilbee. Reap any 

247 # llama-swap we already started against this config (a leaked duplicate 

248 # from a prior race/reload) before spawning, so they cannot accumulate 

249 # and double-book a GPU. 

250 _stop_own_fleet(self._config_path, tuple(self._member_ports)) 

251 ports = _pick_free_ports(1 + len(launches)) 

252 member_ports = dict(zip([launch.model_id for launch in launches], ports[1:], strict=True)) 

253 self._member_ports = sorted(member_ports.values()) 

254 self._launches_payload = [launch.to_state() for launch in launches] 

255 self._launch_by_model = {launch.model_id: launch for launch in launches} 

256 self._estimate_checked.clear() 

257 self._config_path.parent.mkdir(parents=True, exist_ok=True) 

258 self._log_path.parent.mkdir(parents=True, exist_ok=True) 

259 _atomic_write( 

260 self._config_path, 

261 build_swap_config( 

262 launches, 

263 member_ports, 

264 swap=self._group.swaps, 

265 ttl_seconds=ttl_seconds, 

266 engine_log_dir=self._log_path.parent, 

267 ), 

268 ) 

269 self._port = ports[0] 

270 # Capture llama-swap's stdout/stderr to a file so its access log never 

271 # reaches an inherited terminal (a TUI/CLI parent) and garbles the screen. 

272 self._close_log() 

273 self._log_path.parent.mkdir(parents=True, exist_ok=True) 

274 self._log_file = self._log_path.open("ab") 

275 self._log_offset = self._log_path.stat().st_size 

276 self._proc = spawn_bound_child( 

277 [ 

278 str(resolve_llama_swap()), 

279 _CONFIG_FLAG, 

280 str(self._config_path), 

281 _LISTEN_FLAG, 

282 f"{_HOST}:{self._port}", 

283 ], 

284 bind_lifetime=bind_lifetime, 

285 stdout=self._log_file, 

286 stderr=subprocess.STDOUT, 

287 start_new_session=True, 

288 creationflags=_CREATE_NEW_PROCESS_GROUP, 

289 ) 

290 self._write_state() 

291 self._await_health() 

292 

293 def reap_stale(self) -> None: 

294 """Kill every dead or unhealthy recorded engine; see :func:`reap_stale`.""" 

295 reap_stale(self._data_dir) 

296 

297 def _process_identity(self) -> tuple[int, int | None, float | None] | None: 

298 """(pid, pgid, create time) of the swap this manager runs, or None.""" 

299 if self._proc is not None: 

300 pid = self._proc.pid 

301 pgid: int | None = None 

302 if sys.platform != "win32": 

303 with contextlib.suppress(ProcessLookupError): 

304 pgid = os.getpgid(pid) 

305 created_at: float | None = None 

306 with contextlib.suppress(psutil.NoSuchProcess, psutil.AccessDenied): 

307 created_at = psutil.Process(pid).create_time() 

308 return pid, pgid, created_at 

309 return None 

310 

311 def _write_state(self) -> None: 

312 """Record the swap's pid/pgid/create time, member ports, and our identity. 

313 

314 The write is atomic (tmp file then ``os.replace``) so a sibling's reap 

315 scan can never read a torn file and mistake this live record for junk. 

316 """ 

317 identity = self._process_identity() 

318 if identity is None: 

319 return 

320 swap_pid, pgid, created_at = identity 

321 state = { 

322 _STATE_KEY_PID: swap_pid, 

323 _STATE_KEY_PGID: pgid, 

324 _STATE_KEY_CREATED_AT: created_at, 

325 _STATE_KEY_NAME: _LLAMA_SWAP_PROCESS_NAME, 

326 _STATE_KEY_MEMBER_PORTS: self._member_ports, 

327 _STATE_KEY_PROXY_PORT: self._port, 

328 _STATE_KEY_LAUNCHES: self._launches_payload, 

329 _STATE_KEY_ENGINE_PIN: engine_pin(), 

330 } 

331 _atomic_write(self._state_path, json.dumps(state)) 

332 

333 def endpoint(self) -> str: 

334 """Base URL of the llama-swap OpenAI-compatible proxy.""" 

335 if self._port is None: 

336 raise ProviderError( 

337 "The local model engine is not running.", 

338 provider=_PROVIDER, 

339 kind=ProviderErrorKind.SERVER, 

340 ) 

341 return f"http://{_HOST}:{self._port}" 

342 

343 def role_ready(self, role: WorkerRole) -> bool: 

344 """Whether at least one of *role*'s replica servers is loaded and ready.""" 

345 prefix = role_model_prefix(role) 

346 ready = self._ready_models() 

347 self._check_estimates(ready) 

348 return any(model.startswith(prefix) for model in ready) 

349 

350 def _check_estimates(self, ready: set[str]) -> None: 

351 """Compare each newly-ready engine's own report against what it was planned for. 

352 

353 The plan is otherwise open-loop, and a wrong estimate only ever surfaces 

354 as a failed request much later. Once per instance per start: readiness is 

355 polled, and the answer does not change once the engine has loaded. 

356 """ 

357 for model_id in ready - self._estimate_checked: 

358 launch = self._launch_by_model.get(model_id) 

359 self._estimate_checked.add(model_id) 

360 if launch is None: 

361 continue 

362 # Ready means this role's context loaded, so any reduction taken to 

363 # get here has done its job and must not follow the role into the 

364 # next plan, a freed machine, or a model the user switched to. 

365 clear_ctx_downshift(launch.role) 

366 # The engine is ready, so a missing log is not "too early" any more. 

367 if report_missing_log(self._log_path.parent, model_id, launch.role): 

368 continue 

369 if launch.est_vram_bytes <= 0: 

370 continue 

371 check_launch( 

372 self._log_path.parent, 

373 model_id, 

374 launch.role, 

375 launch.model, 

376 launch.est_vram_bytes, 

377 launch.est_vram_by_device, 

378 launch.est_unreported_bytes, 

379 ) 

380 

381 def is_live(self) -> bool: 

382 """Whether the swap process is up and its proxy answers ``/running``.""" 

383 if self._proc is None or self._proc.poll() is not None: 

384 return False 

385 if self._port is None: 

386 return False 

387 return self._proxy_answers() 

388 

389 @property 

390 def running(self) -> bool: 

391 """Whether this manager currently has a spawned llama-swap process.""" 

392 return self._proc is not None 

393 

394 @property 

395 def bound(self) -> bool: 

396 """Whether this manager rides an engine built by another process.""" 

397 return self._bound 

398 

399 def bind(self, state: SwapState) -> bool: 

400 """Use a running engine's proxy without taking any ownership of it. 

401 

402 The engine's own state record stays untouched: the binder writes 

403 nothing, and shutdown() merely drops the binding. 

404 """ 

405 if state.proxy_port is None: 

406 return False 

407 self._port = state.proxy_port 

408 self._member_ports = list(state.member_ports) 

409 if not self._proxy_answers(): 

410 self._port = None 

411 self._member_ports = [] 

412 return False 

413 self._launches_payload = [dict(launch) for launch in state.launches] 

414 self._bound = True 

415 return True 

416 

417 def _proxy_answers(self) -> bool: 

418 """Whether the bound proxy port serves llama-swap's running endpoint. 

419 

420 Shares state_is_healthy's identity check via _running_endpoint_answers, so 

421 bind and reap agree on what "answering" means by construction rather than by 

422 two hand-kept-identical copies. 

423 """ 

424 return _running_endpoint_answers(self.endpoint()) 

425 

426 def shutdown(self) -> None: 

427 """Stop every llama-swap this lilbee owns at our config and reap servers. 

428 

429 Authoritative teardown keyed on config-path identity, not the single 

430 tracked ``Popen``: a warm-up/reset race or a reload can leave several 

431 llama-swap processes this lilbee spawned, any of them reparented to init 

432 (still holding the engine binary open) -- trusting one handle would leak 

433 them. Every llama-swap running against our config is reaped. Unlinks only 

434 this owner's state file; another instance's record stays. 

435 """ 

436 if self._bound: 

437 # Not ours to stop: drop the binding and leave the engine serving. 

438 self._bound = False 

439 self._port = None 

440 self._member_ports = [] 

441 self._launches_payload = [] 

442 return 

443 _stop_own_fleet(self._config_path, tuple(self._member_ports)) 

444 # Nothing is coming back to bind these, so the picker can offer them again. 

445 release_reserved_ports([*self._member_ports, *([self._port] if self._port else [])]) 

446 self._state_path.unlink(missing_ok=True) 

447 if self._proc is not None: 

448 # Free this engine's death pipe so its watcher exits now, not at our death. 

449 release_death_pipe(self._proc.pid) 

450 self._proc = None 

451 self._port = None 

452 self._close_log() 

453 

454 def _close_log(self) -> None: 

455 """Close the captured llama-swap log handle, if one is open.""" 

456 if self._log_file is not None: 

457 with contextlib.suppress(OSError): 

458 self._log_file.close() 

459 self._log_file = None 

460 

461 def _await_health(self) -> None: 

462 """Poll the proxy's /health until it answers, or fail with a clear error.""" 

463 url = f"{self.endpoint()}{_HEALTH_PATH}" 

464 deadline = time.monotonic() + _BOOT_TIMEOUT_S 

465 while time.monotonic() < deadline: 

466 if self._proc is not None and self._proc.poll() is not None: 

467 self._fail("The local model engine exited before it was ready.") 

468 with contextlib.suppress(httpx.HTTPError): 

469 if _probe_client().get(url, timeout=_PROBE_TIMEOUT_S).status_code == httpx.codes.OK: 

470 return 

471 time.sleep(_BOOT_POLL_S) 

472 self._fail("The local model engine did not start in time.") 

473 

474 def _ready_models(self) -> set[str]: 

475 """Model ids whose upstream is loaded and ready, per llama-swap's /running. 

476 

477 A read-only probe: a concurrent shutdown can clear ``_port`` between the 

478 caller's check and ``endpoint()``, raising ProviderError, so that is 

479 suppressed too and the probe reports "nothing ready" rather than throwing. 

480 """ 

481 with contextlib.suppress(httpx.HTTPError, ValueError, KeyError, TypeError, ProviderError): 

482 payload = ( 

483 _probe_client() 

484 .get(f"{self.endpoint()}{_RUNNING_PATH}", timeout=_PROBE_TIMEOUT_S) 

485 .json() 

486 ) 

487 return { 

488 entry[_KEY_MODEL] 

489 for entry in payload[_KEY_RUNNING] 

490 if entry.get(_KEY_STATE) == _STATE_READY 

491 } 

492 return set() 

493 

494 def _boot_log_tail(self) -> str: 

495 """The current boot's captured llama-swap output, capped for an error message.""" 

496 try: 

497 with self._log_path.open("rb") as handle: 

498 handle.seek(self._log_offset) 

499 data = handle.read() 

500 except OSError: 

501 return "" 

502 return data.decode(errors="replace").strip()[-_BOOT_LOG_TAIL_CHARS:] 

503 

504 def _fail(self, message: str) -> None: 

505 """Tear down and raise a user-facing engine-start error carrying the boot log.""" 

506 self.shutdown() 

507 tail = self._boot_log_tail() 

508 if tail: 

509 message = f"{message} Engine log ({self._log_path}):\n{tail}" 

510 raise ProviderError(message, provider=_PROVIDER, kind=ProviderErrorKind.SERVER) 

511 

512 

513# Linux publishes the range here; every other platform is asked via sysctl. 

514_PROC_PORT_RANGE = Path("/proc/sys/net/ipv4/ip_local_port_range") 

515 

516 

517def _port_range_from(path: Path) -> tuple[int, int] | None: 

518 """The two integers in *path*, or ``None`` when it is absent or unreadable.""" 

519 try: 

520 low, high = path.read_text(encoding="utf-8").split()[:2] 

521 return int(low), int(high) 

522 except (OSError, ValueError): 

523 return None 

524 

525 

526def _ephemeral_range() -> tuple[int, int] | None: 

527 """The port range the kernel hands out for unbound sockets, if it says. 

528 

529 ``None`` when neither source answers, which is the signal to fall back to 

530 letting the OS choose. 

531 """ 

532 from_proc = _port_range_from(_PROC_PORT_RANGE) 

533 if from_proc is not None: 

534 return from_proc 

535 try: # macOS and the BSDs, which have no procfs entry for this 

536 out = subprocess.run( 

537 ["/usr/sbin/sysctl", "-n", "net.inet.ip.portrange.first", "net.inet.ip.portrange.last"], 

538 capture_output=True, 

539 text=True, 

540 encoding="utf-8", 

541 errors="replace", 

542 timeout=5, 

543 check=False, 

544 ) 

545 low, high = out.stdout.split()[:2] 

546 return int(low), int(high) 

547 except (OSError, ValueError, subprocess.SubprocessError): 

548 return None 

549 

550 

551# Where lilbee looks for engine ports when the kernel's ephemeral range is known. 

552# Above the registered-service crowd, below every default ephemeral range. 

553_PORT_SEARCH_FLOOR = 20000 

554_PORT_WINDOW_SPAN = 8192 

555# Block width. Each process searches one block, so concurrent lilbees hold 

556# disjoint ranges. A fleet takes one proxy port plus one per member, and embed 

557# and vision replicate per GPU, so 64 covers a 30-GPU host; a wider fleet spills 

558# into the next block. 

559_PORT_BLOCK = 64 

560 

561# Ports handed to a child that has not bound them yet. llama-swap binds a member 

562# port only on that member's first request, so the probe socket is long closed 

563# by then and the port looks free to every later probe. Without this the picker 

564# hands the next group exactly what it gave the last one, every time. 

565_reserved_ports: set[int] = set() 

566_reserved_lock = threading.Lock() 

567 

568 

569def release_reserved_ports(ports: Iterable[int]) -> None: 

570 """Give *ports* back to the picker, once nothing is expected to bind them.""" 

571 with _reserved_lock: 

572 _reserved_ports.difference_update(ports) 

573 

574 

575def _window_span(ceiling: tuple[int, int]) -> int: 

576 """How many ports below the ephemeral floor this host leaves to search.""" 

577 return max(1, min(ceiling[0], _PORT_SEARCH_FLOOR + _PORT_WINDOW_SPAN) - _PORT_SEARCH_FLOOR) 

578 

579 

580def _search_start(ceiling: tuple[int, int]) -> int: 

581 """First port of the block this process owns. 

582 

583 The pid selects a whole block, not an offset: reservation is per-process, and 

584 a fleet takes its ports contiguously, so pid-offset starts one apart overlap 

585 on all but one port. 

586 """ 

587 blocks = max(1, _window_span(ceiling) // _PORT_BLOCK) 

588 return _PORT_SEARCH_FLOOR + (os.getpid() % blocks) * _PORT_BLOCK 

589 

590 

591def _pick_free_ports(count: int) -> list[int]: 

592 """Bind *count* free localhost ports at once and return them. 

593 

594 All sockets stay open until every port is claimed so the OS cannot hand the 

595 same port out twice within one allocation. 

596 

597 Picked from below the kernel's ephemeral range rather than inside it. The 

598 gap between lilbee picking a port and llama-server binding it spans the whole 

599 lazy-spawn wait, and a port inside the ephemeral range can be handed to any 

600 passing outbound connection during that gap; one below it cannot be handed to 

601 anybody, so the only way to lose it is another server binding that exact port 

602 on purpose. Falls back to letting the OS choose when the range is unknown. 

603 """ 

604 ceiling = _ephemeral_range() 

605 sockets = [socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(count)] 

606 try: 

607 for sock in sockets: 

608 _bind_below_ephemeral(sock, ceiling) 

609 return [int(sock.getsockname()[1]) for sock in sockets] 

610 finally: 

611 for sock in sockets: 

612 sock.close() 

613 

614 

615def _bind_below_ephemeral(sock: socket.socket, ceiling: tuple[int, int] | None) -> None: 

616 """Bind *sock* to a free, unreserved port under the ephemeral floor. 

617 

618 Falls back to letting the OS choose when the range is unknown or the window 

619 is used up, which keeps a fleet start working at the cost of returning to the 

620 ephemeral range for those ports. 

621 """ 

622 if ceiling is not None and ceiling[0] > _PORT_SEARCH_FLOOR: 

623 span = _window_span(ceiling) 

624 start = _search_start(ceiling) 

625 for offset in range(span): 

626 port = _PORT_SEARCH_FLOOR + (start - _PORT_SEARCH_FLOOR + offset) % span 

627 with _reserved_lock: 

628 if port in _reserved_ports: 

629 continue 

630 try: 

631 sock.bind((_HOST, port)) 

632 except OSError: 

633 continue 

634 _reserved_ports.add(port) 

635 return 

636 sock.bind((_HOST, 0)) 

637 

638 

639def _live_children(pid: int) -> list[psutil.Process]: 

640 """The process's current descendants, or none when it already exited.""" 

641 try: 

642 children: list[psutil.Process] = psutil.Process(pid).children(recursive=True) 

643 except psutil.NoSuchProcess: 

644 return [] 

645 return children 

646 

647 

648def _reap_survivors(children: list[psutil.Process]) -> None: 

649 """Terminate then kill any captured child that is still running.""" 

650 survivors = [child for child in children if child.is_running()] 

651 for child in survivors: 

652 with contextlib.suppress(psutil.NoSuchProcess): 

653 child.terminate() 

654 _, alive = psutil.wait_procs(survivors, timeout=_ORPHAN_STOP_TIMEOUT_S) 

655 for child in alive: 

656 with contextlib.suppress(psutil.NoSuchProcess): 

657 child.kill() 

658 _await_killed(alive) 

659 

660 

661def _await_killed(procs: list[psutil.Process]) -> None: 

662 """Wait for SIGKILLed processes to exit so their VRAM is free before any probe.""" 

663 if not procs: 

664 return 

665 _, alive = psutil.wait_procs(procs, timeout=_KILL_WAIT_TIMEOUT_S) 

666 for proc in alive: 

667 log.warning("Process %s survived SIGKILL; its VRAM may still be held.", proc.pid) 

668 

669 

670def _processes_named(needle: str) -> Iterator[psutil.Process]: 

671 """Live processes whose executable name contains *needle*. 

672 

673 ``name()`` is a cheap field (comm/proc_name); ``cmdline()`` reads the full 

674 argument vector and on macOS blocks on entitlement-protected binaries. So the 

675 name is the pre-filter and callers pay for ``cmdline()`` only on a match, 

676 which keeps a full-process-table scan from stalling on an unrelated process. 

677 """ 

678 for proc in psutil.process_iter(["name"]): 

679 # process_iter already skips processes that vanish mid-scan and, per its 

680 # ad_value contract, leaves ``name`` as None where it could not be read. 

681 name = proc.info["name"] or "" 

682 if needle in name: 

683 yield proc 

684 

685 

686def _swaps_for_config(config_path: Path) -> list[psutil.Process]: 

687 """Every live llama-swap (any owner) running against *config_path*. 

688 

689 Identity is the ``-config <path>`` argument, which every llama-swap this 

690 lilbee starts carries and which survives reparenting to init -- so this finds 

691 a leaked duplicate or a swap reparented away from us, neither of which a 

692 tracked Popen handle nor a ``children()`` scan would catch. 

693 """ 

694 target = str(config_path) 

695 swaps: list[psutil.Process] = [] 

696 for proc in _processes_named(_LLAMA_SWAP_PROCESS_NAME): 

697 try: 

698 cmdline = proc.cmdline() 

699 except ( 

700 psutil.NoSuchProcess, 

701 psutil.AccessDenied, 

702 psutil.ZombieProcess, 

703 OSError, 

704 SystemError, 

705 ): 

706 # OSError/SystemError: macOS psutil mishandles entitlement-protected 

707 # binaries (sysctl KERN_PROCARGS2), leaking a raw PermissionError or a 

708 # C-extension SystemError instead of an AccessDenied. 

709 continue 

710 # Identity is the -config path; _processes_named already gated on comm. 

711 if target in cmdline: 

712 swaps.append(proc) 

713 return swaps 

714 

715 

716def find_live_state(data_dir: Path, group: SwapGroup) -> SwapState | None: 

717 """The newest recorded state for *group* at *data_dir* (no liveness check). 

718 

719 A record's presence does not prove the engine is up; callers that need that 

720 probe it with ``state_is_healthy``. The name reflects that a record is written 

721 only for a running engine, not that this function verifies it. 

722 """ 

723 best: SwapState | None = None 

724 for state_path in sorted(data_dir.glob(_STATE_FILE_GLOB)): 

725 if f".{group.value}." not in f".{state_path.name}": 

726 continue 

727 state = _load_state(state_path) 

728 if state is None: 

729 continue 

730 if best is None or (state.created_at or 0) > (best.created_at or 0): 

731 best = state 

732 return best 

733 

734 

735def _running_endpoint_answers(base_url: str) -> bool: 

736 """Whether *base_url* serves llama-swap's ``/running`` endpoint (identity, not 

737 just liveness). 

738 

739 Proxy ports are ephemeral: after an engine dies, any unrelated local service 

740 that later binds the recorded port and returns a 2xx/3xx to an unknown path 

741 would pass a bare status check, so a dead record would look healthy forever and 

742 inference clients would bind to a non-engine endpoint. Requiring the ``running`` 

743 JSON payload shape that only llama-swap produces makes the probe identity-checked. 

744 Total: any transport error or non-conforming body reads as "not our engine". 

745 """ 

746 try: 

747 resp = _probe_client().get(f"{base_url}{_RUNNING_PATH}", timeout=_LIVENESS_TIMEOUT) 

748 except (OSError, httpx.HTTPError): 

749 return False 

750 if resp.status_code >= httpx.codes.BAD_REQUEST: 

751 return False 

752 try: 

753 return isinstance(resp.json().get(_KEY_RUNNING), list) 

754 except (ValueError, AttributeError): 

755 return False 

756 

757 

758def state_is_healthy(state: SwapState) -> bool: 

759 """Whether the engine behind *state* answers on its recorded proxy port.""" 

760 if state.proxy_port is None: 

761 return False 

762 return _running_endpoint_answers(f"http://{_HOST}:{state.proxy_port}") 

763 

764 

765def engine_record_exists(data_dir: Path) -> bool: 

766 """Whether any engine state file is present, without probing proxy health. 

767 

768 A filesystem fact, unlike a proxy HTTP probe: it is true for an engine that 

769 is live but momentarily unprobeable (fd exhaustion, host thrash), so the 

770 ladder can clear a recorded engine before building rather than double-build 

771 beside one an HTTP probe failed to see. 

772 """ 

773 return any(data_dir.glob(_STATE_FILE_GLOB)) 

774 

775 

776def stop_engine(data_dir: Path) -> list[str]: 

777 """Stop every engine the dir's state files record, regardless of liveness. 

778 

779 The unconditional off switch behind ``lilbee engine stop`` and the 

780 last-user-out path: each recorded swap is terminated through its state 

781 record (never a Popen handle, so it works on engines this process did 

782 not build) and its file removed. A record whose llama-swap is already dead 

783 still has its llama-servers (each in its own process group) reaped by 

784 recorded port, exactly as reap_stale does -- otherwise the off switch would 

785 leave those orphans holding VRAM and delete the ports needed to find them. 

786 Stale config files for dead owners are cleaned too, and the persistence 

787 opt-in is dropped with the engine it described, so the dir is left as 

788 clean as a reap leaves it. Unparseable files are left alone, as in 

789 reap_stale: they may be a sibling's in-flight write. Returns the group tokens 

790 whose engine was actually alive, so a caller reports only real stops. 

791 """ 

792 _clean_stale_configs(data_dir) 

793 # The persistence opt-in describes the engine instance being stopped, so it 

794 # dies with it. Cleared here rather than at each call site so no stop path 

795 # can leave a mark that makes the next engine sticky-warm. 

796 clear_keep_warm(data_dir) 

797 stopped: list[str] = [] 

798 for state_path in sorted(data_dir.glob(_STATE_FILE_GLOB)): 

799 state = _load_state(state_path) 

800 if state is None: 

801 continue 

802 if _stop_recorded_engine(state): 

803 group = _state_group(state_path.name) 

804 if group is not None: 

805 stopped.append(group) 

806 state_path.unlink(missing_ok=True) 

807 return stopped 

808 

809 

810def reap_stale(data_dir: Path) -> None: 

811 """Kill every dead or unhealthy recorded engine at *data_dir*. 

812 

813 An OOM-killed lilbee leaves llama-swap (and its servers) holding VRAM, 

814 so planning would otherwise see artificially reduced free memory; the 

815 ladder calls this before its GPU probe. Every state file is scanned 

816 (all groups, including legacy names): an engine that is alive AND 

817 answering on its proxy is spared regardless of who started it (a 

818 reload's own healthy groups, or a bindable engine the ladder skipped); 

819 everything else is stopped through its record and its file removed. An 

820 unparseable file is skipped, never deleted: it may be a sibling's 

821 in-flight write. When the swap itself is dead, its servers (each in 

822 its own process group) may still be alive holding VRAM; they are 

823 matched by name plus recorded member port and stopped before the file 

824 is removed. 

825 

826 Module-level (not a method) because it must run before planning decides 

827 which role groups exist, when no per-group manager has been built yet. 

828 """ 

829 _clean_stale_tmp_files(data_dir) 

830 _clean_stale_configs(data_dir) 

831 for state_path in sorted(data_dir.glob(_STATE_FILE_GLOB)): 

832 state = _load_state(state_path) 

833 if state is None: 

834 continue 

835 if state_is_healthy(state): 

836 # An answering engine is in use (bind accepts on exactly this 

837 # test); reaping must never disagree with binding. 

838 continue 

839 _stop_recorded_engine(state) 

840 state_path.unlink(missing_ok=True) 

841 

842 

843def _clean_stale_tmp_files(data_dir: Path) -> None: 

844 """Remove crash-leftover state and config tmp files whose writer is dead.""" 

845 tmp_glob = f"{_STATE_TMP_PREFIX}*{_STATE_TMP_SUFFIX}" 

846 for tmp_path in data_dir.glob(tmp_glob): 

847 writer_pid = _state_owner_pid(tmp_path.name) 

848 if writer_pid is not None and not psutil.pid_exists(writer_pid): 

849 tmp_path.unlink(missing_ok=True) 

850 

851 

852def _clean_stale_configs(data_dir: Path) -> None: 

853 """Remove per-owner config files whose owner lilbee is gone. 

854 

855 The swaps themselves are reaped from the state files; these leftover config 

856 files are just clutter once their writer pid is dead. A live owner's config 

857 (pid still exists) and a pid-less legacy name are left untouched; skipping on 

858 pid reuse only leaves harmless clutter, never deletes a live owner's config. 

859 """ 

860 for config_path in data_dir.glob(_CONFIG_FILE_GLOB): 

861 owner = _config_owner_pid(config_path.name) 

862 if owner is not None and not psutil.pid_exists(owner): 

863 config_path.unlink(missing_ok=True) 

864 

865 

866def _stop_own_fleet(config_path: Path, member_ports: tuple[int, ...]) -> None: