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, log_engine_launch 

35from lilbee.providers.fleet.readback import ( 

36 MEMORY_FLAG, 

37 check_launch, 

38 check_memory_report, 

39 report_missing_log, 

40) 

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

42from lilbee.runtime.engine_lock import clear_keep_warm 

43 

44if TYPE_CHECKING: 

45 from lilbee.providers.fleet.launch import InstanceLaunch 

46 from lilbee.providers.roles import WorkerRole 

47 

48log = logging.getLogger(__name__) 

49 

50_HOST = "127.0.0.1" 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

65_LOGS_SUBDIR = "logs" 

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

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

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

69_STATE_FILENAME_PREFIX = "llama-swap.state." 

70_STATE_FILENAME_SUFFIX = ".json" 

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

72_STATE_FILE_GLOB = f"{_STATE_FILENAME_PREFIX}*" 

73_STATE_KEY_PID = "pid" 

74_STATE_KEY_PGID = "pgid" 

75_STATE_KEY_CREATED_AT = "created_at" 

76_STATE_KEY_NAME = "name" 

77_STATE_KEY_MEMBER_PORTS = "member_ports" 

78_STATE_KEY_PROXY_PORT = "proxy_port" 

79_STATE_KEY_LAUNCHES = "launches" 

80_STATE_KEY_ENGINE_PIN = "engine_pin" 

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

82# reap scan's glob. 

83_STATE_TMP_PREFIX = "." 

84_STATE_TMP_SUFFIX = ".tmp" 

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

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

87_CREATE_TIME_TOLERANCE_S = 1.0 

88_LLAMA_SWAP_PROCESS_NAME = "llama-swap" 

89_LLAMA_SERVER_PROCESS_NAME = "llama-server" 

90_CONFIG_FLAG = "-config" 

91_LISTEN_FLAG = "-listen" 

92_HEALTH_PATH = "/health" 

93_RUNNING_PATH = "/running" 

94_HTTP_TIMEOUT_S = 10.0 

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

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

97_BOOT_TIMEOUT_S = 30.0 

98_BOOT_POLL_S = 0.25 

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

100_BOOT_LOG_TAIL_CHARS = 2000 

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

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

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

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

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

106_STOP_TIMEOUT_S = 2.5 

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

108_ORPHAN_STOP_TIMEOUT_S = 5.0 

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

110# free-memory probe runs. 

111_KILL_WAIT_TIMEOUT_S = 5.0 

112_PROBE_TIMEOUT_S = 5.0 

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

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

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

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

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

118# not usable for inference either. 

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

120 

121 

122@lru_cache(maxsize=1) 

123def _probe_client() -> httpx.Client: 

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

125 

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

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

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

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

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

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

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

133 """ 

134 return httpx.Client(trust_env=False) 

135 

136 

137_PROVIDER = "llama-server" 

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

139_KEY_RUNNING = "running" 

140_KEY_MODEL = "model" 

141_KEY_STATE = "state" 

142_STATE_READY = "ready" 

143 

144 

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

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

147 return getattr(module, name, default) 

148 

149 

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

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

152 

153 

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

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

156 

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

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

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

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

161 reads a torn record. 

162 

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

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

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

166 """ 

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

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

169 os.replace(tmp_path, path) 

170 

171 

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

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

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

175 

176 

177@dataclass(frozen=True) 

178class SwapState: 

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

180 

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

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

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

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

185 """ 

186 

187 pid: int 

188 pgid: int | None 

189 created_at: float | None = None 

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

191 proxy_port: int | None = None 

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

193 engine_pin: str | None = None 

194 

195 

196class SwapManager: 

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

198 

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

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

201 """ 

202 

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

204 self._data_dir = data_dir 

205 self._group = group 

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

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

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

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

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

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

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

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

214 self._log_file: BinaryIO | None = None 

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

216 self._log_offset = 0 

217 self._port: int | None = None 

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

219 # Which member port serves which model id, for the direct GET /memory 

220 # readback. Only launches this manager spawned are in here: a bound 

221 # manager's state record carries the ports but not the mapping, and it 

222 # has no launches to check either. 

223 self._member_port_by_model: dict[str, int] = {} 

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

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

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

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

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

229 self._bound = False 

230 

231 def start( 

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

233 ) -> None: 

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

