diff --git a/CLAUDE.md b/CLAUDE.md index 246916a..5df7d7f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -14,7 +14,9 @@ Bare `ai.example.com/v1/...` at the root also still hits the chat router ## Code -- `proxy.py` — chat proxy. Owns the chat router process, generates `models-preset.ini` at startup from `MODELS × CTX_CHOICES`, load/unload via `/models/load` & `/models/unload`, idle-unloads after 10 min. Has chat logging (`logs/chat-.log`) and SSE chunk capture. Also reverse-proxies `/embedding/*` to `embed_proxy.py` on :8003 (strips the `/embedding` prefix; embed lifecycle stays owned by `embed_proxy.py`). +- `proxy.py` — slim entry point (app lifecycle, route registration). Imports handlers from `proxy_request_handlers` and config from `proxy_config`. +- `proxy_config.py` — model metadata, preset generation, CLI config. +- `proxy_request_handlers.py` — HTTP request handlers, streaming, retry/recovery. - `embed_proxy.py` — slimmed-down twin of `proxy.py` for the embedder. Single model, no chat logging, no SSE parsing. Same load/unload pattern. Supports embeddings and re-ranking (`/v1/embeddings`, `/v1/rerank`). - `watchdog.ps1` / `watchdog-embed.ps1` — thin restart-on-crash supervisors. They forward all extra args to the underlying Python script. - `restart-watchdog.ps1` / `restart-watchdog-embed.ps1` — graceful midnight restarts (WM_CLOSE → cascade shutdown → relaunch). Run via Task Scheduler. diff --git a/chat_logger.py b/chat_logger.py index 3ff3904..af0e1d3 100644 --- a/chat_logger.py +++ b/chat_logger.py @@ -179,7 +179,7 @@ async def _flush_all(self) -> None: await self._flush_text() await self._flush_tool_calls() - # ---- rescue helpers ---- + # ── rescue helpers ── @staticmethod def _split_safe_prefix(text: str, markers: tuple[str, ...]) -> tuple[str, str]: @@ -244,7 +244,7 @@ def _build_synthesized_event(self, parsed: dict) -> bytes: body = json.dumps(event, ensure_ascii=False) return f"data: {body}\r\n\r\n".encode() - # ---- main read loop ---- + # ── main read loop ── async def readany(self) -> bytes: # Loop so we only ever return b"" at true EOF — the consumer treats an @@ -255,38 +255,95 @@ async def readany(self) -> bytes: if out is not None: return out + # ── read loop helpers ── + + async def _handle_eof(self) -> bytes | None: + """Handle EOF: flush buffers and return any leftover.""" + await self._flush_all() + if self._buffer: + leftover = self._buffer + self._buffer = b"" + return leftover + return b"" + + @staticmethod + def _find_event_boundary(buffer: bytes): + """Find the earliest SSE event boundary in *buffer*. + + Returns (idx, sep_len) or None if no boundary found. + """ + crlf_idx = buffer.find(b"\r\n\r\n") + lf_idx = buffer.find(b"\n\n") + if crlf_idx == -1 and lf_idx == -1: + return None + if crlf_idx != -1 and (lf_idx == -1 or crlf_idx <= lf_idx): + return crlf_idx, 4 + return lf_idx, 2 + + @staticmethod + def _extract_payload(event_text: str) -> str: + """Extract the payload from an SSE event text block (last data: line wins).""" + payload = "" + for line in event_text.splitlines(): + if line.startswith("data:"): + payload = line[5:].strip() + return payload + + async def _log_delta(self, obj: dict) -> None: + """Process a parsed SSE event for logging side-effects. + + Updates _current_kind, _current_text, _tool_calls, and + _last_chunk_id on the instance. + """ + choices = obj.get("choices") or [] + if not choices: + return + delta = choices[0].get("delta") or {} + reasoning = delta.get("reasoning_content") + content = delta.get("content") + tool_calls = delta.get("tool_calls") + if reasoning: + if self._current_kind != "thinking": + await self._flush_all() + self._current_kind = "thinking" + self._current_text += reasoning + if content: + if self._current_kind != "content": + await self._flush_all() + self._current_kind = "content" + self._current_text += content + if tool_calls: + await self._flush_text() + for tc in tool_calls: + i = tc.get("index", 0) + slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""}) + fn = tc.get("function") or {} + if fn.get("name"): + slot["name"] = fn["name"] + if fn.get("arguments"): + slot["arguments"] += fn["arguments"] + cid = obj.get("id") + if cid: + self._last_chunk_id = cid + async def _readany_once(self) -> bytes | None: data = await self._wrapped.content.readany() if not data: - await self._flush_all() - if self._buffer: - leftover = self._buffer - self._buffer = b"" - return leftover - return b"" + return await self._handle_eof() self._buffer += data outbound: list[bytes] = [] while True: - crlf_idx = self._buffer.find(b"\r\n\r\n") - lf_idx = self._buffer.find(b"\n\n") - if crlf_idx == -1 and lf_idx == -1: + boundary = self._find_event_boundary(self._buffer) + if boundary is None: break - if crlf_idx != -1 and (lf_idx == -1 or crlf_idx <= lf_idx): - idx, sep_len = crlf_idx, 4 - else: - idx, sep_len = lf_idx, 2 + idx, sep_len = boundary raw_event = self._buffer[: idx + sep_len] self._buffer = self._buffer[idx + sep_len:] event_text = raw_event.decode("utf-8", errors="replace").strip() if not event_text: outbound.append(raw_event) continue - # Extract payload line - payload = "" - for line in event_text.splitlines(): - if line.startswith("data:"): - payload = line[5:].strip() - # Non-data lines, comments, [DONE] → pass through + payload = self._extract_payload(event_text) if not payload: outbound.append(raw_event) continue @@ -301,62 +358,35 @@ async def _readany_once(self) -> bytes | None: outbound.append(raw_event) continue # --- (a) existing logging on original delta --- - choices = obj.get("choices") or [] - if choices: - delta = choices[0].get("delta") or {} - reasoning = delta.get("reasoning_content") - content = delta.get("content") - tool_calls = delta.get("tool_calls") - if reasoning: - if self._current_kind != "thinking": - await self._flush_all() - self._current_kind = "thinking" - self._current_text += reasoning - if content: - if self._current_kind != "content": - await self._flush_all() - self._current_kind = "content" - self._current_text += content - if tool_calls: - await self._flush_text() - for tc in tool_calls: - i = tc.get("index", 0) - slot = self._tool_calls.setdefault(i, {"name": "", "arguments": ""}) - fn = tc.get("function") or {} - if fn.get("name"): - slot["name"] = fn["name"] - if fn.get("arguments"): - slot["arguments"] += fn["arguments"] - # Track chunk id for synthesised events - cid = obj.get("id") - if cid: - self._last_chunk_id = cid + await self._log_delta(obj) # --- (b) build outbound bytes with rescue transform --- outbound_event = self._transform_event(obj) outbound.append(outbound_event) return b"".join(outbound) if outbound else None - def _transform_event(self, obj: dict) -> bytes: - """Transform a single parsed event dict into outbound SSE bytes, - applying the rescue state machine.""" - choices = obj.get("choices") or [] - if not choices: - # No choices — pass through - body = json.dumps(obj, ensure_ascii=False) - return f"data: {body}\r\n\r\n".encode() + # ── transform helpers ── - delta = choices[0].get("delta") or {} - reasoning = delta.get("reasoning_content") + def _transform_pass_through(self, obj: dict) -> bytes: + """No-choices pass-through.""" + body = json.dumps(obj, ensure_ascii=False) + return f"data: {body}\r\n\r\n".encode() - # If no reasoning_content, just rewrite finish_reason if needed - if not reasoning: - if self._rescued_any and choices[0].get("finish_reason") == "stop": - choices[0]["finish_reason"] = "tool_calls" - body = json.dumps(obj, ensure_ascii=False) - return f"data: {body}\r\n\r\n".encode() + def _transform_no_reasoning(self, obj: dict) -> bytes: + """No reasoning_content — rewrite finish_reason if needed.""" + choices = obj["choices"] + if self._rescued_any and choices[0].get("finish_reason") == "stop": + choices[0]["finish_reason"] = "tool_calls" + body = json.dumps(obj, ensure_ascii=False) + return f"data: {body}\r\n\r\n".encode() + + def _run_rescue_loop(self, reasoning_text: str): + """Run the rescue state machine on reasoning_content. - # Run rescue state machine on reasoning_content - work = self._reasoning_holdback + reasoning + Returns (prose_parts, synthesized) — prose_parts is a list[str] of + forwarded reasoning prose, synthesized is a list[bytes] of + tool-call SSE events. + """ + work = self._reasoning_holdback + reasoning_text self._reasoning_holdback = "" prose_parts: list[str] = [] synthesized: list[bytes] = [] @@ -409,14 +439,17 @@ def _transform_event(self, obj: dict) -> bytes: self._rescue_buf += work work = "" - # Build the outbound event + return prose_parts, synthesized + + def _build_outbound_event(self, obj: dict, prose_parts: list[str], synthesized: list[bytes]) -> bytes: + """Build the outbound SSE event(s) after rescue processing.""" + choices = obj["choices"] + delta = choices[0]["delta"] forwarded_reasoning = "".join(prose_parts) if forwarded_reasoning: delta["reasoning_content"] = forwarded_reasoning else: delta.pop("reasoning_content", None) - # If delta is now empty and has no other keys, we still emit the event - # (the caller handles skipping if needed) # Rewrite finish_reason if rescued if self._rescued_any and choices[0].get("finish_reason") == "stop": @@ -429,5 +462,21 @@ def _transform_event(self, obj: dict) -> bytes: result += syn return result + def _transform_event(self, obj: dict) -> bytes: + """Transform a single parsed event dict into outbound SSE bytes, + applying the rescue state machine.""" + choices = obj.get("choices") or [] + if not choices: + return self._transform_pass_through(obj) + + delta = choices[0].get("delta") or {} + reasoning = delta.get("reasoning_content") + if not reasoning: + return self._transform_no_reasoning(obj) + + prose_parts, synthesized = self._run_rescue_loop(reasoning) + return self._build_outbound_event(obj, prose_parts, synthesized) + + def __getattr__(self, name: str) -> object: return getattr(self._wrapped, name) diff --git a/proxy.py b/proxy.py index 8799e4f..3cdc8ab 100644 --- a/proxy.py +++ b/proxy.py @@ -1,648 +1,22 @@ +"""Slim entry point — app lifecycle and route registration for the chat proxy.""" from __future__ import annotations -import argparse import asyncio import contextlib -import json import logging -import secrets -import time -from dataclasses import dataclass -from pathlib import Path -from aiohttp import ClientError, ClientSession, ClientTimeout, web - -from log_paths import ( - DATE_FMT, - current_week_dir, - local_now, -) +from aiohttp import ClientSession, ClientTimeout, web from proxy_base import ( - API_KEY, ClientGone, DeadWorkerError, ForwardedAccessLogger, ProxyConfig, - auth_middleware, client_ip, configure_logging, filter_request_headers, - filter_response_headers, health_handler, idle_watchdog, - _is_dead_worker_response, + ForwardedAccessLogger, ProxyConfig, + auth_middleware, configure_logging, health_handler, idle_watchdog, ) from router_manager import ChatRouterManager -from chat_logger import ChatLogger, SSEChunkLogger - - -ROOT = Path(__file__).resolve().parent -SERVER_EXE = ROOT / "llama.cpp_latest" / "llama-server.exe" -PRESET_PATH = ROOT / "models-preset.ini" - -# Each (model, ctx) pair becomes its own preset section so the router -# exposes them as distinct models that the pi-llama-cpp extension can -# discover and switch between — pi treats ctx changes as model switches, -# which forces a router reload with the new ctx-size baked in. -@dataclass(frozen=True) -class ModelChoice: - label: str # menu display - base_id: str # e.g. "qwen3.6-35b-q3" — ctx suffix appended at render time - model_file: Path # GGUF weights - mmproj_file: Path # multimodal projector - spec_mtp: bool = False # enable built-in MTP speculative decoding (draft-mtp) - - def preset_id(self, ctx: int) -> str: - return f"{self.base_id}-{ctx // 1024}k" - - -MODELS: list[ModelChoice] = [ - ModelChoice( - "Qwen3.6-35B-A3B Q3", - "qwen3.6-35b-q3", - ROOT / "models" / "Qwen3.6-35B-A3B-UD-Q3_K_XL.gguf", - ROOT / "models" / "_aux" / "mmproj-F16.gguf", - ), - ModelChoice( - "Qwen3.6-35B-A3B Q4", - "qwen3.6-35b-q4", - ROOT / "models" / "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", - ROOT / "models" / "_aux" / "mmproj-F16.gguf", - ), - ModelChoice( - "Qwen3.6-27B Q4", - "qwen3.6-27b-q4", - ROOT / "models" / "Qwen3.6-27B-UD-Q4_K_XL.gguf", - ROOT / "models" / "_aux" / "mmproj-27b-BF16.gguf", - ), - ModelChoice( - "Qwen3.6-27B Q4 MTP", - "qwen3.6-27b-q4-mtp", - ROOT / "models" / "Qwen3.6-27B-UD-Q4_K_XL.mtp.gguf", - ROOT / "models" / "_aux" / "mmproj-27b-BF16.gguf", - spec_mtp=True, - ), -] - -CTX_CHOICES: list[int] = [32768, 65536, 98304, 131072] - -PROXY_HOST = "0.0.0.0" -PROXY_PORT = 8001 -SERVER_HOST = "127.0.0.1" -SERVER_PORT = 8002 -EMBED_PROXY_HOST = "::1" -EMBED_PROXY_PORT = 8003 -IDLE_TIMEOUT = 600 # 10 minutes of inactivity before unload -IDLE_CHECK_INTERVAL = 30 # check every 30s -HEALTH_POLL_INTERVAL = 1.0 -BOOT_TIMEOUT = 60 -RETRY_AFTER_SECONDS = 30 - - -class ChatProxyConfig(ProxyConfig): - @property - def server_command(self) -> list[str]: - log_file = current_week_dir(ROOT / "logs") / f"llama-server-{local_now().strftime(DATE_FMT)}.log" - return [ - str(SERVER_EXE), - "--log-file", str(log_file), - "--log-timestamps", - "--log-prefix", - "--models-preset", str(PRESET_PATH), - "--models-max", "1", - "--no-models-autoload", - # --- perf A/B test: chat-template flags (see chat_template_perf_test.md) --- - # Variant D: full new config (jinja + custom template + preserve_thinking kwarg) - "--jinja", - "--chat-template-file", str(ROOT / "chat_template.jinja"), - "--chat-template-kwargs", '{"preserve_thinking":true}', - # ------------------------------------------------------------------------ - "--host", self.server_host, - "--port", str(self.server_port), - "--api-key", self.api_key, - ] - - -DEFAULT_MODEL: ModelChoice = MODELS[0] -DEFAULT_CTX: int = CTX_CHOICES[-1] - - -def pick_setup(model_arg: str | None, ctx_arg: int | None) -> tuple[ModelChoice, int]: - """Resolve (model, context) fallback defaults from CLI args. - - All (model × ctx) combos are exposed as router presets regardless; this - only picks which one to use when a client doesn't specify a model. - """ - model = DEFAULT_MODEL - if model_arg: - match = next( - (m for m in MODELS if m.label == model_arg or m.model_file.stem == model_arg), - None, - ) - if match is None: - labels = ", ".join(m.label for m in MODELS) - raise SystemExit(f"Unknown model {model_arg!r}; available: {labels}") - model = match - - ctx = ctx_arg if ctx_arg is not None else DEFAULT_CTX - if ctx_arg is not None and ctx_arg not in CTX_CHOICES: - logging.warning("--ctx-size %d is outside preset choices %s", ctx_arg, CTX_CHOICES) - - for path in (model.model_file, model.mmproj_file): - if not path.exists(): - raise SystemExit(f"Missing file for {model.label}: {path}") - return model, ctx - - -def _model_preset_section(model: ModelChoice, ctx: int) -> str: - """Return the INI section text for a single (model, ctx) pair.""" - spec = ( - f"spec-type = draft-mtp\n" - f"spec-draft-n-max = 2\n" - f"spec-draft-p-min = 0.0\n" - if model.spec_mtp - else "" - ) - return ( - f"[{model.preset_id(ctx)}]\n" - f"model = {model.model_file.as_posix()}\n" - f"mmproj = {model.mmproj_file.as_posix()}\n" - f"ctx-size = {ctx}\n" - f"n-gpu-layers = 999\n" - f"flash-attn = on\n" - f"cache-type-k = q4_0\n" - f"cache-type-v = q4_0\n" - f"no-mmap = 1\n" - f"parallel = 1\n" - f"jinja = 1\n" - f"temp = 0.6\n" - f"top-p = 0.95\n" - f"top-k = 20\n" - + spec - ) - - -def write_preset(models: list[ModelChoice], ctx_choices: list[int]) -> None: - """Generate models-preset.ini with every (model, ctx) combination. - - Each combo becomes a distinct preset id (e.g. `qwen3.6-35b-q3-128k`), - so picking a different ctx in pi triggers a router reload with the new - context size — the only way to "change ctx" without restarting the proxy. - """ - sections = [ - _model_preset_section(m, ctx) - for m in models - for ctx in ctx_choices - ] - content = "\n".join(sections) + "\n" - PRESET_PATH.write_text(content, encoding="utf-8") - - -def build_config() -> ProxyConfig: - p = argparse.ArgumentParser(description="Router-mode proxy for llama-server") - p.add_argument("--proxy-host", default=PROXY_HOST) - p.add_argument("--proxy-port", type=int, default=PROXY_PORT) - p.add_argument("--server-host", default=SERVER_HOST) - p.add_argument("--server-port", type=int, default=SERVER_PORT) - p.add_argument("--idle-timeout", type=int, default=IDLE_TIMEOUT) - p.add_argument("--idle-check-interval", type=int, default=IDLE_CHECK_INTERVAL) - p.add_argument("--health-poll-interval", type=float, default=HEALTH_POLL_INTERVAL) - p.add_argument("--boot-timeout", type=int, default=BOOT_TIMEOUT) - p.add_argument("--model", default=None, - help="Skip the model picker (use exact label from MODELS)") - p.add_argument("--ctx-size", type=int, default=None, - help="Skip the context picker (any int; menu offers 32k/64k/128k)") - p.add_argument("--api-key", default=API_KEY) - p.add_argument("--embed-host", default=EMBED_PROXY_HOST) - p.add_argument("--embed-port", type=int, default=EMBED_PROXY_PORT) - p.add_argument("--no-chat-log", action="store_true") - args = p.parse_args() - - if not SERVER_EXE.exists(): - raise SystemExit(f"Missing required file: {SERVER_EXE}") - - model, ctx = pick_setup(args.model, args.ctx_size) - write_preset(MODELS, CTX_CHOICES) - default_id = model.preset_id(ctx) - print(f"Default: {model.label} @ {ctx // 1024}k ctx (id: {default_id})") - print(f"Exposed presets: {len(MODELS) * len(CTX_CHOICES)} (one per model×ctx combo)") - - return ChatProxyConfig( - proxy_host=args.proxy_host, - proxy_port=args.proxy_port, - server_host=args.server_host, - server_port=args.server_port, - idle_timeout=args.idle_timeout, - idle_check_interval=args.idle_check_interval, - health_poll_interval=args.health_poll_interval, - boot_timeout=args.boot_timeout, - default_model=default_id, - api_key=args.api_key, - embed_host=args.embed_host, - embed_port=args.embed_port, - chat_log=not args.no_chat_log, - ) - - -def _strip_chat_prefix(path: str) -> str: - """Strip the /chat alias prefix so backend sees plain /v1/... paths.""" - if path == "/chat": - return "/" - if path.startswith("/chat/"): - return path[len("/chat"):] - return path - - -def _inject_cache_prompt(body: bytes | None, method: str, path: str) -> bytes | None: - if method != "POST" or path.rstrip("/") != "/v1/chat/completions" or not body: - return body - try: - payload = json.loads(body) - if not payload.get("cache_prompt"): - payload["cache_prompt"] = True - return json.dumps(payload, separators=(",", ":")).encode() - except (ValueError, TypeError): - pass - return body - - -def _model_from_body(body: bytes | None, fallback: str) -> str: - if not body: - return fallback - try: - payload = json.loads(body) - m = payload.get("model") - if isinstance(m, str) and m: - return m - except (ValueError, TypeError): - pass - return fallback - - -async def proxy_request(request: web.Request) -> web.StreamResponse: - manager: ChatRouterManager = request.app["manager"] - session: ClientSession = request.app["session"] - chat_logger: ChatLogger | None = request.app.get("chat_logger") - - req_id = secrets.token_hex(4) - started = time.monotonic() - manager.begin_request() - - effective_path = _strip_chat_prefix(request.path) - - # Determine whether this request requires a model and parse the body up front. - # Body parsing errors (OOM, malformed) become 503 so clients get Retry-After. - try: - body = await request.read() if request.can_read_body else None - body = _inject_cache_prompt(body, request.method, effective_path) - except Exception as exc: - logging.exception("[req=%s] body read failed", req_id) - manager.end_request() - return web.json_response( - {"error": "backend unavailable", "detail": str(exc)}, - status=503, - headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, - ) - - path = effective_path.rstrip("/") - needs_model = path in ( - "/v1/chat/completions", "/v1/completions", "/v1/embeddings", - "/chat/completions", "/completions", "/embeddings", - ) - model: str | None = None - if needs_model: - model = _model_from_body(body, manager.config.default_model) - - if chat_logger is not None: - await chat_logger.log_request(request.method, request.path, body, req_id) - - query = request.rel_url.query_string - target_url = f"{manager.config.backend_base_url}{effective_path}" - if query: - target_url = f"{target_url}?{query}" - headers = filter_request_headers(request.headers, manager.config.api_key) - - async def _do_forward() -> web.StreamResponse: - """Inner forward/stream; called from within use_model or directly. - - For non-SSE responses with status >= 500 the body is pre-read so we - can detect a dead-worker error BEFORE committing the response to the - client via ``downstream.prepare()``. If a dead-worker marker is found - a ``DeadWorkerError`` is raised for the caller to handle. All other - paths (success, 4xx, SSE) stream zero-copy as before. - """ - try: - upstream_resp = await session.request( - request.method, target_url, headers=headers, - data=body, allow_redirects=False, timeout=None, - ) - is_sse = "text/event-stream" in upstream_resp.headers.get("Content-Type", "") - - # Pre-read non-SSE error bodies BEFORE prepare() so we can inspect - # them and raise DeadWorkerError without having committed headers. - if upstream_resp.status >= 500 and not is_sse: - error_body = await upstream_resp.content.read() - if _is_dead_worker_response(upstream_resp.status, error_body): - raise DeadWorkerError(upstream_resp.status, error_body) - # Real (non-dead-worker) 5xx — forward as-is. - response_headers = filter_response_headers(upstream_resp.headers) - response_headers["X-Request-ID"] = req_id - downstream = web.StreamResponse( - status=upstream_resp.status, reason=upstream_resp.reason, - headers=response_headers, - ) - try: - await downstream.prepare(request) - except ConnectionResetError: - raise ClientGone from None - try: - await downstream.write(error_body) - except ConnectionResetError: - logging.info("[req=%s] client disconnected", req_id) - finally: - with contextlib.suppress(ConnectionResetError, RuntimeError): - await downstream.write_eof() - duration_ms = int((time.monotonic() - started) * 1000) - logging.info( - "[req=%s] %s %s %s -> %s in %dms", - req_id, client_ip(request), request.method, request.rel_url, - downstream.status, duration_ms, - ) - return downstream - - response_headers = filter_response_headers(upstream_resp.headers) - response_headers["X-Request-ID"] = req_id - downstream = web.StreamResponse( - status=upstream_resp.status, reason=upstream_resp.reason, headers=response_headers, - ) - try: - await downstream.prepare(request) - except ConnectionResetError: - raise ClientGone from None - try: - if is_sse and chat_logger: - wrapped = SSEChunkLogger(upstream_resp, chat_logger) - while True: - try: - chunk = await asyncio.wait_for(wrapped.readany(), timeout=25) - except asyncio.TimeoutError: - await downstream.write(b": keep-alive\n\n") - continue - if not chunk: - break - await downstream.write(chunk) - elif is_sse: - while True: - try: - chunk = await asyncio.wait_for(upstream_resp.content.readany(), timeout=25) - except asyncio.TimeoutError: - await downstream.write(b": keep-alive\n\n") - continue - if not chunk: - break - await downstream.write(chunk) - else: - async for chunk in upstream_resp.content.iter_any(): - await downstream.write(chunk) - except ConnectionResetError: - logging.info("[req=%s] client disconnected", req_id) - # Closing upstream is the correct way to signal cancellation, but it - # does NOT prevent the worker crash on its own: llama-server has an - # unfixed cancel→next-request desync (ggml-org/llama.cpp#20921) that - # wedges/crashes the worker — worst with reasoning + speculative on a - # large-context (cancel-during-prefill) request. We can't make the - # cancel safe here, so we flag the worker suspect and let the NEXT - # model request probe it first (see guard_after_cancel). - try: - upstream_resp.close() - except Exception: - pass - if needs_model: - manager.mark_worker_suspect() - finally: - with contextlib.suppress(ConnectionResetError, RuntimeError): - await downstream.write_eof() - # Always close the upstream response to free the router's connection. - # Prevents zombie connections and worker child crashes on disconnect. - # close() is idempotent — safe if already closed in the except block. - try: - upstream_resp.close() - except Exception: - pass - duration_ms = int((time.monotonic() - started) * 1000) - logging.info( - "[req=%s] %s %s %s -> %s in %dms", - req_id, client_ip(request), request.method, request.rel_url, - downstream.status, duration_ms, - ) - return downstream - except DeadWorkerError: - raise # propagate to retry loop - except ClientGone: - # Downstream client (cloudflared / end client) hung up before we finished - # sending the response — raised when downstream.prepare() hit a closing - # transport. Same benign disconnect the write loops already handle, just - # earlier. Not a proxy fault, so log calmly and return 499. - logging.info("[req=%s] client disconnected before response sent", req_id) - # Close the upstream response to release the router's connection. The - # worker had already begun generating (it produced response headers), so - # this is a mid-generation abort too — flag it for the post-cancel guard. - try: - upstream_resp.close() - except Exception: - pass - if needs_model: - manager.mark_worker_suspect() - return web.Response( - status=499, reason="Client Closed Request", - headers={"X-Request-ID": req_id}, - ) - except ClientError as exc: - # Genuine upstream failure (router connection reset/refused, etc.) — a - # real bad gateway. Note: an *upstream* ClientConnectionResetError lands - # here, while a *downstream* one was already converted to ClientGone above. - logging.exception("[req=%s] proxy failure", req_id) - return web.json_response( - {"error": "bad gateway", "detail": str(exc)}, - status=502, headers={"X-Request-ID": req_id}, - ) - - try: - if needs_model: - # Post-cancel guard: if a prior request was aborted mid-generation, - # probe (and recover if needed) the worker before this request hits it, - # avoiding the upstream cancel→next-request crash. No-op when not suspect. - # Runs outside use_model so its recover_worker() doesn't nest _load_lock. - await manager.guard_after_cancel(model) - # Wrap the entire forward inside use_model so _end_forward() only - # fires after the stream is fully drained — no eviction mid-stream. - # Up to 2 attempts: on DeadWorkerError, recover then retry once. - max_attempts = 2 - for attempt in range(1, max_attempts + 1): - try: - async with manager.use_model(model): - return await _do_forward() - except DeadWorkerError as exc: - dead_detected_at = time.monotonic() - logging.warning( - "[req=%s] dead-worker 500 on attempt %d/%d — body: %r", - req_id, attempt, max_attempts, exc.body[:200], - ) - if attempt < max_attempts: - recovered = await manager.recover_worker(model, dead_detected_at) - if not recovered: - logging.error("[req=%s] worker recovery failed — giving up", req_id) - return web.json_response( - {"error": "backend unavailable", "detail": "worker recovery failed"}, - status=503, - headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, - ) - logging.info("[req=%s] worker recovered — retrying request", req_id) - else: - logging.error("[req=%s] dead worker persists after recovery — 503", req_id) - return web.json_response( - {"error": "backend unavailable", "detail": "dead worker after recovery"}, - status=503, - headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, - ) - except Exception as exc: - logging.exception("[req=%s] backend unavailable", req_id) - return web.json_response( - {"error": "backend unavailable", "detail": str(exc)}, - status=503, - headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, - ) - else: - # Non-model endpoints (health, /v1/models, props, …) forward directly. - return await _do_forward() - finally: - manager.end_request() - - -async def models_handler(request: web.Request) -> web.Response: - """Normalize the router's /models payload so clients see a sane status. - - llama.cpp's router keeps `status.failed = true` (with a stale `exit_code`) - on presets that have never had a successful load in the current process — - a residual diagnostic flag rather than a real "this model is broken" - signal. Newer pi-llama-cpp versions short-circuit to FAILED when they - see `failed: true`, so a freshly-booted router shows every preset as - "Retry" in pi. We rewrite the flag to false whenever `value` says the - preset is simply unloaded — `value` is the source of truth. - """ - manager: ChatRouterManager = request.app["manager"] - session: ClientSession = request.app["session"] - effective_path = _strip_chat_prefix(request.path) - query = request.rel_url.query_string - target_url = f"{manager.config.backend_base_url}{effective_path}" - if query: - target_url = f"{target_url}?{query}" - headers = filter_request_headers(request.headers, manager.config.api_key) - async with session.get(target_url, headers=headers) as upstream: - body_text = await upstream.text() - try: - payload = json.loads(body_text) - except (ValueError, TypeError): - return web.Response(status=upstream.status, body=body_text, - content_type=upstream.content_type or "application/json") - for entry in payload.get("data") or []: - status = entry.get("status") - if isinstance(status, dict) and status.get("value") == "unloaded": - status["failed"] = False - status.pop("exit_code", None) - return web.json_response(payload, status=upstream.status) - - -async def props_handler(request: web.Request) -> web.Response: - """Normalize the router's /props response for unloaded models. - - llama.cpp returns HTTP 400 with `{"error":{"code":400,"message":"model is - not loaded",...}}` when /props is asked about a non-loaded preset. The - pi-llama-cpp extension probes /props as a sanity check after seeing a - model in /models with status "unloaded", and a strict reading of its - parser can mis-classify that 400 response as FAILED (shows "Retry" - instead of "Load & switch"). We rewrite to a clean 200 JSON whose - shape matches the exact equality checks in baseModel.getStatus(). - """ - manager: ChatRouterManager = request.app["manager"] - session: ClientSession = request.app["session"] - effective_path = _strip_chat_prefix(request.path) - query = request.rel_url.query_string - target_url = f"{manager.config.backend_base_url}{effective_path}" - if query: - target_url = f"{target_url}?{query}" - headers = filter_request_headers(request.headers, manager.config.api_key) - async with session.get(target_url, headers=headers) as upstream: - body_text = await upstream.text() - if upstream.status == 400 and "model is not loaded" in body_text: - return web.json_response( - {"error": {"code": 400, "message": "model is not loaded"}} - ) - return web.Response( - status=upstream.status, - body=body_text, - content_type=upstream.content_type or "application/json", - ) - - -async def embed_forward(request: web.Request) -> web.StreamResponse: - """Reverse-proxy /embedding/* to the standalone embed_proxy on :8003. - - Strips the prefix so embed_proxy sees plain OpenAI-style paths - (/v1/embeddings, /v1/models, /health, ...). embed_proxy owns its own - router, load/unload, and idle timer — this is a dumb HTTP forwarder. - """ - session: ClientSession = request.app["session"] - config: ProxyConfig = request.app["config"] - req_id = secrets.token_hex(4) - started = time.monotonic() - - tail = request.match_info.get("tail", "") - sub_path = "/" + tail if tail else "/" - query = request.rel_url.query_string - target_url = f"{config.embed_base_url}{sub_path}" - if query: - target_url = f"{target_url}?{query}" - - body = await request.read() if request.can_read_body else None - headers = filter_request_headers(request.headers, config.api_key) - - try: - upstream_resp = await session.request( - request.method, target_url, headers=headers, - data=body, allow_redirects=False, timeout=None, - ) - response_headers = filter_response_headers(upstream_resp.headers) - response_headers["X-Request-ID"] = req_id - downstream = web.StreamResponse( - status=upstream_resp.status, reason=upstream_resp.reason, headers=response_headers, - ) - await downstream.prepare(request) - try: - is_sse = "text/event-stream" in upstream_resp.headers.get("Content-Type", "") - if is_sse: - while True: - try: - chunk = await asyncio.wait_for(upstream_resp.content.readany(), timeout=25) - except asyncio.TimeoutError: - await downstream.write(b": keep-alive\n\n") - continue - if not chunk: - break - await downstream.write(chunk) - else: - async for chunk in upstream_resp.content.iter_any(): - await downstream.write(chunk) - except ConnectionResetError: - logging.info("[embed req=%s] client disconnected", req_id) - finally: - with contextlib.suppress(ConnectionResetError, RuntimeError): - await downstream.write_eof() - duration_ms = int((time.monotonic() - started) * 1000) - logging.info( - "[embed req=%s] %s %s %s -> %s in %dms", - req_id, client_ip(request), request.method, request.path, - downstream.status, duration_ms, - ) - return downstream - except ClientError as exc: - logging.exception("[embed req=%s] forward failure", req_id) - return web.json_response( - {"error": "bad gateway", "detail": str(exc)}, - status=502, headers={"X-Request-ID": req_id}, - ) +from chat_logger import ChatLogger +from proxy_request_handlers import ( + embed_forward, models_handler, props_handler, proxy_request, +) +from proxy_config import ROOT, build_config async def lifecycle_context(app: web.Application): diff --git a/proxy_config.py b/proxy_config.py new file mode 100644 index 0000000..4f2b901 --- /dev/null +++ b/proxy_config.py @@ -0,0 +1,228 @@ +"""Model metadata, preset generation, and CLI configuration for the chat proxy.""" +from __future__ import annotations + +import argparse +import logging +from dataclasses import dataclass +from pathlib import Path + +from log_paths import ( + DATE_FMT, + current_week_dir, + local_now, +) + +from proxy_base import API_KEY, ProxyConfig + + +ROOT = Path(__file__).resolve().parent +SERVER_EXE = ROOT / "llama.cpp_latest" / "llama-server.exe" +PRESET_PATH = ROOT / "models-preset.ini" + +# Each (model, ctx) pair becomes its own preset section so the router +# exposes them as distinct models that the pi-llama-cpp extension can +# discover and switch between — pi treats ctx changes as model switches, +# which forces a router reload with the new ctx-size baked in. +@dataclass(frozen=True) +class ModelChoice: + label: str # menu display + base_id: str # e.g. "qwen3.6-35b-q3" — ctx suffix appended at render time + model_file: Path # GGUF weights + mmproj_file: Path # multimodal projector + spec_mtp: bool = False # enable built-in MTP speculative decoding (draft-mtp) + + def preset_id(self, ctx: int) -> str: + return f"{self.base_id}-{ctx // 1024}k" + + +MODELS: list[ModelChoice] = [ + ModelChoice( + "Qwen3.6-35B-A3B Q3", + "qwen3.6-35b-q3", + ROOT / "models" / "Qwen3.6-35B-A3B-UD-Q3_K_XL.gguf", + ROOT / "models" / "_aux" / "mmproj-F16.gguf", + ), + ModelChoice( + "Qwen3.6-35B-A3B Q4", + "qwen3.6-35b-q4", + ROOT / "models" / "Qwen3.6-35B-A3B-UD-Q4_K_M.gguf", + ROOT / "models" / "_aux" / "mmproj-F16.gguf", + ), + ModelChoice( + "Qwen3.6-27B Q4", + "qwen3.6-27b-q4", + ROOT / "models" / "Qwen3.6-27B-UD-Q4_K_XL.gguf", + ROOT / "models" / "_aux" / "mmproj-27b-BF16.gguf", + ), + ModelChoice( + "Qwen3.6-27B Q4 MTP", + "qwen3.6-27b-q4-mtp", + ROOT / "models" / "Qwen3.6-27B-UD-Q4_K_XL.mtp.gguf", + ROOT / "models" / "_aux" / "mmproj-27b-BF16.gguf", + spec_mtp=True, + ), +] + +CTX_CHOICES: list[int] = [32768, 65536, 98304, 131072] + +PROXY_HOST = "0.0.0.0" +PROXY_PORT = 8001 +SERVER_HOST = "127.0.0.1" +SERVER_PORT = 8002 +EMBED_PROXY_HOST = "::1" +EMBED_PROXY_PORT = 8003 +IDLE_TIMEOUT = 600 # 10 minutes of inactivity before unload +IDLE_CHECK_INTERVAL = 30 # check every 30s +HEALTH_POLL_INTERVAL = 1.0 +BOOT_TIMEOUT = 60 +RETRY_AFTER_SECONDS = 30 + + +class ChatProxyConfig(ProxyConfig): + @property + def server_command(self) -> list[str]: + log_file = current_week_dir(ROOT / "logs") / f"llama-server-{local_now().strftime(DATE_FMT)}.log" + return [ + str(SERVER_EXE), + "--log-file", str(log_file), + "--log-timestamps", + "--log-prefix", + "--models-preset", str(PRESET_PATH), + "--models-max", "1", + "--no-models-autoload", + # --- perf A/B test: chat-template flags (see chat_template_perf_test.md) --- + # Variant D: full new config (jinja + custom template + preserve_thinking kwarg) + "--jinja", + "--chat-template-file", str(ROOT / "chat_template.jinja"), + "--chat-template-kwargs", '{"preserve_thinking":true}', + # ------------------------------------------------------------------------ + "--host", self.server_host, + "--port", str(self.server_port), + "--api-key", self.api_key, + ] + + +DEFAULT_MODEL: ModelChoice = MODELS[0] +DEFAULT_CTX: int = CTX_CHOICES[-1] + + +def pick_setup(model_arg: str | None, ctx_arg: int | None) -> tuple[ModelChoice, int]: + """Resolve (model, context) fallback defaults from CLI args. + + All (model × ctx) combos are exposed as router presets regardless; this + only picks which one to use when a client doesn't specify a model. + """ + model = DEFAULT_MODEL + if model_arg: + match = next( + (m for m in MODELS if m.label == model_arg or m.model_file.stem == model_arg), + None, + ) + if match is None: + labels = ", ".join(m.label for m in MODELS) + raise SystemExit(f"Unknown model {model_arg!r}; available: {labels}") + model = match + + ctx = ctx_arg if ctx_arg is not None else DEFAULT_CTX + if ctx_arg is not None and ctx_arg not in CTX_CHOICES: + logging.warning("--ctx-size %d is outside preset choices %s", ctx_arg, CTX_CHOICES) + + for path in (model.model_file, model.mmproj_file): + if not path.exists(): + raise SystemExit(f"Missing file for {model.label}: {path}") + return model, ctx + + +def _model_preset_section(model: ModelChoice, ctx: int) -> str: + """Return the INI section text for a single (model, ctx) pair.""" + spec = ( + f"spec-type = draft-mtp\n" + f"spec-draft-n-max = 2\n" + f"spec-draft-p-min = 0.0\n" + if model.spec_mtp + else "" + ) + return ( + f"[{model.preset_id(ctx)}]\n" + f"model = {model.model_file.as_posix()}\n" + f"mmproj = {model.mmproj_file.as_posix()}\n" + f"ctx-size = {ctx}\n" + f"n-gpu-layers = 999\n" + f"flash-attn = on\n" + f"cache-type-k = q4_0\n" + f"cache-type-v = q4_0\n" + f"no-mmap = 1\n" + f"parallel = 1\n" + f"jinja = 1\n" + f"temp = 0.6\n" + f"top-p = 0.95\n" + f"top-k = 20\n" + + spec + ) + + +def write_preset(models: list[ModelChoice], ctx_choices: list[int]) -> None: + """Generate models-preset.ini with every (model, ctx) combination. + + Each combo becomes a distinct preset id (e.g. `qwen3.6-35b-q3-128k`), + so picking a different ctx in pi triggers a router reload with the new + context size — the only way to "change ctx" without restarting the proxy. + """ + sections = [ + _model_preset_section(m, ctx) + for m in models + for ctx in ctx_choices + ] + content = "\n".join(sections) + "\n" + PRESET_PATH.write_text(content, encoding="utf-8") + + +def _build_arg_parser() -> argparse.ArgumentParser: + """Construct and return the CLI argument parser for the chat proxy.""" + p = argparse.ArgumentParser(description="Router-mode proxy for llama-server") + p.add_argument("--proxy-host", default=PROXY_HOST) + p.add_argument("--proxy-port", type=int, default=PROXY_PORT) + p.add_argument("--server-host", default=SERVER_HOST) + p.add_argument("--server-port", type=int, default=SERVER_PORT) + p.add_argument("--idle-timeout", type=int, default=IDLE_TIMEOUT) + p.add_argument("--idle-check-interval", type=int, default=IDLE_CHECK_INTERVAL) + p.add_argument("--health-poll-interval", type=float, default=HEALTH_POLL_INTERVAL) + p.add_argument("--boot-timeout", type=int, default=BOOT_TIMEOUT) + p.add_argument("--model", default=None, + help="Skip the model picker (use exact label from MODELS)") + p.add_argument("--ctx-size", type=int, default=None, + help="Skip the context picker (any int; menu offers 32k/64k/128k)") + p.add_argument("--api-key", default=API_KEY) + p.add_argument("--embed-host", default=EMBED_PROXY_HOST) + p.add_argument("--embed-port", type=int, default=EMBED_PROXY_PORT) + p.add_argument("--no-chat-log", action="store_true") + return p + + +def build_config() -> ProxyConfig: + args = _build_arg_parser().parse_args() + + if not SERVER_EXE.exists(): + raise SystemExit(f"Missing required file: {SERVER_EXE}") + + model, ctx = pick_setup(args.model, args.ctx_size) + write_preset(MODELS, CTX_CHOICES) + default_id = model.preset_id(ctx) + print(f"Default: {model.label} @ {ctx // 1024}k ctx (id: {default_id})") + print(f"Exposed presets: {len(MODELS) * len(CTX_CHOICES)} (one per model×ctx combo)") + + return ChatProxyConfig( + proxy_host=args.proxy_host, + proxy_port=args.proxy_port, + server_host=args.server_host, + server_port=args.server_port, + idle_timeout=args.idle_timeout, + idle_check_interval=args.idle_check_interval, + health_poll_interval=args.health_poll_interval, + boot_timeout=args.boot_timeout, + default_model=default_id, + api_key=args.api_key, + embed_host=args.embed_host, + embed_port=args.embed_port, + chat_log=not args.no_chat_log, + ) diff --git a/proxy_request_handlers.py b/proxy_request_handlers.py new file mode 100644 index 0000000..37a451c --- /dev/null +++ b/proxy_request_handlers.py @@ -0,0 +1,459 @@ +"""HTTP request handlers, streaming, and retry/recovery for the chat proxy.""" +from __future__ import annotations + +import asyncio +import contextlib +import json +import logging +import secrets +import time + +from aiohttp import ClientError, ClientSession, ClientTimeout, web + +from proxy_base import ( + ClientGone, DeadWorkerError, ProxyConfig, + client_ip, filter_request_headers, filter_response_headers, + _is_dead_worker_response, +) +from router_manager import ChatRouterManager +from chat_logger import ChatLogger, SSEChunkLogger +from proxy_config import RETRY_AFTER_SECONDS + + +def _strip_chat_prefix(path: str) -> str: + """Strip the /chat alias prefix so backend sees plain /v1/... paths.""" + if path == "/chat": + return "/" + if path.startswith("/chat/"): + return path[len("/chat"):] + return path + + +def _inject_cache_prompt(body: bytes | None, method: str, path: str) -> bytes | None: + if method != "POST" or path.rstrip("/") != "/v1/chat/completions" or not body: + return body + try: + payload = json.loads(body) + if not payload.get("cache_prompt"): + payload["cache_prompt"] = True + return json.dumps(payload, separators=(",", ":")).encode() + except (ValueError, TypeError): + pass + return body + + +def _model_from_body(body: bytes | None, fallback: str) -> str: + if not body: + return fallback + try: + payload = json.loads(body) + m = payload.get("model") + if isinstance(m, str) and m: + return m + except (ValueError, TypeError): + pass + return fallback + + +async def _retry_with_recovery( + manager: ChatRouterManager, model: str, req_id: str, started: float, + forward_fn, +) -> web.StreamResponse: + """Retry loop with DeadWorkerError recovery (max 2 attempts).""" + await manager.guard_after_cancel(model) + max_attempts = 2 + for attempt in range(1, max_attempts + 1): + try: + async with manager.use_model(model): + return await forward_fn() + except DeadWorkerError as exc: + dead_detected_at = time.monotonic() + logging.warning( + "[req=%s] dead-worker 500 on attempt %d/%d — body: %r", + req_id, attempt, max_attempts, exc.body[:200], + ) + if attempt < max_attempts: + recovered = await manager.recover_worker(model, dead_detected_at) + if not recovered: + logging.error("[req=%s] worker recovery failed — giving up", req_id) + return web.json_response( + {"error": "backend unavailable", "detail": "worker recovery failed"}, + status=503, + headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, + ) + logging.info("[req=%s] worker recovered — retrying request", req_id) + else: + logging.error("[req=%s] dead worker persists after recovery — 503", req_id) + return web.json_response( + {"error": "backend unavailable", "detail": "dead worker after recovery"}, + status=503, + headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, + ) + except Exception as exc: + logging.exception("[req=%s] backend unavailable", req_id) + return web.json_response( + {"error": "backend unavailable", "detail": str(exc)}, + status=503, + headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, + ) + + +async def _stream_response( + upstream_resp, request: web.Request, is_sse: bool, chat_logger, + req_id: str, started: float, needs_model: bool, manager: ChatRouterManager, +) -> web.StreamResponse: + """Prepare downstream, stream content (SSE/non-SSE), handle disconnect, log.""" + response_headers = filter_response_headers(upstream_resp.headers) + response_headers["X-Request-ID"] = req_id + downstream = web.StreamResponse( + status=upstream_resp.status, reason=upstream_resp.reason, headers=response_headers, + ) + try: + await downstream.prepare(request) + except ConnectionResetError: + raise ClientGone from None + try: + if is_sse and chat_logger: + wrapped = SSEChunkLogger(upstream_resp, chat_logger) + while True: + try: + chunk = await asyncio.wait_for(wrapped.readany(), timeout=25) + except asyncio.TimeoutError: + await downstream.write(b": keep-alive\n\n") + continue + if not chunk: + break + await downstream.write(chunk) + elif is_sse: + while True: + try: + chunk = await asyncio.wait_for(upstream_resp.content.readany(), timeout=25) + except asyncio.TimeoutError: + await downstream.write(b": keep-alive\n\n") + continue + if not chunk: + break + await downstream.write(chunk) + else: + async for chunk in upstream_resp.content.iter_any(): + await downstream.write(chunk) + except ConnectionResetError: + logging.info("[req=%s] client disconnected", req_id) + # Mid-stream disconnect on a model request: the router may still be + # generating. Without guard_after_cancel the next request to this + # worker can crash (ggml-org/llama.cpp#20921). Mark suspect so the + # retry loop probes/recovers before the next model request. + if needs_model: + manager.mark_worker_suspect() + finally: + with contextlib.suppress(ConnectionResetError, RuntimeError): + await downstream.write_eof() + # Always close the upstream response to free the router's connection. + # Prevents zombie connections and worker child crashes on disconnect. + # close() is idempotent — safe if already closed in the except block. + try: + upstream_resp.close() + except Exception: + pass + duration_ms = int((time.monotonic() - started) * 1000) + logging.info( + "[req=%s] %s %s %s -> %s in %dms", + req_id, client_ip(request), request.method, request.rel_url, + downstream.status, duration_ms, + ) + return downstream + + +async def proxy_request(request: web.Request) -> web.StreamResponse: + manager: ChatRouterManager = request.app["manager"] + session: ClientSession = request.app["session"] + chat_logger: ChatLogger | None = request.app.get("chat_logger") + + req_id = secrets.token_hex(4) + started = time.monotonic() + manager.begin_request() + + effective_path = _strip_chat_prefix(request.path) + + # Determine whether this request requires a model and parse the body up front. + # Body parsing errors (OOM, malformed) become 503 so clients get Retry-After. + try: + body = await request.read() if request.can_read_body else None + body = _inject_cache_prompt(body, request.method, effective_path) + except Exception as exc: + logging.exception("[req=%s] body read failed", req_id) + manager.end_request() + return web.json_response( + {"error": "backend unavailable", "detail": str(exc)}, + status=503, + headers={"Retry-After": str(RETRY_AFTER_SECONDS), "X-Request-ID": req_id}, + ) + + path = effective_path.rstrip("/") + needs_model = path in ( + "/v1/chat/completions", "/v1/completions", "/v1/embeddings", + "/chat/completions", "/completions", "/embeddings", + ) + model: str | None = None + if needs_model: + model = _model_from_body(body, manager.config.default_model) + + if chat_logger is not None: + await chat_logger.log_request(request.method, request.path, body, req_id) + + query = request.rel_url.query_string + target_url = f"{manager.config.backend_base_url}{effective_path}" + if query: + target_url = f"{target_url}?{query}" + headers = filter_request_headers(request.headers, manager.config.api_key) + + async def _do_forward() -> web.StreamResponse: + """Inner forward/stream; called from within use_model or directly. + + For non-SSE responses with status >= 500 the body is pre-read so we + can detect a dead-worker error BEFORE committing the response to the + client via ``downstream.prepare()``. If a dead-worker marker is found + a ``DeadWorkerError`` is raised for the caller to handle. All other + paths (success, 4xx, SSE) stream zero-copy as before. + """ + try: + upstream_resp = await session.request( + request.method, target_url, headers=headers, + data=body, allow_redirects=False, timeout=None, + ) + is_sse = "text/event-stream" in upstream_resp.headers.get("Content-Type", "") + + # Pre-read non-SSE error bodies BEFORE prepare() so we can inspect + # them and raise DeadWorkerError without having committed headers. + if upstream_resp.status >= 500 and not is_sse: + error_body = await upstream_resp.content.read() + if _is_dead_worker_response(upstream_resp.status, error_body): + raise DeadWorkerError(upstream_resp.status, error_body) + # Real (non-dead-worker) 5xx — forward as-is. + response_headers = filter_response_headers(upstream_resp.headers) + response_headers["X-Request-ID"] = req_id + downstream = web.StreamResponse( + status=upstream_resp.status, reason=upstream_resp.reason, + headers=response_headers, + ) + try: + await downstream.prepare(request) + except ConnectionResetError: + raise ClientGone from None + try: + await downstream.write(error_body) + except ConnectionResetError: + logging.info("[req=%s] client disconnected", req_id) + finally: + with contextlib.suppress(ConnectionResetError, RuntimeError): + await downstream.write_eof() + duration_ms = int((time.monotonic() - started) * 1000) + logging.info( + "[req=%s] %s %s %s -> %s in %dms", + req_id, client_ip(request), request.method, request.rel_url, + downstream.status, duration_ms, + ) + return downstream + + return await _stream_response( + upstream_resp, request, is_sse, chat_logger, + req_id, started, needs_model, manager, + ) + except DeadWorkerError: + raise # propagate to retry loop + except ClientGone: + # Downstream client (cloudflared / end client) hung up before we finished + # sending the response — raised when downstream.prepare() hit a closing + # transport. Same benign disconnect the write loops already handle, just + # earlier. Not a proxy fault, so log calmly and return 499. + logging.info("[req=%s] client disconnected before response sent", req_id) + # Close the upstream response to release the router's connection. The + # worker had already begun generating (it produced response headers), so + # this is a mid-generation abort too — flag it for the post-cancel guard. + try: + upstream_resp.close() + except Exception: + pass + if needs_model: + manager.mark_worker_suspect() + return web.Response( + status=499, reason="Client Closed Request", + headers={"X-Request-ID": req_id}, + ) + except ClientError as exc: + # Genuine upstream failure (router connection reset/refused, etc.) — a + # real bad gateway. Note: an *upstream* ClientConnectionResetError lands + # here, while a *downstream* one was already converted to ClientGone above. + logging.exception("[req=%s] proxy failure", req_id) + return web.json_response( + {"error": "bad gateway", "detail": str(exc)}, + status=502, headers={"X-Request-ID": req_id}, + ) + + try: + if needs_model: + return await _retry_with_recovery(manager, model, req_id, started, _do_forward) + else: + # Non-model endpoints (health, /v1/models, props, …) forward directly. + return await _do_forward() + finally: + manager.end_request() + + +async def models_handler(request: web.Request) -> web.Response: + """Normalize the router's /models payload so clients see a sane status. + + llama.cpp's router keeps `status.failed = true` (with a stale `exit_code`) + on presets that have never had a successful load in the current process — + a residual diagnostic flag rather than a real "this model is broken" + signal. Newer pi-llama-cpp versions short-circuit to FAILED when they + see `failed: true`, so a freshly-booted router shows every preset as + "Retry" in pi. We rewrite the flag to false whenever `value` says the + preset is simply unloaded — `value` is the source of truth. + """ + manager: ChatRouterManager = request.app["manager"] + session: ClientSession = request.app["session"] + effective_path = _strip_chat_prefix(request.path) + query = request.rel_url.query_string + target_url = f"{manager.config.backend_base_url}{effective_path}" + if query: + target_url = f"{target_url}?{query}" + headers = filter_request_headers(request.headers, manager.config.api_key) + async with session.get(target_url, headers=headers) as upstream: + body_text = await upstream.text() + try: + payload = json.loads(body_text) + except (ValueError, TypeError): + return web.Response(status=upstream.status, body=body_text, + content_type=upstream.content_type or "application/json") + for entry in payload.get("data") or []: + status = entry.get("status") + if isinstance(status, dict) and status.get("value") == "unloaded": + status["failed"] = False + status.pop("exit_code", None) + return web.json_response(payload, status=upstream.status) + + +async def props_handler(request: web.Request) -> web.Response: + """Normalize the router's /props response for unloaded models. + + llama.cpp returns HTTP 400 with `{"error":{"code":400,"message":"model is + not loaded",...}}` when /props is asked about a non-loaded preset. The + pi-llama-cpp extension probes /props as a sanity check after seeing a + model in /models with status "unloaded", and a strict reading of its + parser can mis-classify that 400 response as FAILED (shows "Retry" + instead of "Load & switch"). We rewrite to a clean 200 JSON whose + shape matches the exact equality checks in baseModel.getStatus(). + """ + manager: ChatRouterManager = request.app["manager"] + session: ClientSession = request.app["session"] + effective_path = _strip_chat_prefix(request.path) + query = request.rel_url.query_string + target_url = f"{manager.config.backend_base_url}{effective_path}" + if query: + target_url = f"{target_url}?{query}" + headers = filter_request_headers(request.headers, manager.config.api_key) + async with session.get(target_url, headers=headers) as upstream: + body_text = await upstream.text() + if upstream.status == 400 and "model is not loaded" in body_text: + return web.json_response( + {"error": {"code": 400, "message": "model is not loaded"}} + ) + return web.Response( + status=upstream.status, + body=body_text, + content_type=upstream.content_type or "application/json", + ) + + +def _build_embed_target(request: web.Request, config: ProxyConfig) -> str: + """Build the target URL for an /embedding/* forward. + + Strips the /embedding prefix so embed_proxy sees plain OpenAI-style + paths (/v1/embeddings, /v1/models, /health, …). + """ + tail = request.match_info.get("tail", "") + sub_path = "/" + tail if tail else "/" + query = request.rel_url.query_string + target_url = f"{config.embed_base_url}{sub_path}" + if query: + target_url = f"{target_url}?{query}" + return target_url + + +async def _stream_embed_response( + upstream_resp, + request: web.Request, + req_id: str, + started: float, +) -> web.StreamResponse: + """Create downstream, prepare, stream upstream content, and log. + + Handles both SSE (with keep-alive) and non-SSE responses. + ConnectionResetError on the downstream write is caught and logged. + """ + response_headers = filter_response_headers(upstream_resp.headers) + response_headers["X-Request-ID"] = req_id + downstream = web.StreamResponse( + status=upstream_resp.status, reason=upstream_resp.reason, headers=response_headers, + ) + await downstream.prepare(request) + try: + is_sse = "text/event-stream" in upstream_resp.headers.get("Content-Type", "") + if is_sse: + while True: + try: + chunk = await asyncio.wait_for(upstream_resp.content.readany(), timeout=25) + except asyncio.TimeoutError: + await downstream.write(b": keep-alive\n\n") + continue + if not chunk: + break + await downstream.write(chunk) + else: + async for chunk in upstream_resp.content.iter_any(): + await downstream.write(chunk) + except ConnectionResetError: + logging.info("[embed req=%s] client disconnected", req_id) + finally: + with contextlib.suppress(ConnectionResetError, RuntimeError): + await downstream.write_eof() + duration_ms = int((time.monotonic() - started) * 1000) + logging.info( + "[embed req=%s] %s %s %s -> %s in %dms", + req_id, client_ip(request), request.method, request.path, + downstream.status, duration_ms, + ) + return downstream + + +async def embed_forward(request: web.Request) -> web.StreamResponse: + """Reverse-proxy /embedding/* to the standalone embed_proxy on :8003. + + Strips the prefix so embed_proxy sees plain OpenAI-style paths + (/v1/embeddings, /v1/models, /health, ...). embed_proxy owns its own + router, load/unload, and idle timer — this is a dumb HTTP forwarder. + """ + session: ClientSession = request.app["session"] + config: ProxyConfig = request.app["config"] + req_id = secrets.token_hex(4) + started = time.monotonic() + + target_url = _build_embed_target(request, config) + body = await request.read() if request.can_read_body else None + headers = filter_request_headers(request.headers, config.api_key) + + try: + upstream_resp = await session.request( + request.method, target_url, headers=headers, + data=body, allow_redirects=False, timeout=None, + ) + return await _stream_embed_response( + upstream_resp, request, req_id, started, + ) + except ClientError as exc: + logging.exception("[embed req=%s] forward failure", req_id) + return web.json_response( + {"error": "bad gateway", "detail": str(exc)}, + status=502, headers={"X-Request-ID": req_id}, + ) diff --git a/router_manager.py b/router_manager.py index 000eb4c..f39e6b3 100644 --- a/router_manager.py +++ b/router_manager.py @@ -127,12 +127,7 @@ async def _load_locked(self, model: str) -> None: self._loaded = model return logging.info("Loading model: %s", model) - url = f"{self.config.backend_base_url}/models/load" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - async with self.session.post(url, headers=headers, json={"model": model}) as r: + async with self._post_json("/models/load", {"model": model}) as r: if r.status >= 400: body = await r.text() if "already running" in body: @@ -140,29 +135,17 @@ async def _load_locked(self, model: str) -> None: return raise RuntimeError(f"load returned {r.status}: {body}") deadline = time.monotonic() + self.LOAD_TIMEOUT - while time.monotonic() < deadline: - status = await self._status(model) - if status == "loaded": - self._loaded = model - logging.info("Model loaded: %s", model) - return - if status == "failed": - raise RuntimeError(f"model {model} failed to load") - await asyncio.sleep(0.5) - raise TimeoutError(f"model {model} did not load in {self.LOAD_TIMEOUT}s") + if await self._poll_until_loaded(model, deadline): + self._loaded = model + logging.info("Model loaded: %s", model) async def unload(self, reason: str) -> None: if self._loaded is None: return model = self._loaded logging.info("Unloading %s (%s)", model, reason) - url = f"{self.config.backend_base_url}/models/unload" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } try: - async with self.session.post(url, headers=headers, json={"model": model}) as r: + async with self._post_json("/models/unload", {"model": model}) as r: if r.status >= 400: logging.warning("unload returned %s", r.status) except ClientError as e: @@ -170,9 +153,7 @@ async def unload(self, reason: str) -> None: self._loaded = None async def _status(self, model: str) -> str: - url = f"{self.config.backend_base_url}/v1/models" - headers = {"Authorization": f"Bearer {self.config.api_key}"} - async with self.session.get(url, headers=headers) as r: + async with self._get_json("/v1/models") as r: data = await r.json() for entry in data.get("data", []): if entry.get("id") == model: @@ -189,6 +170,44 @@ async def unload_if_idle(self) -> None: if idle >= self.config.idle_timeout: await self.unload(f"idle for {int(idle)}s") + # ── HTTP helpers ────────────────────────────────────────────────────── + + def _auth_headers(self, content_type: str = "application/json") -> dict: + return { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": content_type, + } + + def _post_json(self, path: str, json_data: dict, + timeout: float | None = None): + url = f"{self.config.backend_base_url}{path}" + kwargs: dict = {"headers": self._auth_headers(), "json": json_data} + if timeout is not None: + kwargs["timeout"] = ClientTimeout(total=timeout) + return self.session.post(url, **kwargs) + + def _get_json(self, path: str, timeout: float | None = None): + url = f"{self.config.backend_base_url}{path}" + kwargs: dict = {"headers": {"Authorization": f"Bearer {self.config.api_key}"}} + if timeout is not None: + kwargs["timeout"] = ClientTimeout(total=timeout) + return self.session.get(url, **kwargs) + + async def _poll_until_loaded(self, model: str, deadline: float) -> bool: + """Poll *model* status until loaded, failed, or *deadline* expires. + + Returns True on success (model reached "loaded"). Raises RuntimeError + on "failed" and TimeoutError on deadline expiry. + """ + while time.monotonic() < deadline: + status = await self._status(model) + if status == "loaded": + return True + if status == "failed": + raise RuntimeError(f"model {model} failed to load") + await asyncio.sleep(0.5) + raise TimeoutError(f"model {model} did not load in {self.LOAD_TIMEOUT}s") + # ── Chat subclass ─────────────────────────────────────────────────────────── @@ -246,12 +265,7 @@ async def _load_locked(self, model: str) -> None: self._loaded_at = time.monotonic() return logging.info("Loading model: %s", model) - url = f"{self.config.backend_base_url}/models/load" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } - async with self.session.post(url, headers=headers, json={"model": model}) as r: + async with self._post_json("/models/load", {"model": model}) as r: if r.status >= 400: body = await r.text() if "already running" in body: @@ -260,17 +274,10 @@ async def _load_locked(self, model: str) -> None: return raise RuntimeError(f"load returned {r.status}: {body}") deadline = time.monotonic() + self.LOAD_TIMEOUT - while time.monotonic() < deadline: - status = await self._status(model) - if status == "loaded": - self._loaded = model - self._loaded_at = time.monotonic() - logging.info("Model loaded: %s", model) - return - if status == "failed": - raise RuntimeError(f"model {model} failed to load") - await asyncio.sleep(0.5) - raise TimeoutError(f"model {model} did not load in {self.LOAD_TIMEOUT}s") + if await self._poll_until_loaded(model, deadline): + self._loaded = model + self._loaded_at = time.monotonic() + logging.info("Model loaded: %s", model) @contextlib.asynccontextmanager async def use_model(self, model: str): @@ -307,6 +314,167 @@ async def use_model(self, model: str): finally: self._end_forward() + # ── recover_worker helpers ──────────────────────────────────────────── + + async def _recover_check_peer( + self, model: str, dead_detected_at: float, + ) -> bool: + """Return True if a peer coroutine already recovered *model*.""" + if self._loaded == model and self._loaded_at > dead_detected_at: + logging.info( + "[recover_worker] %s already recovered by peer (loaded_at=%.3f > detected=%.3f)", + model, self._loaded_at, dead_detected_at, + ) + return True + return False + + async def _recover_unload( + self, model: str, url_unload: str, auth_headers: dict, + ) -> None: + """Best-effort unload POST — force the router to tear down the dead worker entry.""" + try: + async with self.session.post( + url_unload, headers=auth_headers, json={"model": model}, + timeout=ClientTimeout(total=10), + ) as r: + body_text = await r.text() + if r.status >= 400: + logging.warning("[recover_worker] unload returned %s: %s", r.status, body_text) + else: + logging.info("[recover_worker] unload OK for %s", model) + except (ClientError, asyncio.TimeoutError) as e: + logging.warning("[recover_worker] unload error (ignored): %s", e) + + async def _recover_wait_unloaded(self, model: str) -> None: + """Poll until the router confirms *model* is no longer loaded (15 s deadline).""" + unload_deadline = time.monotonic() + 15 + while time.monotonic() < unload_deadline: + try: + status = await self._status(model) + if status != "loaded": + logging.info("[recover_worker] router confirms %s is %s", model, status) + break + except Exception: + pass + await asyncio.sleep(0.5) + else: + logging.warning( + "[recover_worker] router still shows loaded after 15s — proceeding anyway" + ) + + async def _recover_load_once( + self, model: str, load_attempt: int, max_load_attempts: int, + load_url: str, auth_headers: dict, + ) -> bool | str: + """Execute one load attempt with rapid-poll and stability check. + + Returns True on success (worker stable), False on terminal failure + (POST error or status "failed"), "retry_now" on queued-exit flash + or mid-load death (retry immediately, no unload/sleep), + "retry_after_unload" on poll timeout (caller should unload + sleep). + """ + logging.info( + "[recover_worker] load attempt %d/%d for %s", + load_attempt, max_load_attempts, model, + ) + try: + async with self.session.post( + load_url, headers=auth_headers, json={"model": model}, + timeout=ClientTimeout(total=10), + ) as r: + body_text = await r.text() + if r.status >= 400 and "already running" not in body_text: + logging.warning("[recover_worker] load returned %s: %s", r.status, body_text) + except (ClientError, asyncio.TimeoutError) as e: + logging.error("[recover_worker] load POST failed: %s", e) + return False + + # Poll until status leaves "loading" (either loaded or unloaded) + poll_deadline = time.monotonic() + 120 # per_load_timeout + prev_status = "" + seen_loading = False + while time.monotonic() < poll_deadline: + try: + status = await self._status(model) + except Exception: + await asyncio.sleep(0.5) + continue + if status != prev_status: + logging.info("[recover_worker] %s status: %s", model, status) + prev_status = status + if status == "loading": + seen_loading = True + if status == "loaded": + # Give the router 1.5s to process any pending exit signal + # before declaring victory. + await asyncio.sleep(1.5) + status2 = await self._status(model) + if status2 == "loaded": + self._loaded = model + self._loaded_at = time.monotonic() + logging.info( + "[recover_worker] worker stable for %s (attempt %d)", + model, load_attempt, + ) + return True + logging.warning( + "[recover_worker] fresh worker for %s exited immediately " + "(status after 1.5s: %s, attempt %d/%d)", + model, status2, load_attempt, max_load_attempts, + ) + return "retry_now" # queued exit consumed — try next load attempt + if status == "failed": + logging.error("[recover_worker] model %s failed to load", model) + return False + if status == "unloaded" and seen_loading: + # The worker started loading but was then stopped before ever + # reaching "loaded" — a queued exit signal consumed it (the + # router force-kills it after its ~10s stop timeout) or it + # crashed mid-load. For a slow-loading model the kill lands + # before "loaded" is ever observed, so the loaded→flash case + # above never fires. Retry the next load immediately instead + # of spinning here for the full per_load_timeout. + logging.warning( + "[recover_worker] fresh worker for %s died during load " + "(status: unloaded, attempt %d/%d) — retrying", + model, load_attempt, max_load_attempts, + ) + return "retry_now" # queued exit consumed — try next load attempt + await asyncio.sleep(0.25) + + logging.error( + "[recover_worker] timed out waiting for %s to load (attempt %d/%d)", + model, load_attempt, max_load_attempts, + ) + return "retry_after_unload" # timeout — caller may retry unload before next attempt + + async def _recover_retry_unload( + self, model: str, url_unload: str, auth_headers: dict, + ) -> None: + """Retry unload after a timed-out load attempt, then pause 2 s.""" + logging.info("[recover_worker] retrying unload before next load attempt") + try: + async with self.session.post( + url_unload, headers=auth_headers, json={"model": model}, + timeout=ClientTimeout(total=10), + ) as r: + body_text = await r.text() + if r.status >= 400: + logging.warning( + "[recover_worker] retry-unload returned %s: %s", + r.status, body_text, + ) + else: + logging.info( + "[recover_worker] retry-unload OK for %s", model + ) + except (ClientError, asyncio.TimeoutError) as e: + logging.warning( + "[recover_worker] retry-unload error (ignored): %s", e + ) + # Brief pause to let the router settle before the next load. + await asyncio.sleep(2) + async def recover_worker(self, model: str, dead_detected_at: float) -> bool: """Force-cycle the router worker for *model* after a dead-worker 500. @@ -323,13 +491,7 @@ async def recover_worker(self, model: str, dead_detected_at: float) -> bool: error so the caller can return 503. """ async with self._load_lock: - # Re-check: if _loaded_at was updated AFTER the dead-worker was - # detected, a peer coroutine already completed recovery. - if self._loaded == model and self._loaded_at > dead_detected_at: - logging.info( - "[recover_worker] %s already recovered by peer (loaded_at=%.3f > detected=%.3f)", - model, self._loaded_at, dead_detected_at, - ) + if await self._recover_check_peer(model, dead_detected_at): return True logging.warning( @@ -345,35 +507,11 @@ async def recover_worker(self, model: str, dead_detected_at: float) -> bool: "Authorization": f"Bearer {self.config.api_key}", "Content-Type": "application/json", } - try: - async with self.session.post( - url_unload, headers=auth_headers, json={"model": model}, - timeout=ClientTimeout(total=10), - ) as r: - body_text = await r.text() - if r.status >= 400: - logging.warning("[recover_worker] unload returned %s: %s", r.status, body_text) - else: - logging.info("[recover_worker] unload OK for %s", model) - except (ClientError, asyncio.TimeoutError) as e: - logging.warning("[recover_worker] unload error (ignored): %s", e) + await self._recover_unload(model, url_unload, auth_headers) # Wait for the router to reflect the unloaded state so that # _load_locked doesn't see "loaded" and short-circuit. - unload_deadline = time.monotonic() + 15 - while time.monotonic() < unload_deadline: - try: - status = await self._status(model) - if status != "loaded": - logging.info("[recover_worker] router confirms %s is %s", model, status) - break - except Exception: - pass - await asyncio.sleep(0.5) - else: - logging.warning( - "[recover_worker] router still shows loaded after 15s — proceeding anyway" - ) + await self._recover_wait_unloaded(model) # Reload — spawn fresh worker(s), draining any buffered exit # signals the router queued during unload. Each exit signal @@ -388,113 +526,23 @@ async def recover_worker(self, model: str, dead_detected_at: float) -> bool: # the loaded→unloaded flash and retries the load immediately. load_url = f"{self.config.backend_base_url}/models/load" max_load_attempts = 4 - per_load_timeout = 120 # seconds to wait for status != "loading" for load_attempt in range(1, max_load_attempts + 1): - logging.info( - "[recover_worker] load attempt %d/%d for %s", - load_attempt, max_load_attempts, model, + result = await self._recover_load_once( + model, load_attempt, max_load_attempts, + load_url, auth_headers, ) - try: - async with self.session.post( - load_url, headers=auth_headers, json={"model": model}, - timeout=ClientTimeout(total=10), - ) as r: - body_text = await r.text() - if r.status >= 400 and "already running" not in body_text: - logging.warning("[recover_worker] load returned %s: %s", r.status, body_text) - except (ClientError, asyncio.TimeoutError) as e: - logging.error("[recover_worker] load POST failed: %s", e) + if result is True: + return True + if result is False: return False - - # Poll until status leaves "loading" (either loaded or unloaded) - poll_deadline = time.monotonic() + per_load_timeout - prev_status = "" - seen_loading = False # have we observed this attempt actually start? - while time.monotonic() < poll_deadline: - try: - status = await self._status(model) - except Exception: - await asyncio.sleep(0.5) - continue - if status != prev_status: - logging.info("[recover_worker] %s status: %s", model, status) - prev_status = status - if status == "loading": - seen_loading = True - if status == "loaded": - # Give the router 1.5s to process any pending exit signal - # before declaring victory. - await asyncio.sleep(1.5) - status2 = await self._status(model) - if status2 == "loaded": - self._loaded = model - self._loaded_at = time.monotonic() - logging.info( - "[recover_worker] worker stable for %s (attempt %d)", - model, load_attempt, - ) - return True - logging.warning( - "[recover_worker] fresh worker for %s exited immediately " - "(status after 1.5s: %s, attempt %d/%d)", - model, status2, load_attempt, max_load_attempts, - ) - break # queued exit consumed — try next load attempt - if status == "failed": - logging.error("[recover_worker] model %s failed to load", model) - return False - if status == "unloaded" and seen_loading: - # The worker started loading but was then stopped before ever - # reaching "loaded" — a queued exit signal consumed it (the - # router force-kills it after its ~10s stop timeout) or it - # crashed mid-load. For a slow-loading model the kill lands - # before "loaded" is ever observed, so the loaded→flash case - # above never fires. Retry the next load immediately instead - # of spinning here for the full per_load_timeout. - logging.warning( - "[recover_worker] fresh worker for %s died during load " - "(status: unloaded, attempt %d/%d) — retrying", - model, load_attempt, max_load_attempts, - ) - break # queued exit consumed — try next load attempt - await asyncio.sleep(0.25) - else: - logging.error( - "[recover_worker] timed out waiting for %s to load (attempt %d/%d)", - model, load_attempt, max_load_attempts, + # retryable: "retry_now" → next attempt immediately; + # "retry_after_unload" → unload + 2s pause first + # (only if attempts remain) + if result == "retry_after_unload" and load_attempt < max_load_attempts: + await self._recover_retry_unload( + model, url_unload, auth_headers, ) - # Don't return False here — let the for-loop try the next attempt. - # The router may need a fresh /models/load POST to clear a stuck - # "loading" state from a previous crash. - # Retry the unload to give the router a chance to tear down - # any zombie worker entry from the timed-out attempt. - if load_attempt < max_load_attempts: - logging.info( - "[recover_worker] retrying unload before next load attempt" - ) - try: - async with self.session.post( - url_unload, headers=auth_headers, json={"model": model}, - timeout=ClientTimeout(total=10), - ) as r: - body_text = await r.text() - if r.status >= 400: - logging.warning( - "[recover_worker] retry-unload returned %s: %s", - r.status, body_text, - ) - else: - logging.info( - "[recover_worker] retry-unload OK for %s", model - ) - except (ClientError, asyncio.TimeoutError) as e: - logging.warning( - "[recover_worker] retry-unload error (ignored): %s", e - ) - # Brief pause to let the router settle before the next load. - await asyncio.sleep(2) - continue logging.error("[recover_worker] all %d load attempts failed for %s", max_load_attempts, model) return False @@ -510,11 +558,6 @@ async def _probe_worker(self, model: str) -> bool: cancel. Returns True if the worker answers normally; False if it returns a dead-worker error, errors out, or times out (a busy-finishing-an-orphan or genuinely-wedged worker both warrant a recovery cycle).""" - url = f"{self.config.backend_base_url}/v1/chat/completions" - headers = { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - } payload = { "model": model, "messages": [{"role": "user", "content": "ping"}], @@ -522,9 +565,9 @@ async def _probe_worker(self, model: str) -> bool: "stream": False, } try: - async with self.session.post( - url, headers=headers, json=payload, - timeout=ClientTimeout(total=GUARD_PROBE_TIMEOUT), + async with self._post_json( + "/v1/chat/completions", payload, + timeout=GUARD_PROBE_TIMEOUT, ) as r: body = await r.read() if _is_dead_worker_response(r.status, body): diff --git a/tests/test_toolcall_rescue.py b/tests/test_toolcall_rescue.py index be3eab0..c1f9530 100644 --- a/tests/test_toolcall_rescue.py +++ b/tests/test_toolcall_rescue.py @@ -6,7 +6,7 @@ import json import unittest -import proxy +from chat_logger import SSEChunkLogger # --------------------------------------------------------------------------- @@ -92,14 +92,14 @@ def _split_sse_events(raw: bytes) -> list[dict]: return events -def _make_logger(chunks: list[bytes]) -> proxy.SSEChunkLogger: - return proxy.SSEChunkLogger(FakeUpstream(chunks), FakeChatLogger()) +def _make_logger(chunks: list[bytes]) -> SSEChunkLogger: + return SSEChunkLogger(FakeUpstream(chunks), FakeChatLogger()) async def _drive(chunks: list[bytes]) -> tuple[str, FakeChatLogger]: """Run SSEChunkLogger through all chunks, return (combined output, logger).""" logger = FakeChatLogger() - wrapped = proxy.SSEChunkLogger(FakeUpstream(chunks), logger) + wrapped = SSEChunkLogger(FakeUpstream(chunks), logger) out: list[bytes] = [] while True: piece = await wrapped.readany() @@ -399,7 +399,7 @@ def test_subevent_chunk_boundaries(self) -> None: blob += b"data: [DONE]\n\n" content = ChunkedContent(blob, slice_size=13) - wrapped = proxy.SSEChunkLogger( + wrapped = SSEChunkLogger( FakeUpstream([]), # headers unused; we swap content FakeChatLogger(), ) @@ -450,7 +450,7 @@ def test_rescue_subevent_chunks(self) -> None: blob += b"data: [DONE]\n\n" content = ChunkedContent(blob, slice_size=13) - wrapped = proxy.SSEChunkLogger( + wrapped = SSEChunkLogger( FakeUpstream([]), FakeChatLogger(), )