Coverage for src/lilbee/cli/launchers/server.py: 100%
198 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Server-lifecycle helpers shared by every ``lilbee launch <client>`` command."""
3from __future__ import annotations
5import contextlib
6import json
7import logging
8import os
9import shutil
10import socket
11import subprocess
12import sys
13import time
14from typing import IO
16import httpx
17import typer
19from lilbee.cli.app import console
20from lilbee.cli.commands.servers import port_file
21from lilbee.cli.launchers.warm_render import render_warm
22from lilbee.core.config import cfg
23from lilbee.modelhub.registry import ModelRegistry
24from lilbee.parent_monitor import PARENT_PID_ENV
25from lilbee.providers.fleet.child_guard import spawn_bound_child
26from lilbee.providers.fleet.swap_config import cold_load_timeout_s
27from lilbee.server.auth import server_json_path
29log = logging.getLogger(__name__)
31LOOPBACK = "127.0.0.1"
32"""Loopback address used for launcher-spawned sessions and the URLs we hand to clients."""
34_SERVER_BOOT_TIMEOUT_S = 60.0
35_SERVER_POLL_INTERVAL_S = 0.5
36# Floor on the cold model-load wait; chat_warm_budget_s() scales it up with the weights.
37_WARM_TIMEOUT_S = 600.0
38_HEALTH_PROBE_TIMEOUT_S = 2.0
39_HTTP_OK = 200
40_HEALTH_PATH = "/api/health"
41_TERMINATE_GRACE_S = 10
42_KILL_GRACE_S = 5
43# Spawn attempts; free_port()'s released probe port can be stolen before the server binds.
44_SPAWN_ATTEMPTS = 3
47def running_server_session() -> tuple[str, int] | None:
48 """Return ``(token, port)`` for a server already running on this machine, else None."""
49 session_path = server_json_path()
50 port_path = port_file()
51 if not session_path.exists() or not port_path.exists():
52 return None
53 try:
54 data = json.loads(session_path.read_text(encoding="utf-8"))
55 token = data.get("token")
56 port = int(port_path.read_text(encoding="utf-8").strip())
57 except (json.JSONDecodeError, OSError, ValueError):
58 return None
59 if not isinstance(token, str) or not token:
60 return None
61 return token, port
64def free_port() -> int:
65 """Return an unused TCP port on the loopback interface."""
66 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
67 s.bind((LOOPBACK, 0))
68 return int(s.getsockname()[1])
71def _session_token() -> str | None:
72 """The bearer token from server.json, or None if it is not readable yet."""
73 try:
74 data = json.loads(server_json_path().read_text(encoding="utf-8"))
75 except (json.JSONDecodeError, UnicodeDecodeError, OSError):
76 return None
77 token = data.get("token")
78 return token if isinstance(token, str) and token else None
81def _probe_health(port: int) -> dict[str, object] | None:
82 """GET ``/api/health`` once; return the parsed body on 200, else None.
84 The single place the probe URL, timeout, error handling, and status check
85 live, so the three public probes below stay consistent.
87 Health needs the token like every other route: it reports the chat
88 engine's last error, which carries model paths and loader failures. The
89 token is re-read per attempt rather than captured once, because these
90 probes poll a server that is still starting and server.json does not exist
91 until its lifespan has run. No token yet means no server yet, which is the
92 same answer a refused connection gives.
93 """
94 token = _session_token()
95 if token is None:
96 return None
97 try:
98 resp = httpx.get(
99 f"http://{LOOPBACK}:{port}{_HEALTH_PATH}",
100 timeout=_HEALTH_PROBE_TIMEOUT_S,
101 headers={"Authorization": f"Bearer {token}"},
102 )
103 except httpx.HTTPError:
104 return None
105 if resp.status_code != _HTTP_OK:
106 return None
107 try:
108 body = resp.json()
109 except ValueError:
110 return {}
111 return body if isinstance(body, dict) else {}
114def health_ok(port: int) -> bool:
115 """Single-shot ``/api/health`` probe; True iff a 200 comes back fast."""
116 return _probe_health(port) is not None
119def wait_for_health(port: int, timeout_s: float = _SERVER_BOOT_TIMEOUT_S) -> bool:
120 """Poll ``/api/health`` until it answers 200 or *timeout_s* elapses."""
121 deadline = time.monotonic() + timeout_s
122 while time.monotonic() < deadline:
123 if health_ok(port):
124 return True
125 time.sleep(_SERVER_POLL_INTERVAL_S)
126 return False
129def chat_ready(port: int) -> bool:
130 """Single-shot probe: True iff ``/api/health`` reports the chat engine warm."""
131 body = _probe_health(port)
132 return bool(body and body.get("chat_ready", False))
135def served_chat_ctx(port: int) -> int | None:
136 """The chat window ``/api/health`` reports, or None if unknown/unreachable.
138 A launcher passes this to the client so it trims history to the model's
139 actual window instead of overflowing on a long agentic session.
140 """
141 body = _probe_health(port)
142 if body is None:
143 return None
144 ctx = body.get("chat_ctx")
145 return ctx if isinstance(ctx, int) and ctx > 0 else None
148def planned_chat_ctx() -> int | None:
149 """The per-slot window the fleet will serve the configured chat model, or None.
151 Mirrors the fleet's own single-GPU chat sizing, so it answers before the
152 engine is up: the same ``cfg.num_ctx`` short-circuit, then the same
153 :func:`resolve_chat_ctx` against the same budget the fleet sizes with, which
154 is the memory the GPU reports rather than the host's (see
155 :func:`lilbee.providers.fleet.planning.plan_sizing_budget`).
157 A tensor-split chat is sized by the fleet against per-device headroom
158 instead, so this can over-report there; it is only a fallback for a chat
159 engine that is not up yet, and the served window wins once it is. A
160 remote-served chat model has no local window to compute.
161 """
162 from lilbee.providers.base import ProviderError
163 from lilbee.providers.engine_params import resolve_chat_ctx, resolve_model_path
164 from lilbee.providers.fleet.planning import plan_sizing_budget
165 from lilbee.providers.gguf_meta import read_gguf_metadata
166 from lilbee.providers.model_ref import parse_model_ref
168 ref = str(cfg.chat_model)
169 if not ref or not parse_model_ref(ref).is_local:
170 return None
171 if cfg.num_ctx is not None:
172 return cfg.num_ctx
173 try:
174 path = resolve_model_path(ref)
175 return resolve_chat_ctx(
176 path, read_gguf_metadata(path), available_bytes=plan_sizing_budget()
177 )
178 except (ProviderError, OSError, ValueError):
179 # Sizing needs the model file and its GGUF header; an absent or unreadable
180 # one leaves the window unknown rather than failing the launch.
181 log.debug("planned_chat_ctx failed for %s", ref, exc_info=True)
182 return None
185def client_chat_ctx(port: int) -> int | None:
186 """The chat window to advertise to a launched client, warning when it is small.
188 The chat role builds lazily, so a launcher that hands off before the engine
189 is warm gets nothing from ``/api/health``; fall back to the window the fleet
190 plans to serve rather than leaving the client with no window at all. A window
191 below ``cfg.chat_n_ctx_target`` means the host could not back what was asked
192 for, which changes how much history an agent can keep, so say so.
193 """
194 ctx = served_chat_ctx(port)
195 if ctx is None:
196 ctx = planned_chat_ctx()
197 if ctx is not None and ctx < cfg.chat_n_ctx_target:
198 typer.secho(
199 f"Warning: the chat model is served with a {ctx:,}-token context, below the "
200 f"configured chat_n_ctx_target of {cfg.chat_n_ctx_target:,}. Either the model "
201 "was trained on a smaller window, or its weights leave too little of the "
202 "memory budget for the KV cache. A longer-context model, a smaller "
203 "quantization, or a higher gpu_memory_fraction raises it.",
204 err=True,
205 fg=typer.colors.YELLOW,
206 )
207 return ctx
210def chat_warm_budget_s() -> float:
211 """Warm wait scaled to the chat model's on-disk weights at the engine's cold-load rate."""
212 try:
213 shards = ModelRegistry(cfg.models_dir).shard_paths(str(cfg.chat_model))
214 except (KeyError, ValueError):
215 return _WARM_TIMEOUT_S
216 total_bytes = sum(shard.stat().st_size for shard in shards)
217 return max(_WARM_TIMEOUT_S, float(cold_load_timeout_s(total_bytes)))
220def wait_for_chat_warm(port: int, timeout_s: float | None = None) -> bool:
221 """Block until the chat model is loaded, showing granular warm progress.
223 The server warms the chat role on a background thread at startup, so a client
224 launched the instant the HTTP port binds would otherwise hit an
225 apparently-dead stream during the cold model load. Streams ``/api/warm/stream``
226 to render a real read-phase byte bar then an engine-load spinner; falls back
227 to a plain readiness poll when that stream can't be opened.
228 Returns True once the chat engine reports ready, or False if the budget
229 (weights-scaled via :func:`chat_warm_budget_s` unless given) elapses first;
230 the caller proceeds either way, so a still-loading model just warms on the
231 first call.
232 """
233 if timeout_s is None:
234 timeout_s = chat_warm_budget_s()
235 if chat_ready(port):
236 return True
237 streamed = render_warm(f"http://{LOOPBACK}:{port}", timeout_s)
238 if streamed is not None:
239 # The stream ran (ready, error, or its own timeout); don't double-spend
240 # the budget on a second poll. The caller proceeds on False regardless.
241 return streamed
242 return _poll_chat_ready(port, timeout_s)
245def _poll_chat_ready(port: int, timeout_s: float) -> bool:
246 """Fallback warm wait when the progress stream is unavailable: poll readiness."""
247 deadline = time.monotonic() + timeout_s
248 with console.status("Warming the chat model..."):
249 while time.monotonic() < deadline:
250 if chat_ready(port):
251 return True
252 time.sleep(_SERVER_POLL_INTERVAL_S)
253 return False
256def spawn_server(
257 port: int, *, env_overrides: dict[str, str] | None = None
258) -> subprocess.Popen[bytes]:
259 """Spawn ``lilbee serve --port <port>`` as a background subprocess.
261 Prefers the ``lilbee`` binary on PATH so frozen builds (Nuitka standalone)
262 spawn the binary directly. Falls back to ``sys.executable -m lilbee`` for
263 pip / editable installs where the entry point shims to the same form.
265 ``env_overrides`` are layered onto the inherited environment for the child
266 (e.g. ``LILBEE_CHAT_N_CTX_TARGET`` to size the served window for a launched
267 agent); ``None`` inherits the parent environment unchanged.
269 Stdout/stderr go to ``cfg.data_dir / "logs" / "launcher-serve.log"`` (size
270 capped at 5 MB) so a crash mid-session leaves a trace instead of disappearing.
271 Set ``LILBEE_LAUNCHER_SERVE_QUIET=1`` to restore the previous DEVNULL behavior.
272 """
273 lilbee_bin = shutil.which("lilbee")
274 # On Windows, pip/uv may install a ``lilbee.cmd`` wrapper instead of a bare
275 # executable. Popen(shell=False) raises PermissionError on .cmd files, so
276 # fall through to the sys.executable -m lilbee form in that case.
277 _bin_is_cmd = sys.platform == "win32" and (
278 lilbee_bin is not None and lilbee_bin.lower().endswith(".cmd")
279 )
280 cmd = (
281 [lilbee_bin, "serve", "--port", str(port)]
282 if lilbee_bin is not None and not _bin_is_cmd
283 else [sys.executable, "-m", "lilbee", "serve", "--port", str(port)]
284 )
286 log_file: IO[bytes] | None = None
287 if os.environ.get("LILBEE_LAUNCHER_SERVE_QUIET"):
288 stdout: int | IO[bytes] = subprocess.DEVNULL
289 stderr: int | IO[bytes] = subprocess.DEVNULL
290 else:
291 log_dir = cfg.data_dir / "logs"
292 log_dir.mkdir(parents=True, exist_ok=True)
293 log_path = log_dir / "launcher-serve.log"
294 # Truncate when the file passes 5 MB so a long-lived session doesn't
295 # accumulate the chat-completion firehose into the data dir indefinitely.
296 # On Windows the file may still be held open by a previous session, so
297 # fall through to append mode when unlink is denied.
298 if log_path.exists() and log_path.stat().st_size > 5 * 1024 * 1024:
299 with contextlib.suppress(OSError):
300 log_path.unlink()
301 log_file = log_path.open("ab")
302 stdout = log_file
303 stderr = subprocess.STDOUT
305 # LILBEE_PARENT_PID arms serve's parent-death watcher, so a hard-killed
306 # launcher (whose finally never runs) does not orphan serve holding server_lock.
307 child_env = {**os.environ, **(env_overrides or {}), PARENT_PID_ENV: str(os.getpid())}
309 try:
310 return spawn_bound_child(
311 cmd,
312 stdout=stdout,
313 stderr=stderr,
314 env=child_env,
315 )
316 finally:
317 # Popen dups the fd into the child; the parent's handle is no longer
318 # needed and would otherwise leak for the launcher's whole lifetime.
319 if log_file is not None:
320 log_file.close()
323def stop_spawned_server(proc: subprocess.Popen[bytes]) -> None:
324 """Terminate *proc* gracefully, escalating to kill if it ignores SIGTERM."""
325 if proc.poll() is not None:
326 return
327 proc.terminate()
328 try:
329 proc.wait(timeout=_TERMINATE_GRACE_S)
330 except subprocess.TimeoutExpired:
331 proc.kill()
332 proc.wait(timeout=_KILL_GRACE_S)
335def ensure_server_running(
336 *, env_overrides: dict[str, str] | None = None
337) -> tuple[tuple[str, int], subprocess.Popen[bytes] | None]:
338 """Return ``(session, spawned_proc)`` for a usable lilbee server.
340 Reuses an already-running server when its session files are healthy.
341 Otherwise spawns a fresh server on a free port. The returned ``spawned_proc``
342 is ``None`` when an existing server was reused; the caller is responsible
343 for stopping a spawned process when it is done with it.
345 ``env_overrides`` reach a freshly spawned child (e.g. a launcher sizing the
346 served window); a reused server keeps whatever window it booted with.
347 """
348 existing = running_server_session()
349 if existing is not None and health_ok(existing[1]):
350 return existing, None
351 last_port = 0
352 for _ in range(_SPAWN_ATTEMPTS):
353 # Honor a user-pinned port so a persisted agent config keeps a valid URL;
354 # fall back to a free port when unset (0).
355 last_port = cfg.server_port or free_port()
356 spawned = _spawn_and_wait(last_port, env_overrides=env_overrides)
357 if spawned is not None:
358 return _session_for_spawned(spawned), spawned
359 typer.secho(
360 f"lilbee server failed to start on port {last_port}; check the logs.",
361 err=True,
362 fg=typer.colors.RED,
363 )
364 raise typer.Exit(1)
367def _spawn_and_wait(
368 port: int, *, env_overrides: dict[str, str] | None = None
369) -> subprocess.Popen[bytes] | None:
370 """Spawn a server on *port* and wait for health; None when it never comes up."""
371 spawned = spawn_server(port, env_overrides=env_overrides)
372 with console.status(f"Starting lilbee server on port {port}..."):
373 healthy = wait_for_health(port)
374 if healthy:
375 return spawned
376 stop_spawned_server(spawned)
377 return None
380def _session_for_spawned(spawned: subprocess.Popen[bytes]) -> tuple[str, int]:
381 """Read the session a freshly-healthy server wrote, stopping it when missing."""
382 session = running_server_session()
383 if session is None:
384 stop_spawned_server(spawned)
385 typer.secho(
386 "lilbee server started but did not write a session file; cannot continue.",
387 err=True,
388 fg=typer.colors.RED,
389 )
390 raise typer.Exit(1)
391 return session