235 

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

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

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

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

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

241 

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

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

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

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

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

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

248 shortens it for the roles that are warmed. 

249 

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

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

252 """ 

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

254 # probe already saw the real free memory. 

255 self.reap_stale() 

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

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

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

259 # and double-book a GPU. 

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

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

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

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

264 self._member_port_by_model = dict(member_ports) 

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

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

267 self._estimate_checked.clear() 

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

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

270 _atomic_write( 

271 self._config_path, 

272 build_swap_config( 

273 launches, 

274 member_ports, 

275 swap=self._group.swaps, 

276 ttl_seconds=ttl_seconds, 

277 engine_log_dir=self._log_path.parent, 

278 ), 

279 ) 

280 self._port = ports[0] 

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

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

283 self._close_log() 

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

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

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

287 self._proc = spawn_bound_child( 

288 [ 

289 str(resolve_llama_swap()), 

290 _CONFIG_FLAG, 

291 str(self._config_path), 

292 _LISTEN_FLAG, 

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

294 ], 

295 bind_lifetime=bind_lifetime, 

296 stdout=self._log_file, 

297 stderr=subprocess.STDOUT, 

298 start_new_session=True, 

299 creationflags=_CREATE_NEW_PROCESS_GROUP, 

300 ) 

301 self._write_state() 

302 self._await_health() 

303 for launch in launches: 

304 log_engine_launch(launch) 

305 

306 def reap_stale(self) -> None: 

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

308 reap_stale(self._data_dir) 

309 

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

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

312 if self._proc is not None: 

313 pid = self._proc.pid 

314 pgid: int | None = None 

315 if sys.platform != "win32": 

316 with contextlib.suppress(ProcessLookupError): 

317 pgid = os.getpgid(pid) 

318 created_at: float | None = None 

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

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

321 return pid, pgid, created_at 

322 return None 

323 

324 def _write_state(self) -> None: 

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

326 

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

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

329 """ 

330 identity = self._process_identity() 

331 if identity is None: 

332 return 

333 swap_pid, pgid, created_at = identity 

334 state = { 

335 _STATE_KEY_PID: swap_pid, 

336 _STATE_KEY_PGID: pgid, 

337 _STATE_KEY_CREATED_AT: created_at, 

338 _STATE_KEY_NAME: _LLAMA_SWAP_PROCESS_NAME, 

339 _STATE_KEY_MEMBER_PORTS: self._member_ports, 

340 _STATE_KEY_PROXY_PORT: self._port, 

341 _STATE_KEY_LAUNCHES: self._launches_payload, 

342 _STATE_KEY_ENGINE_PIN: engine_pin(), 

343 } 

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

345 

346 def endpoint(self) -> str: 

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

348 if self._port is None: 

349 raise ProviderError( 

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

351 provider=_PROVIDER, 

352 kind=ProviderErrorKind.SERVER, 

353 ) 

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

355 

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

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

358 prefix = role_model_prefix(role) 

359 ready = self._ready_models() 

360 self._check_estimates(ready) 

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

362 

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

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

365 

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

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

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

369 """ 

370 for model_id in ready - self._estimate_checked: 

371 launch = self._launch_by_model.get(model_id) 

372 self._estimate_checked.add(model_id) 

373 if launch is None: 

374 continue 

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

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

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

378 clear_ctx_downshift(launch.role) 

379 # A launch carrying --memory serves its own report on GET /memory 

380 # and was given no trace log to read (swap_config), so the log-side 

381 # checks would misfire on it by construction. 

382 if MEMORY_FLAG in launch.argv: 

383 self._check_memory_estimate(model_id, launch) 

384 continue 

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

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

387 continue 

388 if launch.est_vram_bytes <= 0: 

389 continue 

390 check_launch( 

391 self._log_path.parent, 

392 model_id, 

393 launch.role, 

394 launch.model, 

395 launch.est_vram_bytes, 

396 launch.est_vram_by_device, 

397 launch.est_unreported_bytes, 

398 ) 

399 

400 def _check_memory_estimate(self, model_id: str, launch: InstanceLaunch) -> None: 

401 """Run the estimate check off the engine's own ``GET /memory``. 

402 

403 Straight to the member's port rather than through llama-swap: the model 

404 is ready, so the server owns its port, and the proxy adds only a routing 

