| """The replica router (D-031): one public Space in front of two private model Spaces. |
| |
| Design and limits: docs/SERVING.md "Scaling beyond one server". No torch here: the router |
| holds no model, so it has a small image and starts quickly. |
| |
| The decisions do no I/O: `classify` says what a replica's reply means, `Picker` chooses |
| where a request goes. `ReplicaRouter` carries a visitor's request to a replica and back. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import asyncio |
| import json |
| import logging |
| import math |
| import os |
| import re |
| import time |
| from collections.abc import AsyncIterator, Awaitable, Callable |
| from contextlib import asynccontextmanager |
| from dataclasses import asdict, dataclass |
| from enum import Enum |
|
|
| import httpx |
| from fastapi import FastAPI |
| from fastapi import Request as HTTPRequest |
| from fastapi.responses import HTMLResponse, Response, StreamingResponse |
| from starlette.requests import ClientDisconnect |
|
|
| from tinyllm.obs import EngineSnapshot, Metrics, RequestRecord, itl_stats, setup_logging |
| from tinyllm.web_common import (CLIENT_HEADER, HeadersMiddleware, RateLimiter, client_hash, error_response, read_page, |
| resolve_salt) |
|
|
| log = logging.getLogger("tinyllm.router") |
|
|
|
|
| class Verdict(Enum): |
| """What one reply from a replica means. See `classify`.""" |
| OK = "ok" |
| NOT_READY = "not_ready" |
| ASLEEP = "asleep" |
| DOWN = "down" |
| MISCONFIGURED = "misconfigured" |
| BUSY = "busy" |
| PASS = "pass" |
|
|
|
|
| |
| RETRYABLE = frozenset({Verdict.NOT_READY, Verdict.ASLEEP, Verdict.DOWN, Verdict.MISCONFIGURED, Verdict.BUSY}) |
|
|
|
|
| @dataclass(frozen=True) |
| class Reply: |
| """A replica's answer. `status` is None for a connection error or timeout.""" |
| status: int | None |
| content_type: str = "" |
| text: str = "" |
|
|
|
|
| def _error_type(text: str) -> str | None: |
| try: |
| return json.loads(text)["error"]["type"] |
| except (ValueError, KeyError, TypeError): |
| return None |
|
|
|
|
| def classify(reply: Reply) -> Verdict: |
| """The reply table in docs/SERVING.md, as code. Anything unrecognised is DOWN: never ready, |
| never woken, and never shown to a visitor, because it didn't come from our server.""" |
| status = reply.status |
| if status is None or status == 502 or 300 <= status < 400: |
| return Verdict.DOWN |
| if status in (401, 403): |
| return Verdict.MISCONFIGURED |
| kind = reply.content_type.split(";")[0].strip().lower() |
| if kind == "text/html": |
| if status == 404: |
| return Verdict.MISCONFIGURED |
| if status == 503 and "paused" in reply.text.lower(): |
| return Verdict.ASLEEP |
| return Verdict.DOWN |
| if kind not in ("application/json", "text/event-stream"): |
| return Verdict.DOWN |
| if 200 <= status < 300: |
| return Verdict.OK |
| if status == 503: |
| return Verdict.NOT_READY |
| if status == 429 and _error_type(reply.text) == "server_busy": |
| return Verdict.BUSY |
| return Verdict.PASS |
|
|
|
|
| @dataclass(eq=False) |
| class Replica: |
| """What the router knows about one replica.""" |
| name: str |
| url: str |
| verdict: Verdict | None = None |
| outstanding: int = 0 |
| n_slots: int = 0 |
| sent_since_poll: int = 0 |
| polled_at: float | None = None |
| space_id: str = "" |
| woken_at: float | None = None |
| model: dict | None = None |
|
|
| @property |
| def ready(self) -> bool: |
| return self.verdict is Verdict.OK |
|
|
| @property |
| def load(self) -> int: |
| return self.outstanding + self.sent_since_poll |
|
|
| def refused(self, verdict: Verdict) -> None: |
| """A request came back with a RETRYABLE verdict. Only a full queue leaves it ready.""" |
| self.sent_since_poll = max(0, self.sent_since_poll - 1) |
| if verdict is not Verdict.BUSY: |
| self.verdict = verdict |
|
|
| def polled(self, verdict: Verdict, now: float, outstanding: int = 0, n_slots: int = 0) -> None: |
| self.verdict, self.polled_at = verdict, now |
| self.outstanding, self.n_slots, self.sent_since_poll = outstanding, n_slots, 0 |
|
|
|
|
| class Picker: |
| """Shortest queue among the ready replicas; ties go round-robin. Event loop only.""" |
|
|
| def __init__(self, replicas: list[Replica]): |
| self.replicas = replicas |
| self._last = -1 |
|
|
| def pick(self, exclude: tuple[Replica, ...] = ()) -> Replica | None: |
| """The replica to send to, counted as sent; None if no ready replica is left.""" |
| n = len(self.replicas) |
| in_turn = [self.replicas[(self._last + 1 + i) % n] for i in range(n)] |
| ready = [r for r in in_turn if r.ready and r not in exclude] |
| if not ready: |
| return None |
| choice = min(ready, key=lambda r: r.load) |
| self._last = self.replicas.index(choice) |
| choice.sent_since_poll += 1 |
| return choice |
|
|
|
|
| |
|
|
| MAX_ATTEMPTS = 2 |
| NO_REPLICA = "none" |
| _PASSED_HEADERS = ("content-type", "cache-control", "retry-after") |
| _STREAM_ERROR = (b'data: {"error": {"message": "the replica failed", "type": "server_error", "code": 502}}\n\n' |
| b"data: [DONE]\n\n") |
|
|
|
|
| @dataclass |
| class RouterConfig: |
| rate_per_min: float = 20.0 |
| rate_burst: int = 5 |
| trust_proxy: bool = False |
| proxy_hops: int = 1 |
| salt: str | None = None |
| connect_timeout_s: float = 3.0 |
| read_timeout_s: float = 65.0 |
| max_body_bytes: int = 1 << 20 |
| waking_retry_s: int = 30 |
| poll_interval_s: float = 5.0 |
| poll_timeout_s: float = 2.0 |
| idle_stop_s: float = 600.0 |
| wake_cooldown_s: float = 300.0 |
| model_label: str | None = None |
| drain_s: float = 30.0 |
|
|
|
|
| class _ClientGone(Exception): |
| """The visitor disconnected.""" |
|
|
|
|
| class _BodyTooLarge(Exception): |
| pass |
|
|
|
|
| class _NoAnswer(Exception): |
| """A replica accepted the request and then went silent. It keeps the request; no retry.""" |
|
|
|
|
| def _json(data: bytes) -> dict: |
| try: |
| out = json.loads(data) |
| except ValueError: |
| return {} |
| return out if isinstance(out, dict) else {} |
|
|
|
|
| def _wants_stream(body: bytes) -> bool: |
| return _json(body).get("stream") is True |
|
|
|
|
| async def _read_body(request: HTTPRequest, limit: int) -> bytes: |
| declared = request.headers.get("content-length", "") |
| if declared.isdigit() and int(declared) > limit: |
| raise _BodyTooLarge |
| parts, size = [], 0 |
| async for chunk in request.stream(): |
| size += len(chunk) |
| if size > limit: |
| raise _BodyTooLarge |
| parts.append(chunk) |
| return b"".join(parts) |
|
|
|
|
| async def _unless_gone(request: HTTPRequest, work: Awaitable): |
| """Await `work`; if the visitor leaves first, cancel it (which closes the replica's connection).""" |
| task = asyncio.ensure_future(work) |
| try: |
| while True: |
| done, _ = await asyncio.wait({task}, timeout=0.25) |
| if done: |
| return task.result() |
| if await request.is_disconnected(): |
| raise _ClientGone |
| finally: |
| if not task.done(): |
| task.cancel() |
| await asyncio.gather(task, return_exceptions=True) |
|
|
|
|
| class _StreamMeter: |
| """Reads the events of a stream as they pass through, for the request record.""" |
|
|
| def __init__(self): |
| self.buffer = b"" |
| self.times: list[float] = [] |
| self.prompt_tokens = self.completion_tokens = 0 |
| self.finish: str | None = None |
| self.error: dict | None = None |
| self.done = False |
|
|
| def feed(self, chunk: bytes, now: float) -> None: |
| self.buffer += chunk |
| *events, self.buffer = self.buffer.split(b"\n\n") |
| for event in events: |
| self._event(event, now) |
|
|
| def _event(self, event: bytes, now: float) -> None: |
| if not event.startswith(b"data: "): |
| return |
| payload = event[len(b"data: "):] |
| if payload == b"[DONE]": |
| self.done = True |
| return |
| d = _json(payload) |
| if isinstance(d.get("error"), dict): |
| self.error = d["error"] |
| elif d.get("choices"): |
| self.times.append(now) |
| self.finish = d["choices"][0].get("finish_reason") or self.finish |
| usage = d.get("usage") or {} |
| self.prompt_tokens = usage.get("prompt_tokens", self.prompt_tokens) |
| self.completion_tokens = usage.get("completion_tokens", self.completion_tokens) |
|
|
|
|
| def _reply(resp: httpx.Response | None) -> Reply: |
| return Reply(None) if resp is None else Reply(resp.status_code, resp.headers.get("content-type", ""), resp.text) |
|
|
|
|
| def _load(stats: bytes) -> tuple[int, int] | None: |
| """(outstanding, n_slots) from a replica's /stats, or None if it doesn't have them.""" |
| variants = _json(stats).get("variants") |
| if not isinstance(variants, dict) or not variants: |
| return None |
| try: |
| return sum(int(v["outstanding"]) for v in variants.values()), sum(int(v["n_slots"]) for v in variants.values()) |
| except (KeyError, TypeError, ValueError): |
| return None |
|
|
|
|
| def _model(stats: bytes) -> dict | None: |
| """What a replica serves, from its /stats. A replica has one variant.""" |
| variants = _json(stats).get("variants") |
| if not isinstance(variants, dict) or not variants: |
| return None |
| name, v = next(iter(variants.items())) |
| return {"id": name, "version": v.get("version"), "norm_backend": v.get("norm_backend")} if isinstance(v, dict) else None |
|
|
|
|
| class ReplicaRouter: |
| def __init__(self, replicas: list[Replica], cfg: RouterConfig, token: str | None = None, |
| metrics: Metrics | None = None, wake: Callable[[Replica], None] | None = None): |
| self.cfg = cfg |
| self.replicas = replicas |
| self.picker = Picker(replicas) |
| self.limiter = RateLimiter(cfg.rate_per_min, cfg.rate_burst) |
| self.salt, self.salt_source = resolve_salt(cfg.salt) |
| self.metrics = metrics or Metrics() |
| self.http: httpx.AsyncClient | None = None |
| self._token = token |
| self.wake = wake |
| self.last_traffic = 0.0 |
| self._poller: asyncio.Task | None = None |
| self._fresh = asyncio.Event() |
| self._background: set[asyncio.Task] = set() |
| self._tokens_since_snap = 0 |
| self._last_snap = time.monotonic() |
|
|
| def record(self, rec: RequestRecord, arrival: float) -> None: |
| rec.latency_s = time.monotonic() - arrival |
| self.metrics.observe(rec) |
| self._tokens_since_snap += rec.completion_tokens |
| log.info("request", extra={"fields": rec.to_dict()}) |
|
|
| def error(self, rec: RequestRecord, arrival: float, status: int, message: str, err_type: str, |
| headers: dict | None = None) -> Response: |
| rec.status, rec.error, rec.finish_reason = status, message, "timeout" if status == 504 else "error" |
| self.record(rec, arrival) |
| return error_response(status, message, err_type, headers) |
|
|
| def gone(self, rec: RequestRecord, arrival: float) -> Response: |
| rec.status, rec.finish_reason = 499, "abort" |
| self.record(rec, arrival) |
| return Response(status_code=499) |
|
|
| def _auth(self) -> dict: |
| return {"authorization": f"Bearer {self._token}"} if self._token else {} |
|
|
| async def traffic(self, request: HTTPRequest) -> None: |
| """A visitor's request is going to be routed. Only this starts polling, so an open status tab |
| can't keep replicas awake. After an idle spell what we know is stale: wait for one poll round.""" |
| self.last_traffic = time.monotonic() |
| if self._poller is None: |
| self._fresh = asyncio.Event() |
| self._poller = asyncio.create_task(self._poll_loop()) |
| if not self._fresh.is_set(): |
| await _unless_gone(request, self._fresh.wait()) |
|
|
| async def _poll_loop(self) -> None: |
| try: |
| while True: |
| await self.poll_round() |
| self._fresh.set() |
| if time.monotonic() - self.last_traffic >= self.cfg.idle_stop_s: |
| return |
| await asyncio.sleep(self.cfg.poll_interval_s) |
| finally: |
| self._poller = None |
| self._fresh.set() |
|
|
| async def poll_round(self) -> None: |
| results = await asyncio.gather(*(self.poll(r) for r in self.replicas), return_exceptions=True) |
| for replica, result in zip(self.replicas, results): |
| if isinstance(result, Exception): |
| log.error("poll failed", extra={"fields": {"replica": replica.name, "error": type(result).__name__}}) |
| replica.polled(Verdict.DOWN, time.monotonic()) |
| self._snapshot() |
|
|
| async def _get(self, replica: Replica, path: str) -> httpx.Response | None: |
| """None if there was no answer. The exception's text names the private URL, so it's dropped.""" |
| try: |
| return await self.http.get(f"{replica.url}{path}", headers=self._auth(), timeout=self.cfg.poll_timeout_s) |
| except httpx.TransportError: |
| return None |
|
|
| async def poll(self, replica: Replica) -> None: |
| try: |
| ready, stats = await asyncio.wait_for(asyncio.gather(self._get(replica, "/ready"), self._get(replica, "/stats")), |
| self.cfg.poll_timeout_s) |
| except asyncio.TimeoutError: |
| ready = stats = None |
| before, verdict, now = replica.verdict, classify(_reply(ready)), time.monotonic() |
| if verdict is Verdict.OK: |
| readable = classify(_reply(stats)) is Verdict.OK |
| load = _load(stats.content) if readable else None |
| replica.model = (_model(stats.content) if readable else None) or replica.model |
| |
| replica.polled(verdict, now, *(load or (replica.load, replica.n_slots))) |
| else: |
| replica.polled(verdict, now) |
| if verdict is not before: |
| level = logging.ERROR if verdict is Verdict.MISCONFIGURED else logging.INFO |
| log.log(level, "replica state", extra={"fields": {"replica": replica.name, "state": verdict.value}}) |
|
|
| def _snapshot(self) -> None: |
| """For the status page's load chart. A gap in the chart means no replica was ready.""" |
| now = time.monotonic() |
| elapsed = max(now - self._last_snap, 1e-9) |
| tokens, self._tokens_since_snap, self._last_snap = self._tokens_since_snap, 0, now |
| up = [r for r in self.replicas if r.ready] |
| if not up: |
| return |
| busy = sum(min(r.outstanding, r.n_slots) for r in up) |
| queued = sum(r.outstanding for r in up) - busy |
| snap = EngineSnapshot(time.time(), queued, busy, sum(r.n_slots for r in up), round(tokens / elapsed, 2), |
| 0.0) |
| self.metrics.snapshot(snap) |
| if busy or queued or tokens: |
| log.info("engine", extra={"fields": {"variant": "router", **asdict(snap)}}) |
|
|
| def _known_model(self) -> dict | None: |
| return next((r.model for r in self.replicas if r.model), None) |
|
|
| def models(self) -> dict: |
| model, label = self._known_model(), self.cfg.model_label |
| data = [] |
| if model or label: |
| data = [{"id": model["id"] if model else label, "object": "model", "owned_by": "tinyllm", |
| "version": model["version"] if model else label}] |
| return {"object": "list", "data": data} |
|
|
| def describe(self, out: dict) -> dict: |
| """Adds the fields the pages read. Names and states only: never a URL or a Space id.""" |
| model, up, now = self._known_model() or {}, [r for r in self.replicas if r.ready], time.monotonic() |
| name = model.get("id") or self.cfg.model_label or "router" |
| out["variants"] = {name: {"version": model.get("version") or self.cfg.model_label, |
| "outstanding": sum(r.load for r in up), "n_slots": sum(r.n_slots for r in up), |
| "failed": None, "norm_backend": model.get("norm_backend")}} |
| out["polling"] = self._poller is not None |
| out["replicas"] = [{"name": r.name, "state": r.verdict.value if r.verdict else None, "queue": r.load, |
| "n_slots": r.n_slots, |
| "polled_s_ago": None if r.polled_at is None else round(now - r.polled_at, 1)} |
| for r in self.replicas] |
| return out |
|
|
| def wake_asleep(self, now: float) -> None: |
| """Only a replica HF says is paused: restarting a running one would kill its requests.""" |
| if self.wake is None: |
| return |
| for r in self.replicas: |
| if r.verdict is Verdict.ASLEEP and (r.woken_at is None or now - r.woken_at >= self.cfg.wake_cooldown_s): |
| r.woken_at = now |
| task = asyncio.create_task(self._wake(r)) |
| self._background.add(task) |
| task.add_done_callback(self._background.discard) |
|
|
| async def _wake(self, replica: Replica) -> None: |
| log.info("waking replica", extra={"fields": {"replica": replica.name}}) |
| try: |
| await asyncio.to_thread(self.wake, replica) |
| except Exception as e: |
| log.error("wake failed", extra={"fields": {"replica": replica.name, "error": type(e).__name__}}) |
|
|
| async def stop(self) -> None: |
| tasks = [*self._background, *([self._poller] if self._poller else [])] |
| for task in tasks: |
| task.cancel() |
| await asyncio.gather(*tasks, return_exceptions=True) |
|
|
| async def completions(self, request: HTTPRequest) -> Response: |
| arrival = time.monotonic() |
| client = client_hash(request, self.cfg.trust_proxy, self.cfg.proxy_hops, self.salt) |
| rec = RequestRecord(request.state.request_id, 0, variant=NO_REPLICA, client=client, arrival_unix=time.time()) |
| wait = self.limiter.check(client, arrival) |
| if wait: |
| return self.error(rec, arrival, 429, "rate limit exceeded", "rate_limit_error", |
| {"Retry-After": str(math.ceil(wait))}) |
| try: |
| body = await _read_body(request, self.cfg.max_body_bytes) |
| except _BodyTooLarge: |
| return self.error(rec, arrival, 413, "request body is too large", "invalid_request_error") |
| except ClientDisconnect: |
| return self.gone(rec, arrival) |
| rec.stream = _wants_stream(body) |
| try: |
| await self.traffic(request) |
| except _ClientGone: |
| return self.gone(rec, arrival) |
| return await self._route(request, rec, arrival, body) |
|
|
| async def _route(self, request: HTTPRequest, rec: RequestRecord, arrival: float, body: bytes) -> Response: |
| tried: tuple[Replica, ...] = () |
| busy = False |
| try: |
| for _ in range(MAX_ATTEMPTS): |
| replica = self.picker.pick(exclude=tried) |
| if replica is None: |
| break |
| tried += (replica,) |
| rec.variant = replica.name |
| verdict, resp, data = await self._send(request, replica, body, rec.client) |
| if verdict not in RETRYABLE: |
| return await self._answer(request, rec, arrival, replica, verdict, resp, data) |
| replica.refused(verdict) |
| busy = busy or verdict is Verdict.BUSY |
| except _ClientGone: |
| return self.gone(rec, arrival) |
| except _NoAnswer: |
| return self.error(rec, arrival, 504, "the replica did not answer in time", "server_error") |
| if busy or not any(r.ready for r in self.replicas): |
| self.wake_asleep(time.monotonic()) |
| if not tried: |
| rec.variant = NO_REPLICA |
| if busy: |
| return self.error(rec, arrival, 429, "server busy: queue full, retry later", "server_busy", |
| {"Retry-After": "1"}) |
| return self.error(rec, arrival, 503, "no replica is ready: waking up, retry shortly", "waking_up", |
| {"Retry-After": str(self.cfg.waking_retry_s)}) |
|
|
| def _outgoing(self, request: HTTPRequest, client: str) -> dict: |
| """Built from nothing, so nothing else the visitor sent (cookies, Authorization) reaches a replica.""" |
| headers = {"content-type": request.headers.get("content-type", "application/json"), |
| "accept-encoding": "identity", "x-request-id": request.state.request_id, CLIENT_HEADER: client} |
| return {**headers, **self._auth()} |
|
|
| async def _send(self, request: HTTPRequest, replica: Replica, body: bytes, |
| client: str) -> tuple[Verdict, httpx.Response | None, bytes]: |
| """Until the reply's headers (and its body, unless it's a 2xx). The exception text names the |
| private URL, so only its type is logged.""" |
| req = self.http.build_request("POST", f"{replica.url}/v1/completions", content=body, |
| headers=self._outgoing(request, client)) |
| try: |
| resp = await _unless_gone(request, self.http.send(req, stream=True)) |
| content_type = resp.headers.get("content-type", "") |
| if classify(Reply(resp.status_code, content_type)) is Verdict.OK: |
| return Verdict.OK, resp, b"" |
| data = await _unless_gone(request, resp.aread()) |
| except httpx.ReadTimeout as e: |
| raise _NoAnswer from e |
| except httpx.TransportError as e: |
| log.warning("replica unreachable", extra={"fields": {"replica": replica.name, "error": type(e).__name__}}) |
| return classify(Reply(None)), None, b"" |
| return classify(Reply(resp.status_code, content_type, data.decode("utf-8", "replace"))), resp, data |
|
|
| async def _answer(self, request: HTTPRequest, rec: RequestRecord, arrival: float, replica: Replica, |
| verdict: Verdict, resp: httpx.Response, data: bytes) -> Response: |
| rec.model_version = resp.headers.get("x-model-version", "") |
| if rec.model_version: |
| request.state.model_version = rec.model_version |
| headers = {k: resp.headers[k] for k in _PASSED_HEADERS if k in resp.headers} |
| headers["X-Replica"] = replica.name |
| if verdict is Verdict.OK and "text/event-stream" in resp.headers.get("content-type", ""): |
| return StreamingResponse(self._relay(resp, rec, arrival), status_code=resp.status_code, headers=headers) |
| if verdict is Verdict.OK: |
| try: |
| data = await _unless_gone(request, resp.aread()) |
| except httpx.TransportError: |
| return self.error(rec, arrival, 502, "the replica failed", "server_error") |
| d = _json(data) |
| rec.status = resp.status_code |
| if verdict is Verdict.OK: |
| |
| rec.ttft_s = time.monotonic() - arrival |
| usage, choices = d.get("usage") or {}, d.get("choices") or [{}] |
| rec.prompt_tokens = usage.get("prompt_tokens", 0) |
| rec.completion_tokens = usage.get("completion_tokens", 0) |
| rec.finish_reason = choices[0].get("finish_reason") |
| else: |
| err = d.get("error") |
| rec.error = err.get("message") if isinstance(err, dict) else None |
| rec.finish_reason = "timeout" if rec.status == 504 else "error" |
| self.record(rec, arrival) |
| return Response(data, status_code=resp.status_code, headers=headers) |
|
|
| async def _relay(self, resp: httpx.Response, rec: RequestRecord, arrival: float) -> AsyncIterator[bytes]: |
| """Pass a stream through. Nothing can be retried now: the visitor already has part of the story.""" |
| meter, finished = _StreamMeter(), False |
| try: |
| try: |
| async for chunk in resp.aiter_bytes(): |
| meter.feed(chunk, time.monotonic()) |
| yield chunk |
| except httpx.TransportError as e: |
| log.warning("replica stream failed", extra={"fields": {"replica": rec.variant, "error": type(e).__name__}}) |
| if meter.done: |
| rec.status = meter.error.get("code", 500) if meter.error else 200 |
| rec.error = meter.error.get("message") if meter.error else None |
| rec.finish_reason = ("timeout" if rec.status == 504 else "error") if meter.error else meter.finish |
| else: |
| rec.status, rec.error, rec.finish_reason = 502, "the replica failed", "error" |
| finished = True |
| yield (b"\n\n" if meter.buffer else b"") + _STREAM_ERROR |
| finished = True |
| finally: |
| if not finished: |
| rec.status, rec.finish_reason = 499, "abort" |
| if meter.times: |
| rec.ttft_s = meter.times[0] - arrival |
| rec.itl_p50_s, rec.itl_p95_s = itl_stats(meter.times) |
| rec.prompt_tokens, rec.completion_tokens = meter.prompt_tokens, meter.completion_tokens |
| self.record(rec, arrival) |
| await resp.aclose() |
|
|
|
|
| def hf_wake(token: str | None) -> Callable[[Replica], None]: |
| """A wake function that restarts the replica's Space on the Hub. It blocks: run it in a thread.""" |
| def wake(replica: Replica) -> None: |
| from huggingface_hub import HfApi |
| HfApi(token=token).restart_space(replica.space_id) |
| return wake |
|
|
|
|
| def create_router_app(replicas: list[Replica], cfg: RouterConfig | None = None, token: str | None = None, |
| metrics: Metrics | None = None, wake: Callable[[Replica], None] | None = None, |
| shipper=None) -> FastAPI: |
| """Only the routes listed here exist, so no other path of a replica can be reached through the router.""" |
| cfg = cfg or RouterConfig() |
| router = ReplicaRouter(replicas, cfg, token, metrics, wake) |
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| timeout = httpx.Timeout(connect=cfg.connect_timeout_s, read=cfg.read_timeout_s, write=10.0, |
| pool=cfg.connect_timeout_s) |
| |
| async with httpx.AsyncClient(timeout=timeout, trust_env=False, follow_redirects=False) as client: |
| router.http = client |
| log.info("server ready", extra={"fields": {"replicas": [r.name for r in replicas], "wake": wake is not None, |
| "slo_ttft_s": router.metrics.slo_ttft_s, "salt": router.salt_source, |
| "proxy_hops": cfg.proxy_hops if cfg.trust_proxy else 0}}) |
| yield |
| await router.stop() |
| log.info("stopped") |
| if shipper is not None: |
| await asyncio.to_thread(shipper.close, cfg.drain_s) |
|
|
| app = FastAPI(title="tinyllm router", lifespan=lifespan, docs_url=None, redoc_url=None, openapi_url=None) |
| app.state.router = router |
| app.add_middleware(HeadersMiddleware, default_version="router") |
|
|
| |
| @app.get("/health") |
| async def health(): |
| return {"status": "ok"} |
|
|
| @app.get("/ready") |
| async def ready(): |
| return {"ready": True} |
|
|
| @app.get("/v1/models") |
| async def models(): |
| return router.models() |
|
|
| @app.get("/stats") |
| async def stats(): |
| return router.describe(router.metrics.stats()) |
|
|
| history_cache: dict = {} |
|
|
| @app.get("/stats/history") |
| async def stats_history(): |
| """For the status page: computed at most once a second, however many pages poll it.""" |
| now = time.monotonic() |
| if history_cache.get("t", -1.0) < now - 1.0: |
| history_cache.update(t=now, out=router.describe(router.metrics.history())) |
| return history_cache["out"] |
|
|
| @app.get("/", response_class=HTMLResponse) |
| async def demo_page(): |
| return HTMLResponse(read_page("demo.html")) |
|
|
| @app.get("/status", response_class=HTMLResponse) |
| async def status_page(): |
| return HTMLResponse(read_page("status.html")) |
|
|
| @app.post("/v1/completions") |
| async def completions(request: HTTPRequest): |
| return await router.completions(request) |
|
|
| return app |
|
|
|
|
| |
|
|
| REPLICA_TOKEN_ENV = "TINYLLM_REPLICA_TOKEN" |
| _REPLICA_NAME = re.compile(r"[A-Za-z0-9_.-]{1,32}") |
| _SPACE_ID = re.compile(r"[\w.-]+/[\w.-]+") |
|
|
|
|
| def parse_replica(spec: str) -> Replica: |
| """`name=url` or `name=url,owner/space`; the Space id is what waking needs.""" |
| name, sep, rest = spec.partition("=") |
| url, _, space = rest.partition(",") |
| url = url.rstrip("/") |
| if not (sep and _REPLICA_NAME.fullmatch(name) and url.startswith(("http://", "https://")) |
| and (not space or _SPACE_ID.fullmatch(space))): |
| raise ValueError(f"expected name=url[,owner/space], got {spec!r}") |
| return Replica(name, url, space_id=space) |
|
|
|
|
| def main(argv: list[str] | None = None) -> None: |
| p = argparse.ArgumentParser(description="tinyllm replica router (D-031)") |
| p.add_argument("--replica", action="append", required=True, type=parse_replica, metavar="NAME=URL[,OWNER/SPACE]", |
| help="one per replica; waking needs the Space id") |
| p.add_argument("--host", default="127.0.0.1") |
| p.add_argument("--port", type=int, default=8000) |
| p.add_argument("--rate-limit", type=float, default=20.0, help="requests/min per visitor; 0 = off") |
| p.add_argument("--rate-burst", type=int, default=5) |
| p.add_argument("--trust-proxy", action="store_true", help="visitor address from X-Forwarded-For") |
| p.add_argument("--proxy-hops", type=int, default=1, |
| help="with --trust-proxy: trusted proxies in front (visitor = that many entries from the right)") |
| p.add_argument("--slo-ttft-s", type=float, default=1.0, help="TTFT SLO threshold on /status") |
| p.add_argument("--model-label", default=None, help="shown as the model until a replica has been polled") |
| p.add_argument("--drain-s", type=float, default=30.0) |
| p.add_argument("--log-level", default="info", choices=["debug", "info", "warning", "error"]) |
| p.add_argument("--log-file", default=None, help="also write JSON lines here") |
| p.add_argument("--log-dir", default=None, help="write the shipped log lines as hourly JSONL here (logship.py)") |
| p.add_argument("--log-dataset", default=None, metavar="REPO_ID", help="with --log-dir: upload to this private HF dataset") |
| p.add_argument("--log-replica", default=None, help="name in the log file names; default $SPACE_ID or local") |
| args = p.parse_args(argv) |
| if len({r.name for r in args.replica}) != len(args.replica): |
| p.error("--replica names must be different") |
| if args.log_dataset and not args.log_dir: |
| p.error("--log-dataset needs --log-dir") |
| if args.proxy_hops < 1: |
| p.error("--proxy-hops must be at least 1") |
| token = os.environ.get(REPLICA_TOKEN_ENV) or None |
| if token and not all(r.space_id for r in args.replica): |
| p.error(f"with ${REPLICA_TOKEN_ENV} set, every --replica needs its Space id, so it can be woken") |
|
|
| setup_logging(args.log_level, args.log_file) |
| shipper = None |
| if args.log_dir: |
| from tinyllm.logship import LogShipper |
| shipper = LogShipper(args.log_dir, args.log_dataset, args.log_replica).attach() |
| shipper.start_upload() |
| cfg = RouterConfig(rate_per_min=args.rate_limit, rate_burst=args.rate_burst, trust_proxy=args.trust_proxy, |
| proxy_hops=args.proxy_hops, model_label=args.model_label, drain_s=args.drain_s) |
| app = create_router_app(args.replica, cfg, token, Metrics(slo_ttft_s=args.slo_ttft_s), |
| hf_wake(token) if token else None, shipper) |
|
|
| import uvicorn |
| |
| uvicorn.run(app, host=args.host, port=args.port, log_level="warning", access_log=False, |
| proxy_headers=False, timeout_graceful_shutdown=int(args.drain_s)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