405 layer that has nothing to route on for a bare GET. An endpoint that does 

406 not answer after the engine took the flag is reported, not swallowed -- 

407 it is this mode's analog of a ready engine that wrote no log. 

408 """ 

409 port = self._member_port_by_model.get(model_id) 

410 if port is None: # bound to another process's engine; nothing was planned here 

411 return 

412 try: 

413 resp = _probe_client().get(f"http://{_HOST}:{port}/memory", timeout=_PROBE_TIMEOUT_S) 

414 payload = resp.json() if resp.status_code == httpx.codes.OK else None 

415 except (httpx.HTTPError, ValueError): 

416 payload = None 

417 if payload is None: 

418 log.warning( 

419 "The %s engine was launched with %s but its /memory endpoint did not " 

420 "answer, so its memory use could not be checked against the estimate. " 

421 "Placement estimates for this model are unverified.", 

422 launch.role.value, 

423 MEMORY_FLAG, 

424 ) 

425 return 

426 if launch.est_vram_bytes <= 0: 

427 return 

428 check_memory_report( 

429 launch.role, 

430 launch.model, 

431 launch.est_vram_bytes, 

432 launch.est_vram_by_device, 

433 payload, 

434 ) 

435 

436 def is_live(self) -> bool: 

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

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

439 return False 

440 if self._port is None: 

441 return False 

442 return self._proxy_answers() 

443 

444 @property 

445 def running(self) -> bool: 

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

447 return self._proc is not None 

448 

449 @property 

450 def bound(self) -> bool: 

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

452 return self._bound 

453 

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

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

456 

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

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

459 """ 

460 if state.proxy_port is None: 

461 return False 

462 self._port = state.proxy_port 

463 self._member_ports = list(state.member_ports) 

464 if not self._proxy_answers(): 

465 self._port = None 

466 self._member_ports = [] 

467 return False 

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

469 self._bound = True 

470 return True 

471 

472 def _proxy_answers(self) -> bool: 

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

474 

475 Shares state_is_healthy's identity check via _running_endpoint_answers, so 

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

477 two hand-kept-identical copies. 

478 """ 

479 return _running_endpoint_answers(self.endpoint()) 

480 

481 def shutdown(self) -> None: 

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

483 

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

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

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

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

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

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

490 """ 

491 if self._bound: 

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

493 self._bound = False 

494 self._port = None 

495 self._member_ports = [] 

496 self._launches_payload = [] 

497 return 

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

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

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

501 self._state_path.unlink(missing_ok=True) 

502 if self._proc is not None: 

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

504 release_death_pipe(self._proc.pid) 

505 self._proc = None 

506 self._port = None 

507 self._close_log() 

508 

509 def _close_log(self) -> None: 

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

511 if self._log_file is not None: 

512 with contextlib.suppress(OSError): 

513 self._log_file.close() 

514 self._log_file = None 

515 

516 def _await_health(self) -> None: 

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

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

519 deadline = time.monotonic() + _BOOT_TIMEOUT_S 

520 while time.monotonic() < deadline: 

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

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

523 with contextlib.suppress(httpx.HTTPError): 

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

525 return 

526 time.sleep(_BOOT_POLL_S) 

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

528 

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

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

531 

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

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

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

535 """ 

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

537 payload = ( 

538 _probe_client() 

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

540 .json() 

541 ) 

542 return { 

543 entry[_KEY_MODEL] 

544 for entry in payload[_KEY_RUNNING] 

545 if entry.get(_KEY_STATE) == _STATE_READY 

546 } 

547 return set() 

548 

549 def _boot_log_tail(self) -> str: 

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

551 try: 

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

553 handle.seek(self._log_offset) 

554 data = handle.read() 

555 except OSError: 

556 return "" 

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

558 

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

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

561 self.shutdown() 

562 tail = self._boot_log_tail() 

563 if tail: 

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

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

566 

567 

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

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

570 

571 

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

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

574 try: 

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

576 return int(low), int(high) 

577 except (OSError, ValueError): 

578 return None 

579 

580 

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

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

583 

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

585 letting the OS choose. 

586 """ 

587 from_proc = _port_range_from(_PROC_PORT_RANGE) 

588 if from_proc is not None: 

589 return from_proc 

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

591 out = subprocess.run( 

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

593 capture_output=True, 

594 text=True, 

595 encoding="utf-8", 

596 errors="replace", 

597 timeout=5, 

598 check=False, 

599 ) 

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

601 return int(low), int(high) 

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

603 return None 

604 

605 

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

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

608_PORT_SEARCH_FLOOR = 20000 

609_PORT_WINDOW_SPAN = 8192 

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

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

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

613# into the next block. 

614_PORT_BLOCK = 64 

615 

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

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

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

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

620_reserved_ports: set[int] = set() 

621_reserved_lock = threading.Lock() 

622 

623 

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

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

626 with _reserved_lock: 

627 _reserved_ports.difference_update(ports) 

628 

629 

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

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

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

633 

634 

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

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

637 

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

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

640 on all but one port. 

641 """ 

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

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

644 

645 

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

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

648 

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

650 same port out twice within one allocation. 

651 

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

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

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

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

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

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

658 """ 

659 ceiling = _ephemeral_range() 

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

661 try: 

662 for sock in sockets: 

663 _bind_below_ephemeral(sock, ceiling) 

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

665 finally: 

666 for sock in sockets: 

667 sock.close() 

668 

669 

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

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

672 

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

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

675 ephemeral range for those ports. 

676 """ 

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

678 span = _window_span(ceiling) 

679 start = _search_start(ceiling) 

680 for offset in range(span): 

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

682 with _reserved_lock: 

683 if port in _reserved_ports: 

684 continue 

685 try: 

686 sock.bind((_HOST, port)) 

687 except OSError: 

688 continue 

689 _reserved_ports.add(port) 

690 return 

691 sock.bind((_HOST, 0)) 

692 

693 

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

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

696 try: 

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

698 except psutil.NoSuchProcess: 

699 return [] 

700 return children 

701 

702 

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

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

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

706 for child in survivors: 

707 with contextlib.suppress(psutil.NoSuchProcess): 

708 child.terminate() 

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

710 for child in alive: 

711 with contextlib.suppress(psutil.NoSuchProcess): 

712 child.kill() 

713 _await_killed(alive) 

714 

715 

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

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

718 if not procs: 

719 return 

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

721 for proc in alive: 

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

723 

724 

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

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

727 

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

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

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

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

732 """ 

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

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

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

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

737 if needle in name: 

738 yield proc 

739 

740 

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

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

743 

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

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

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

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

748 """ 

749 target = str(config_path) 

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

751 for proc in _processes_named(_LLAMA_SWAP_PROCESS_NAME): 

752 try: 

753 cmdline = proc.cmdline() 

754 except ( 

755 psutil.NoSuchProcess, 

756 psutil.AccessDenied, 

757 psutil.ZombieProcess, 

758 OSError, 

759 SystemError, 

760 ): 

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

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

763 # C-extension SystemError instead of an AccessDenied. 

764 continue 

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

766 if target in cmdline: 

767 swaps.append(proc) 

768 return swaps 

769 

770 

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

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

773 

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

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

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

777 """ 

778 best: SwapState | None = None 

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

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

781 continue 

782 state = _load_state(state_path) 

783 if state is None: 

784 continue 

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

786 best = state 

787 return best 

788 

789 

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

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

792 just liveness). 

793 

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

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

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

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

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

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

800 """ 

801 try: 

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

803 except (OSError, httpx.HTTPError): 

804 return False 

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

806 return False 

807 try: 

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

809 except (ValueError, AttributeError): 

810 return False 

811 

812 

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

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

815 if state.proxy_port is None: 

816 return False 

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

818 

819 

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

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

822 

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

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

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

826 beside one an HTTP probe failed to see. 

827 """ 

828 return any(data_dir.glob(_STATE_FILE_GLOB)) 

829 

830 

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

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

833 

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

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

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

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

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

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

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

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

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

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

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

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

846 """ 

847 _clean_stale_configs(data_dir) 

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

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

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

851 clear_keep_warm(data_dir) 

852 stopped: list[str] = [] 

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

854 state = _load_state(state_path) 

855 if state is None: 

856 continue 

857 if _stop_recorded_engine(state): 

858 group = _state_group(state_path.name) 

859 if group is not None: 

860 stopped.append(group) 

861 state_path.unlink(missing_ok=True) 

862 return stopped 

863 

864