From 40707175527f6b1ae537cee41b0d41accc793677 Mon Sep 17 00:00:00 2001 From: shangfeng-liu Date: Wed, 14 Oct 2026 13:18:13 +0800 Subject: [PATCH 1/4] fix(server): recover backend workers after runtime cuda oom Replace the poisoned backend process group after a typed runtime CUDA OOM while preserving the API process and bounded new admissions. Retire all old workers and reject traffic from older epochs before opening the new generation. Constraint: Recovery defaults off and dispatched requests are never replayed. Scope-risk: moderate --- README.md | 1 + docs/oom_recovery.md | 24 ++ python/freetoken/core.py | 1 + python/freetoken/message/backend.py | 3 + python/freetoken/message/frontend.py | 2 + python/freetoken/message/tokenizer.py | 7 + python/freetoken/scheduler/config.py | 3 + python/freetoken/scheduler/prefill.py | 2 + python/freetoken/scheduler/scheduler.py | 11 +- python/freetoken/scheduler/utils.py | 1 + python/freetoken/server/accounting.py | 2 +- python/freetoken/server/anthropic_api.py | 6 +- python/freetoken/server/api_server.py | 202 ++++++++-- python/freetoken/server/args.py | 26 ++ python/freetoken/server/control_api.py | 11 +- python/freetoken/server/generation.py | 3 +- python/freetoken/server/launch.py | 47 ++- python/freetoken/server/openai_api.py | 18 +- python/freetoken/server/responses_api.py | 6 +- python/freetoken/server/stats.py | 17 + python/freetoken/server/supervisor.py | 105 ++++- python/freetoken/tokenizer/server.py | 14 +- python/freetoken/tokenizer/tokenize.py | 3 +- .../scheduler/test_abort_inflight_prefill.py | 2 +- tests/scheduler/test_cost_accounting_core.py | 33 ++ tests/server/test_anthropic_api.py | 2 +- tests/server/test_effort_dialect.py | 2 +- tests/server/test_generation_accounting.py | 2 +- tests/server/test_message_wire.py | 26 +- tests/server/test_oom_recovery.py | 377 ++++++++++++++++++ tests/server/test_openai_api.py | 2 +- tests/server/test_prepare_stop_accounting.py | 12 +- tests/server/test_rebuild_maintenance.py | 5 +- tests/server/test_responses_api.py | 2 +- tests/server/test_supervisor.py | 286 ++++++++++++- 35 files changed, 1194 insertions(+), 72 deletions(-) create mode 100644 docs/oom_recovery.md create mode 100644 tests/server/test_oom_recovery.py diff --git a/README.md b/README.md index 6d9ab4ae47..21d40de064 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ For More details: - [Install FreeToken](https://github.com/FlashML-org/FreeToken/blob/main/docs/install.md) - [Quick start](https://github.com/FlashML-org/FreeToken/blob/main/docs/quickstart.md) +- [CUDA OOM recovery](docs/oom_recovery.md) - [Supported models](https://github.com/FlashML-org/FreeToken/blob/main/docs/models.md) - [CLI reference](https://github.com/FlashML-org/FreeToken/blob/main/docs/cli.md) - [Repairing old FTW checkpoints](https://github.com/FlashML-org/FreeToken/blob/main/docs/ftw-hotfix.md) diff --git a/docs/oom_recovery.md b/docs/oom_recovery.md new file mode 100644 index 0000000000..c5b7075b47 --- /dev/null +++ b/docs/oom_recovery.md @@ -0,0 +1,24 @@ +# CUDA OOM recovery + +FreeToken normally leaves backend-worker recovery to the outer daemon. This remains the default. +Set `--oom-recovery` to opt into backend-process-group replacement up to the configured limit after a **post-ready** scheduler +`torch.cuda.OutOfMemoryError`: + +```bash +ft serve --model /path/to/model --oom-recovery +``` + +The replacement limit, readiness wait, and bounded admission queue are configurable with +`--oom-recovery-max-restarts` (default `1`), `--oom-recovery-ready-timeout` (default `900` +seconds), and `--oom-recovery-wait-queue` (default `128`). Requests already dispatched when the +OOM occurs receive one terminal retryable error and are never replayed. New requests can wait for +the fresh readiness handshake; queue overflow, timeout, or replacement failure returns 503. + +Only a typed worker receipt for `torch.cuda.OutOfMemoryError` starts this path. Other worker +failures keep the existing fail-closed process-exit behavior. Before allocation starts, every old +scheduler/tokenizer process is terminated and reaped. Each replacement has a new backend epoch; +workers drop buffered IPC work from retired epochs, including requests, aborts, and cache rebuild +traffic. + +`/health` and `/v1/stats` expose `maintenance`, `backend_epoch`, `restart_attempts`, and +`restart_waiters` under their recovery information. diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 812f9a4349..9fc69d18ff 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -47,6 +47,7 @@ class Req: mrope_delta: int = 0 # The whole prompt's length; 0 = this request's input_ids (a ChunkedReq's end at its chunk). prompt_len: int = 0 + epoch: int = 0 # --- hybrid-radix (GDN linear-state) per-request slots; None for non-hybrid models or # until allocated from LinearStatePool. Set by the scheduler (P2). --- diff --git a/python/freetoken/message/backend.py b/python/freetoken/message/backend.py index ccc1d7921a..886756f78f 100644 --- a/python/freetoken/message/backend.py +++ b/python/freetoken/message/backend.py @@ -79,11 +79,13 @@ class UserMsg(BaseBackendMsg): # precomputed [3, len(input_ids)] mrope positions and decode delta; None for text-only requests and 1-D rope models mrope_positions: torch.Tensor | None = None mrope_delta: int = 0 + epoch: int = 0 @dataclass class AbortBackendMsg(BaseBackendMsg): uid: int + epoch: int = 0 @dataclass @@ -95,3 +97,4 @@ class CacheRebuildBackendMsg(BaseBackendMsg): num_mamba_slots: int | None = None num_swa_pages: int | None = None mode: str = "if_idle" # only "if_idle" is supported; "drain" is deferred (rejected) + epoch: int = 0 diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b9..5cf393f6af 100644 --- a/python/freetoken/message/frontend.py +++ b/python/freetoken/message/frontend.py @@ -56,6 +56,7 @@ class UserReply(BaseFrontendMsg): finish_reason: str | None = None # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence'). matched_stop: str | None = None + epoch: int = 0 @dataclass @@ -68,3 +69,4 @@ class CacheRebuildReply(BaseFrontendMsg): mamba_slots: int = 0 num_swa_pages: int = 0 error: str | None = None + epoch: int = 0 diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index d52d658e63..715800b47b 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -48,6 +48,7 @@ class DetokenizeMsg(BaseTokenizerMsg): swa_total_tokens: int = 0 # Bytes this engine process holds on the GPU (torch reserved pool). 0 on CPU. gpu_mem_bytes: int = 0 + epoch: int = 0 @dataclass @@ -63,6 +64,7 @@ class PromptAdmittedMsg(BaseTokenizerMsg): uid: int prompt_tokens: int cached_tokens: int = 0 + epoch: int = 0 @dataclass @@ -74,11 +76,13 @@ class TokenizeMsg(BaseTokenizerMsg): tools: List[Dict[str, Any]] | None = None images: List[bytes] | None = None inline_system_policy: str | None = None + epoch: int = 0 @dataclass class AbortMsg(BaseTokenizerMsg): uid: int + epoch: int = 0 @dataclass @@ -90,6 +94,7 @@ class CacheRebuildMsg(BaseTokenizerMsg): num_mamba_slots: int | None = None num_swa_pages: int | None = None mode: str = "if_idle" + epoch: int = 0 @dataclass @@ -102,6 +107,7 @@ class CacheRebuildResultMsg(BaseTokenizerMsg): mamba_slots: int = 0 num_swa_pages: int = 0 error: str | None = None + epoch: int = 0 @dataclass @@ -116,3 +122,4 @@ class ErrorReplyMsg(BaseTokenizerMsg): # can react without parsing prose. Today only "context_length_exceeded" (prompt longer than # the servable context); None = no specific class, the message is all there is. code: str | None = None + epoch: int = 0 diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index b2bbe0aefe..a9b0fb01d5 100644 --- a/python/freetoken/scheduler/config.py +++ b/python/freetoken/scheduler/config.py @@ -18,6 +18,9 @@ class SchedulerConfig(EngineConfig): offline_mode: bool = False decode_log_interval: int = 40 special_token_ckpt: bool = False + # Parent-owned worker generation marker. ServerArgs inherits it; keeping the default at the + # scheduler boundary preserves offline and CPU callers that construct SchedulerConfig. + backend_epoch: int = 0 # networking config _unique_suffix: str = field(default_factory=_get_pid_suffix) diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py index de4bddfb4a..8d74a46fb8 100644 --- a/python/freetoken/scheduler/prefill.py +++ b/python/freetoken/scheduler/prefill.py @@ -216,6 +216,7 @@ def _add_one_req( cache_handle=cache_handle, sampling_params=pending_req.sampling_params, prompt_len=pending_req.input_len, + epoch=pending_req.epoch, ) req.mm_items = pending_req.mm_items req.mrope_positions_full = pending_req.mrope_positions_full @@ -288,6 +289,7 @@ def add_one_req(self, req: UserMsg) -> None: mm_items=req.mm_items, mrope_positions_full=req.mrope_positions, mrope_delta=req.mrope_delta, + epoch=req.epoch, ) ) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index d057d3b908..f2c34039c4 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -379,6 +379,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: finish_reason=finish_reason, matched_stop=matched_stop, stop_strs=req.sampling_params.stop_strs or None, + epoch=req.epoch, ) ) @@ -486,6 +487,9 @@ def _gpu_mem_bytes(self) -> int: return torch.cuda.memory_reserved(self.device) def _process_one_msg(self, msg: BaseBackendMsg) -> None: + if getattr(msg, "epoch", self.config.backend_epoch) != self.config.backend_epoch: + logger.warning_rank0("Dropping stale backend message for epoch %s", getattr(msg, "epoch", None)) + return if isinstance(msg, BatchBackendMsg): for msg in msg.data: self._process_one_msg(msg) @@ -507,6 +511,7 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: ErrorReplyMsg( uid=msg.uid, error="image input is not supported by this server", + epoch=msg.epoch, ) ] ) @@ -533,6 +538,7 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: ), # OpenAI's standard class for this, for clients that read a code. code="context_length_exceeded", + epoch=msg.epoch, ) ] ) @@ -654,6 +660,7 @@ def _reply_rebuild(self, request_id: str, status: str, error: str | None = None) mamba_slots=geo["num_mamba_slots"] or 0, num_swa_pages=geo["num_swa_pages"] or 0, error=error, + epoch=self.config.backend_epoch, ) ] ) @@ -882,7 +889,7 @@ def _report_prompt_admissions(self, batch: Batch) -> None: return self.send_result( [ - PromptAdmittedMsg(uid=uid, prompt_tokens=prompt_tokens, cached_tokens=cached_tokens) + PromptAdmittedMsg(uid=uid, prompt_tokens=prompt_tokens, cached_tokens=cached_tokens, epoch=self.config.backend_epoch) for uid, prompt_tokens, cached_tokens in batch.prompt_admissions ] ) @@ -893,7 +900,7 @@ def _flush_abort_acks(self) -> None: return uids = sorted(pending) pending.clear() - self.send_result([ErrorReplyMsg(uid=uid, error="request aborted") for uid in uids]) + self.send_result([ErrorReplyMsg(uid=uid, error="request aborted", epoch=self.config.backend_epoch) for uid in uids]) def _forward(self, forward_input: ForwardInput) -> ForwardOutput: batch, sample_args, input_mapping, output_mapping = forward_input diff --git a/python/freetoken/scheduler/utils.py b/python/freetoken/scheduler/utils.py index 83b79f2a3a..7194a16548 100644 --- a/python/freetoken/scheduler/utils.py +++ b/python/freetoken/scheduler/utils.py @@ -20,6 +20,7 @@ class PendingReq: mm_items: list | None = None mrope_positions_full: torch.Tensor | None = None mrope_delta: int = 0 + epoch: int = 0 @property def input_len(self) -> int: diff --git a/python/freetoken/server/accounting.py b/python/freetoken/server/accounting.py index ba2f309ddb..81b25ca1c4 100644 --- a/python/freetoken/server/accounting.py +++ b/python/freetoken/server/accounting.py @@ -66,7 +66,7 @@ async def prepare_stop_accounting( maintenance = getattr(state, "maintenance_state", "serving") if maintenance == "rebuilding": raise AccountingDrainError("cache rebuild is in progress; retry stop after it finishes") - if maintenance not in {"loading", "serving", "stopping", "failed"}: + if maintenance not in {"loading", "serving", "restarting", "stopping", "failed"}: raise AccountingDrainError(f"engine cannot prepare stop from state {maintenance!r}") # This assignment and FrontendManager.new_user's check run on the same event loop, making diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index d9a6972f09..4ba3a4ba5f 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -20,6 +20,8 @@ from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse, StreamingResponse +from .accounting import AdmissionClosedError + from .anthropic_models import ( AnthropicContentBlock, AnthropicCountTokensRequest, @@ -85,7 +87,7 @@ async def v1_messages(req: AnthropicMessagesRequest, request: Request): log_request("/v1/messages", req, request) state = get_state() mstate = getattr(state, "maintenance_state", "serving") - if mstate != "serving": + if mstate not in {"serving", "restarting"}: detail = "model is still loading" if mstate == "loading" else "cache rebuild in progress" return _anthropic_error_response(503, "overloaded_error", detail) return await handle_anthropic_messages(req, request, state, get_model_sampling()) @@ -123,6 +125,8 @@ async def handle_anthropic_messages( inline_system_policy=getattr(state.config, "anthropic_inline_system", "auto"), ) uid = await submit_generation(spec, state) + except AdmissionClosedError as exc: + return _anthropic_error_response(503, "overloaded_error", str(exc)) except ValueError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index a9b16adfd7..0d6082735c 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -2,6 +2,7 @@ import asyncio import contextlib +import concurrent.futures import hashlib import hmac import json @@ -114,6 +115,36 @@ def _reap_backend_workers(processes: List[Any], timeout: float = 5.0) -> None: continue +def _run_on_loop(loop: Any, action: Callable[[], None], *, timeout: float = 5.0) -> None: + """Run a lifecycle mutation on ``loop`` and acknowledge its terminal outcome. + + A timeout cancels only an unclaimed callback. Once the loop has claimed it, callers wait for + its real result so they never infer a failed lifecycle transition after it actually mutated. + """ + if loop is None: + action() + return + done: concurrent.futures.Future[None] = concurrent.futures.Future() + + def _apply() -> None: + if not done.set_running_or_notify_cancel(): + return + try: + action() + except BaseException as exc: # noqa: BLE001 + done.set_exception(exc) + else: + done.set_result(None) + + loop.call_soon_threadsafe(_apply) + try: + done.result(timeout=timeout) + except concurrent.futures.TimeoutError: + if done.cancel(): + raise + done.result() + + def _unwrap_msg(msg: BaseFrontendMsg) -> List[UserReply]: if isinstance(msg, BatchFrontendMsg): result = [] @@ -139,6 +170,7 @@ class FrontendManager: uid_counter: int = 0 initialized: bool = False ack_map: Dict[int, List[UserReply]] = field(default_factory=dict) + request_epochs: Dict[int, int] = field(default_factory=dict) event_map: Dict[int, asyncio.Event] = field(default_factory=dict) # Stable identity for this serve process. Generated before the backend is ready so every # /health state (loading/ok/error) and /v1/stats can identify the same engine generation. @@ -184,6 +216,10 @@ class FrontendManager: # handler) tears these down itself, AFTER setting _SHUTTING_DOWN, so the supervisor observes # the shutdown flag before the ensuing deaths. See _terminate_backend_workers. backend_processes: List[Any] = field(default_factory=list) + backend_epoch: int = 0 + restart_attempts: int = 0 + _ready_event: Any = None + _restart_waiters: int = 0 # Event loop the listener runs on, captured when the listener starts (_create_listener_once). # Lets a cross-thread caller — the supervisor thread's failure callback — marshal rebuild # future resolution back onto the loop (asyncio Futures are not thread-safe). None until the @@ -234,7 +270,19 @@ def _warm() -> None: threading.Thread(target=_warm, daemon=True, name="frontend-tokenizer-warm").start() - def new_user(self) -> int: + async def new_user(self) -> int: + if self.maintenance_state == "restarting": + if self._ready_event is None: + self._ready_event = asyncio.Event() + if self._restart_waiters >= self.config.oom_recovery_wait_queue: + raise AdmissionClosedError("server restarting queue is full; retry") + self._restart_waiters += 1 + try: + await asyncio.wait_for(self._ready_event.wait(), timeout=self.config.oom_recovery_ready_timeout) + except asyncio.TimeoutError as exc: + raise AdmissionClosedError("worker restart timed out; retry") from exc + finally: + self._restart_waiters -= 1 if self.maintenance_state != "serving": raise AdmissionClosedError( f"server unavailable: engine is {self.maintenance_state}" @@ -242,23 +290,37 @@ def new_user(self) -> int: uid = self.uid_counter self.uid_counter += 1 self.ack_map[uid] = [] + self.request_epochs[uid] = self.backend_epoch self.event_map[uid] = asyncio.Event() self.stats.on_new_user(uid) return uid + def request_epoch(self, uid: int) -> int: + return self.request_epochs.get(uid, self.backend_epoch) + + def wake_failed_restart_waiters(self) -> None: + if self._loop is not None and self._ready_event is not None: + self._loop.call_soon_threadsafe(self._ready_event.set) + async def listen(self): while True: msg = await self.recv_tokenizer.get() if isinstance(msg, CacheRebuildReply): + if msg.epoch != self.backend_epoch: + continue self._resolve_rebuild(msg) continue for msg in _unwrap_msg(msg): - # Global accounting follows actual admitted/sampled work even after the HTTP - # client disconnects and abort_user removes its ack queue. Delivery to a live - # request remains gated below, but observation must happen first. + # Retired-epoch output must not leak into current accounting. Current-epoch + # accounting intentionally survives a client disconnect: the scheduler did the + # work even if abort_user already removed this request's delivery queue. + if msg.epoch != self.backend_epoch: + continue self.stats.observe(msg) if msg.uid not in self.ack_map: continue + if msg.epoch != self.request_epochs.get(msg.uid): + continue self.ack_map[msg.uid].append(msg) self.event_map[msg.uid].set() @@ -318,9 +380,25 @@ def _resolve_all() -> None: # Loop already closed (shutdown racing the crash): nothing left to wake. pass + def fail_dispatched_for_oom(self, message: str) -> None: + """Terminally fail live requests once; no partial work is replayed into a new epoch.""" + if self._loop is None: + return + def _fail() -> None: + for uid, replies in list(self.ack_map.items()): + if any(reply.finished for reply in replies): + continue + replies.append(UserReply(uid=uid, incremental_output="", finished=True, + error=message, error_code="insufficient_workspace", epoch=self.request_epochs[uid])) + self.event_map[uid].set() + self._loop.call_soon_threadsafe(_fail) + def _create_listener_once(self): if not self.initialized: self._loop = asyncio.get_running_loop() + self._ready_event = asyncio.Event() + if self.maintenance_state == "serving": + self._ready_event.set() asyncio.create_task(self.listen()) self.initialized = True @@ -349,6 +427,7 @@ async def wait_for_ack(self, uid: int): finally: self.ack_map.pop(uid, None) self.event_map.pop(uid, None) + self.request_epochs.pop(uid, None) async def stream_generate(self, uid: int): async for ack in self.wait_for_ack(uid): @@ -381,7 +460,8 @@ async def abort_user(self, uid: int): del self.event_map[uid] self.stats.on_abort(uid) logger.warning("Aborting request for user %s", uid) - await self.send_one(AbortMsg(uid=uid)) + self.request_epochs.pop(uid, None) + await self.send_one(AbortMsg(uid=uid, epoch=self.backend_epoch)) def shutdown(self): self.send_tokenizer.stop() @@ -552,6 +632,7 @@ async def dispatch_rebuild( through ``POST /v1/cache/rebuild`` (``ft ctl cache``, the desktop panel, the shell's ``/cache``), which does the pre-flight maintenance_state checks (409/503 short-circuits).""" request_id = str(uuid.uuid4()) + epoch = state.backend_epoch fut = asyncio.get_running_loop().create_future() state.rebuild_futures[request_id] = fut state.maintenance_state = "rebuilding" @@ -564,6 +645,7 @@ async def dispatch_rebuild( num_mamba_slots=num_mamba_slots, num_swa_pages=num_swa_pages, mode=mode, + epoch=state.backend_epoch, ) ) except Exception as e: # noqa: BLE001 @@ -571,8 +653,14 @@ async def dispatch_rebuild( # untouched. Roll the gate back to serving (else a transient ZMQ error would latch # maintenance forever with no reply ever arriving to clear it) and surface the error. state.rebuild_futures.pop(request_id, None) - state.maintenance_state = "serving" + if state.backend_epoch == epoch and state.maintenance_state == "rebuilding": + state.maintenance_state = "serving" return {"status": "failed", "error": f"failed to dispatch rebuild: {e!r}"} + # send_one yields to the event loop. A restart/prepare-stop may have won while this + # rebuild was being queued; never leave its future registered or overwrite that gate. + if state.backend_epoch != epoch or state.maintenance_state != "rebuilding": + state.rebuild_futures.pop(request_id, None) + return {"status": "failed", "error": "cache rebuild superseded by engine lifecycle"} try: return await asyncio.wait_for(fut, timeout=timeout) except asyncio.TimeoutError: @@ -868,12 +956,15 @@ async def generate(req: GenerateRequest, request: Request): logger.debug("Received generate request %s", req) log_request("/generate", req, request) state = get_global_state() - if state.maintenance_state != "serving": + if state.maintenance_state not in {"serving", "restarting"}: detail = "model is still loading" if state.maintenance_state == "loading" else "cache rebuild in progress" return JSONResponse({"error": f"server unavailable: {detail}"}, status_code=503) if req.max_tokens < 1: return JSONResponse({"error": f"max_tokens must be at least 1, got {req.max_tokens}"}, status_code=400) - uid = state.new_user() + try: + uid = await state.new_user() + except AdmissionClosedError as exc: + return JSONResponse({"error": str(exc)}, status_code=503) await state.send_one( TokenizeMsg( uid=uid, @@ -882,6 +973,7 @@ async def generate(req: GenerateRequest, request: Request): ignore_eos=req.ignore_eos, max_tokens=req.max_tokens, ), + epoch=state.request_epoch(uid), ) ) @@ -1027,22 +1119,79 @@ def run_api_server(config: ServerArgs, start_backend: Callable[[], "Any"], run_s # setting _SHUTTING_DOWN) rather than relying on OS signal-delivery order. _GLOBAL_STATE.backend_processes = list(getattr(handle, "processes", None) or []) + def _on_loop(action: Callable[[], None], *, timeout: float = 5.0) -> None: + """Serialize supervisor lifecycle mutations with request coroutines. + + Before the listener has captured a loop (initial loading) this server has no admitted + work, so direct mutation is safe. Thereafter wait for the loop callback to complete; + retirement must not begin until its gate/future transition is visible to handlers. + """ + _run_on_loop(_GLOBAL_STATE._loop, action, timeout=timeout) + def _on_ready() -> None: # A stop requested while weights were loading has already sealed admission. The backend # may finish its ready handshake before SIGTERM arrives; never reopen that gate after the # daemon has received a final accounting snapshot. - if _GLOBAL_STATE.maintenance_state == "loading": - _GLOBAL_STATE.maintenance_state = "serving" - _GLOBAL_STATE.ready_at = time.monotonic() - logger.info(f"API server is ready to serve on {host}:{port}") + def _apply() -> None: + if _GLOBAL_STATE.maintenance_state in {"loading", "restarting"} and not _SHUTTING_DOWN.is_set(): + _GLOBAL_STATE.maintenance_state = "serving" + _GLOBAL_STATE.ready_at = time.monotonic() + logger.info(f"API server is ready to serve on {host}:{port}") + if _GLOBAL_STATE._ready_event is not None: + _GLOBAL_STATE._ready_event.set() + _on_loop(_apply) + + def _restart_backend(epoch: int): + fresh = start_backend(epoch) + _GLOBAL_STATE.backend_processes = list(getattr(fresh, "processes", None) or []) + return fresh + + def _on_restarting(message: str, attempt: int) -> None: + def _apply() -> None: + if _SHUTTING_DOWN.is_set() or _GLOBAL_STATE.maintenance_state == "stopping": + raise RuntimeError("OOM restart cancelled by shutdown") + _GLOBAL_STATE.maintenance_state = "restarting" + _GLOBAL_STATE.restart_attempts = attempt + _GLOBAL_STATE.backend_epoch += 1 + _GLOBAL_STATE.unit_bytes = None + _GLOBAL_STATE.free_vram_bytes = 0 + _GLOBAL_STATE.cache_floors = None + _GLOBAL_STATE.cache_pools = None + _GLOBAL_STATE.gpus = [] + _GLOBAL_STATE.last_rebuild = None + _GLOBAL_STATE.stats.clear_runtime_gauges() + if _GLOBAL_STATE._ready_event is not None: + _GLOBAL_STATE._ready_event.clear() + # Resolve all loop-owned state atomically with the epoch transition, before old + # workers are retired. No cache rebuild coroutine can re-open this gate afterward. + for uid, replies in list(_GLOBAL_STATE.ack_map.items()): + if not any(reply.finished for reply in replies): + replies.append(UserReply(uid=uid, incremental_output="", finished=True, + error="worker_restarting_after_cuda_oom: " + message, + error_code="insufficient_workspace", epoch=_GLOBAL_STATE.request_epochs[uid])) + _GLOBAL_STATE.event_map[uid].set() + for request_id, future in list(_GLOBAL_STATE.rebuild_futures.items()): + _GLOBAL_STATE.rebuild_futures.pop(request_id, None) + if not future.done(): + future.set_result({"status": "failed", "error": "worker_restarting_after_cuda_oom"}) + _on_loop(_apply) def _on_failure(message: str) -> None: - _GLOBAL_STATE.fatal_error = message - _GLOBAL_STATE.maintenance_state = "failed" + def _apply() -> None: + if _SHUTTING_DOWN.is_set() or _GLOBAL_STATE.maintenance_state == "stopping": + return + _GLOBAL_STATE.fatal_error = message + _GLOBAL_STATE.maintenance_state = "failed" + for request_id, future in list(_GLOBAL_STATE.rebuild_futures.items()): + _GLOBAL_STATE.rebuild_futures.pop(request_id, None) + if not future.done(): + future.set_result({"status": "failed", "error": message}) + if _GLOBAL_STATE._ready_event is not None: + _GLOBAL_STATE._ready_event.set() + _on_loop(_apply) logger.error("Backend supervisor: %s", message) # No CacheRebuildReply will ever arrive from a dead backend, so wake any caller blocked # in dispatch_rebuild's await now — otherwise it strands until the full rebuild timeout. - _GLOBAL_STATE.fail_pending_rebuilds(message) # Then take the whole serve down (see _exit_after_backend_death). Shell mode is excluded: # a person is sitting at that TUI, the API is theirs alone, and its stop path is ^C. if not run_shell: @@ -1053,14 +1202,16 @@ def _on_meta(meta: dict) -> None: # sizes allocated at load, delivered once on the ack path; surfaced by cache_geometry # (unit_bytes + the limits block + the pre-first-chat pool seed). Unpack the extras # aside so unit_bytes keeps its original three-key shape; unknown keys, if any, are inert. - meta = dict(meta or {}) - _GLOBAL_STATE.free_vram_bytes = int(meta.pop("free_vram_bytes", 0) or 0) - _GLOBAL_STATE.cache_floors = meta.pop("floors", None) - _GLOBAL_STATE.cache_pools = meta.pop("pools", None) - _GLOBAL_STATE.swa_full_tokens_ratio = float(meta.pop("swa_full_tokens_ratio", 0.0) or 0.0) - _GLOBAL_STATE.cache_budget_bytes = int(meta.pop("cache_budget_bytes", 0) or 0) - _GLOBAL_STATE.gpus = list(meta.pop("gpus", None) or []) - _GLOBAL_STATE.unit_bytes = meta + def _apply() -> None: + values = dict(meta or {}) + _GLOBAL_STATE.free_vram_bytes = int(values.pop("free_vram_bytes", 0) or 0) + _GLOBAL_STATE.cache_floors = values.pop("floors", None) + _GLOBAL_STATE.cache_pools = values.pop("pools", None) + _GLOBAL_STATE.swa_full_tokens_ratio = float(values.pop("swa_full_tokens_ratio", 0.0) or 0.0) + _GLOBAL_STATE.cache_budget_bytes = int(values.pop("cache_budget_bytes", 0) or 0) + _GLOBAL_STATE.gpus = list(values.pop("gpus", None) or []) + _GLOBAL_STATE.unit_bytes = values + _on_loop(_apply) # Early-bind: supervise the backend on a daemon thread so uvicorn can bind # immediately and /health can report loading progress. Shell mode wants exactly the same @@ -1075,6 +1226,11 @@ def _on_meta(meta: dict) -> None: # uvicorn's lifespan shutdown sets this on SIGINT/SIGTERM, so the workers' # expected exit during stop is not reported as a crash. "is_shutting_down": _SHUTTING_DOWN.is_set, + "restart_backend": _restart_backend if config.oom_recovery else None, + "on_restarting": _on_restarting, + "max_restarts": config.oom_recovery_max_restarts if config.oom_recovery else 0, + "restart_ready_timeout": config.oom_recovery_ready_timeout, + "can_restart": lambda: not _SHUTTING_DOWN.is_set() and _GLOBAL_STATE.maintenance_state != "stopping", }, name="freetoken-backend-supervisor", daemon=True, diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 7c0ed9b047..9c65aa55b5 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -2,6 +2,7 @@ import argparse import json +import math import os from dataclasses import dataclass, field from typing import List, Tuple @@ -73,6 +74,11 @@ class ServerArgs(SchedulerConfig): gpu: tuple[str, ...] = () # full UUIDs resolved from --gpu, entry i = TP rank i; None = NVML unavailable, each worker then resolves its raw entry against CUDA's own enumeration gpu_assigned: "tuple[str, ...] | None" = None + # Opt-in recovery of a CUDA worker that OOMs after the first ready handshake. + oom_recovery: bool = False + oom_recovery_max_restarts: int = 1 + oom_recovery_ready_timeout: float = 900.0 + oom_recovery_wait_queue: int = 128 @property def share_tokenizer(self) -> bool: @@ -162,6 +168,15 @@ def _positive_int(value: str) -> int: raise argparse.ArgumentTypeError("must be >= 1") return n + def _positive_finite_float(value: str) -> float: + try: + n = float(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be a positive finite number") from exc + if not math.isfinite(n) or n <= 0: + raise argparse.ArgumentTypeError("must be a positive finite number") + return n + def _lazy_gpu_arg(value: str) -> tuple[str, ...]: from freetoken.gpu_select import gpu_arg @@ -394,6 +409,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.num_tokenizer, help="The number of tokenizer processes to launch. 0 means the tokenizer is shared with the detokenizer.", ) + parser.add_argument("--oom-recovery", action="store_true", default=ServerArgs.oom_recovery, + help="Restart backend workers once after a post-ready CUDA OOM (default: off).") + parser.add_argument("--oom-recovery-max-restarts", type=_positive_int, + default=ServerArgs.oom_recovery_max_restarts, + help="Maximum CUDA-OOM backend restarts (default: 1).") + parser.add_argument("--oom-recovery-ready-timeout", type=_positive_finite_float, + default=ServerArgs.oom_recovery_ready_timeout, + help="Seconds new requests may wait for an OOM replacement (default: 900).") + parser.add_argument("--oom-recovery-wait-queue", type=_positive_int, + default=ServerArgs.oom_recovery_wait_queue, + help="Maximum requests waiting for OOM recovery readiness (default: 128).") parser.add_argument( "--max-prefill-length", diff --git a/python/freetoken/server/control_api.py b/python/freetoken/server/control_api.py index 7158e4e5fa..4405dd1585 100644 --- a/python/freetoken/server/control_api.py +++ b/python/freetoken/server/control_api.py @@ -17,8 +17,14 @@ def build_health(state: Any, version: str) -> dict: """Full-lifecycle health doc: loading -> ok -> error.""" instance_id = getattr(state, "instance_id", None) fatal = getattr(state, "fatal_error", None) + recovery = { + "maintenance": getattr(state, "maintenance_state", "serving"), + "backend_epoch": int(getattr(state, "backend_epoch", 0) or 0), + "restart_attempts": int(getattr(state, "restart_attempts", 0) or 0), + "restart_waiters": int(getattr(state, "_restart_waiters", 0) or 0), + } if fatal: - return {"status": "error", "message": fatal, "instance_id": instance_id} + return {"status": "error", "message": fatal, "instance_id": instance_id, **recovery} mstate = getattr(state, "maintenance_state", "serving") config = getattr(state, "config", None) @@ -35,6 +41,7 @@ def build_health(state: Any, version: str) -> dict: }, "model": model, "instance_id": instance_id, + **recovery, } ready_at = getattr(state, "ready_at", None) @@ -44,8 +51,8 @@ def build_health(state: Any, version: str) -> dict: "model": model, "instance_id": instance_id, "uptime_s": uptime_s, - "maintenance": mstate, "version": version, + **recovery, } diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 2e342ff899..a29b5d9b8b 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -299,7 +299,7 @@ async def submit_generation(spec: GenSpec, state: Any) -> int: calls this — it takes the neutral spec, not a wire request type.""" refs = collect_image_refs(spec.messages) images = await _resolve_images(refs, state) if refs else None - uid = state.new_user() + uid = await state.new_user() await state.send_one( TokenizeMsg( uid=uid, @@ -309,6 +309,7 @@ async def submit_generation(spec: GenSpec, state: Any) -> int: tools=spec.template_tools, images=images, inline_system_policy=spec.inline_system_policy, + epoch=getattr(state, "request_epoch", lambda _uid: 0)(uid), ) ) return uid diff --git a/python/freetoken/server/launch.py b/python/freetoken/server/launch.py index bb0c256f46..328cf3e846 100644 --- a/python/freetoken/server/launch.py +++ b/python/freetoken/server/launch.py @@ -15,13 +15,14 @@ from .supervisor import BackendHandle -def _report_startup_error(ack_queue: mp.Queue, exc: BaseException) -> None: +def _report_startup_error(ack_queue: mp.Queue, exc: BaseException, *, cuda_oom: bool = False) -> None: """Tell the parent WHY this worker is dying — push an ("error", reason) ack before it exits, so the supervisor reports the real cause (e.g. a config ValueError) instead of the generic "backend worker … exited during load". Best-effort and flushed (close + join_thread), since the process is about to terminate; a failure to report must never mask the original error.""" try: - ack_queue.put(("error", f"{type(exc).__name__}: {exc}")) + kind = "cuda_oom" if cuda_oom else "error" + ack_queue.put((kind, f"{type(exc).__name__}: {exc}")) ack_queue.close() ack_queue.join_thread() except Exception: # noqa: BLE001 -- reporting is a nicety; never shadow the real exception @@ -114,6 +115,11 @@ def _run_scheduler(args: ServerArgs, ack_queue: mp.Queue[str]) -> None: try: scheduler.run_forever() + except torch.cuda.OutOfMemoryError as exc: + # A runtime OOM poisons this scheduler generation. Tell the parent with a typed + # signal, then exit; it must retire every peer before fresh CUDA allocation. + _report_startup_error(ack_queue, exc, cuda_oom=True) + raise except KeyboardInterrupt: logger = init_logger(__name__) if args.tp_info.is_primary(): @@ -150,10 +156,11 @@ def launch_server( f"{', '.join(server_args.gpu_assigned) if server_args.gpu_assigned else 'resolved at CUDA init (no NVML)'}" ) - def start_subprocess() -> "BackendHandle": + def start_subprocess(backend_epoch: int = 0) -> "BackendHandle": import multiprocessing as mp from .supervisor import BackendHandle + from .supervisor import _retire_oom_workers mp.set_start_method("spawn", force=True) detach = server_args.shell_mode # see _detach_process_group @@ -162,16 +169,36 @@ def start_subprocess() -> "BackendHandle": ack_queue: mp.Queue = mp.Queue() processes: list[mp.Process] = [] + def _start(process: mp.Process) -> None: + try: + process.start() + except Exception as exc: + # start() can fail after assigning a PID. Retire that process too, and fail + # closed if any previously-created child cannot be reaped. + started = [*processes] + if process.pid is not None: + started.append(process) + try: + _retire_oom_workers(started) + except Exception as cleanup: + raise RuntimeError(f"backend factory cleanup failed: {cleanup}") from exc + try: + ack_queue.close() + ack_queue.join_thread() + except Exception: + pass + raise + processes.append(process) + for i in range(world_size): - new_args = replace(server_args, tp_info=DistributedInfo(i, world_size)) + new_args = replace(server_args, tp_info=DistributedInfo(i, world_size), backend_epoch=backend_epoch) p = mp.Process( target=_run_scheduler, args=(new_args, ack_queue), daemon=False, name=f"freetoken-TP{i}-scheduler", ) - p.start() - processes.append(p) + _start(p) num_tokenizers = server_args.num_tokenizer p = mp.Process( @@ -187,12 +214,12 @@ def start_subprocess() -> "BackendHandle": "create": server_args.tokenizer_create_addr, "tokenizer_id": num_tokenizers, "ack_queue": ack_queue, + "backend_epoch": backend_epoch, }, daemon=False, name="freetoken-detokenizer-0", ) - p.start() - processes.append(p) + _start(p) for i in range(num_tokenizers): p = mp.Process( target=_run_tokenize_worker, @@ -207,12 +234,12 @@ def start_subprocess() -> "BackendHandle": "create": server_args.tokenizer_create_addr, "tokenizer_id": i, "ack_queue": ack_queue, + "backend_epoch": backend_epoch, }, daemon=False, name=f"freetoken-tokenizer-{i}", ) - p.start() - processes.append(p) + _start(p) # Expected ready acks: 1 primary scheduler + num_tokenizers + 1 detokenizer. return BackendHandle( diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 99a4283549..d53e0d01ca 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -22,6 +22,7 @@ ) from .function_call_parser import ToolCallItem from .request_logger import log_request +from .accounting import AdmissionClosedError from .generation import ( DEFAULT_MAX_OUTPUT_TOKENS, ContentDelta, @@ -97,7 +98,7 @@ def _maintenance_gate(state: Any) -> JSONResponse | None: runtime cache "rebuild"/"failed" so clients (and the desktop) get an actionable message. None when serving.""" mstate = getattr(state, "maintenance_state", "serving") - if mstate == "serving": + if mstate in {"serving", "restarting"}: return None if mstate == "loading": msg = "model is still loading" @@ -200,6 +201,8 @@ async def handle_chat_completion( try: uid = await submit_generation(spec, state) + except AdmissionClosedError as exc: + return JSONResponse({"error": str(exc)}, status_code=503) except GenerationError as exc: return create_error_response(str(exc), code=exc.code) @@ -409,11 +412,14 @@ async def handle_completion( if req.stream: if len(prompts) != 1: return create_error_response("Streaming completions only support a single text prompt") - uid = state.new_user() + try: + uid = await state.new_user() + except AdmissionClosedError as exc: + return create_error_response(str(exc), code="server_restarting", status_code=503) await state.send_one( TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling( req, model_sampling, default_max_tokens=default_max_tokens - )) + ), epoch=getattr(state, "request_epoch", lambda _uid: 0)(uid)) ) chunks = stream_completion_chunks(uid, req, state) if request is not None: @@ -425,7 +431,10 @@ async def handle_completion( completion_tokens = 0 cached_tokens = 0 for index, prompt in enumerate(prompts): - uid = state.new_user() + try: + uid = await state.new_user() + except AdmissionClosedError as exc: + return create_error_response(str(exc), code="server_restarting", status_code=503) await state.send_one( TokenizeMsg( uid=uid, @@ -433,6 +442,7 @@ async def handle_completion( sampling_params=_resolve_sampling( req, model_sampling, default_max_tokens=default_max_tokens ), + epoch=getattr(state, "request_epoch", lambda _uid: 0)(uid), ) ) text = "" diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 71d6eed873..977c7e8702 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -25,6 +25,8 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse + +from .accounting import AdmissionClosedError from openai.types.responses import ( Response, ResponseCompletedEvent, @@ -118,7 +120,7 @@ async def v1_responses(req: ResponsesRequest, request: Request): log_request("/v1/responses", req, request) state = get_state() mstate = getattr(state, "maintenance_state", "serving") - if mstate != "serving": + if mstate not in {"serving", "restarting"}: detail = "model is still loading" if mstate == "loading" else "cache rebuild in progress" return _error_response(503, detail) if req.background: @@ -157,6 +159,8 @@ async def handle_responses( reasoning_parser=getattr(state.config, "reasoning_parser", None), ) uid = await submit_generation(spec, state) + except AdmissionClosedError as exc: + return _error_response(503, str(exc), "server_restarting") except GenerationError as exc: return _error_response(400, str(exc), exc.code) except ValueError as exc: diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 3d48ee96a0..051a6a5e04 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -55,6 +55,17 @@ def on_abort(self, uid: int) -> None: if uid in self._inflight: self._aborting.add(uid) + def clear_runtime_gauges(self) -> None: + """Forget worker-generation-specific gauges while a replacement loads. + + Lifetime request and token counters intentionally remain process-scoped; only values + sampled from the retired CUDA allocator/cache are cleared. + """ + self.kv_used_pages = self.kv_total_pages = 0 + self.mamba_used_slots = self.mamba_total_slots = 0 + self.swa_used_tokens = self.swa_total_tokens = 0 + self.vram_bytes = 0 + def observe(self, reply: Any, now: float | None = None) -> None: t = time.monotonic() if now is None else now if getattr(reply, "completion_tokens_delta", 0) > 0: @@ -165,6 +176,12 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "decode_tps": round(tr.decode_tps(), 1), "prefill_tps": round(tr.prefill_tps(), 1), }, + "recovery": { + "maintenance": getattr(state, "maintenance_state", "serving"), + "backend_epoch": int(getattr(state, "backend_epoch", 0) or 0), + "restart_attempts": int(getattr(state, "restart_attempts", 0) or 0), + "restart_waiters": int(getattr(state, "_restart_waiters", 0) or 0), + }, "requests": { "active": tr.active, "completed": tr.completed, diff --git a/python/freetoken/server/supervisor.py b/python/freetoken/server/supervisor.py index 322419a65f..24f18e1983 100644 --- a/python/freetoken/server/supervisor.py +++ b/python/freetoken/server/supervisor.py @@ -81,6 +81,17 @@ def _as_error(msg: Any) -> str | None: return None +def _as_cuda_oom(msg: Any) -> str | None: + """Only the worker's typed torch.cuda.OutOfMemoryError receipt enables recovery. + + A human exception string is deliberately insufficient: ordinary failures which happen to + contain 'OOM'/'out of memory' retain the existing terminal process-exit behavior. + """ + if isinstance(msg, tuple) and msg and msg[0] == "cuda_oom": + return str(msg[1]) if len(msg) > 1 else "CUDA out of memory" + return None + + def _drain_pending_error(get: Callable[[float], Any], attempts: int = 5, timeout: float = 0.02) -> str | None: """After a worker death, poll the ack queue briefly for a final ("error", reason) ack still in flight (the queue feeder flushes on process exit), so the real reason wins over the @@ -96,12 +107,47 @@ def _drain_pending_error(get: Callable[[float], Any], attempts: int = 5, timeout return None +def _drain_pending_failure(get: Callable[[float], Any]) -> tuple[str | None, bool]: + """Return the final worker reason plus whether it was a typed CUDA OOM.""" + for _ in range(5): + try: + msg = get(0.02) + except Empty: + continue + if (oom := _as_cuda_oom(msg)) is not None: + return oom, True + if (err := _as_error(msg)) is not None: + return err, False + return None, False + + +def _retire_oom_workers(processes: List[Any]) -> None: + """Terminate, join, and if necessary kill every failed worker before replacement CUDA use.""" + survivors = [] + for process in processes: + try: + if process.is_alive(): + process.terminate() + process.join(timeout=10) + if process.is_alive(): + process.kill() + process.join(timeout=2) + if process.is_alive(): + survivors.append(getattr(process, "name", "unknown")) + except Exception: # noqa: BLE001 + survivors.append(getattr(process, "name", "unknown")) + if survivors: + raise WorkerDied("OOM workers could not be stopped: " + ", ".join(survivors)) + + def drain_ready( handle: BackendHandle, progress: LoadProgress, get: Callable[[float], Any] | None = None, poll: float = 1.0, on_meta: Callable[[dict], None] | None = None, + timeout_s: float | None = None, + is_shutting_down: Callable[[], bool] | None = None, ) -> None: """Block until ``expected_acks`` non-progress acks arrive, applying ("progress", …) tuples to ``progress``. Between acks, poll worker liveness; raise WorkerDied if a worker @@ -115,7 +161,12 @@ def get(timeout: float) -> Any: return handle.ack_queue.get(timeout=timeout) ready = 0 + deadline = time.monotonic() + timeout_s if timeout_s is not None else None while ready < handle.expected_acks: + if is_shutting_down is not None and is_shutting_down(): + raise WorkerDied("backend startup cancelled by shutdown") + if deadline is not None and time.monotonic() >= deadline: + raise WorkerDied("backend readiness timed out") try: msg = get(poll) except Empty: @@ -129,6 +180,11 @@ def get(timeout: float) -> Any: err = _as_error(msg) if err is not None: raise WorkerDied(err) + oom = _as_cuda_oom(msg) + if oom is not None: + # Recovery is deliberately post-initial-ready only. A load-time CUDA OOM has no + # healthy generation to retire and follows the ordinary startup-failure path. + raise WorkerDied(oom) if isinstance(msg, tuple) and msg and msg[0] == "progress": _tag, desc, done, total = msg progress.update(desc, done, total) @@ -147,6 +203,11 @@ def run_backend_supervisor( poll: float = 1.0, on_meta: Callable[[dict], None] | None = None, is_shutting_down: Callable[[], bool] | None = None, + restart_backend: Callable[[int], BackendHandle] | None = None, + on_restarting: Callable[[str, int], None] | None = None, + max_restarts: int = 1, + restart_ready_timeout: float | None = None, + can_restart: Callable[[], bool] | None = None, ) -> None: """Drain to readiness (watching liveness), flip the gate via ``on_ready``, then watch forever: a worker that dies AFTER ready (a scheduler crash while the frontend keeps @@ -164,6 +225,9 @@ def run_backend_supervisor( def _shutting_down() -> bool: return bool(is_shutting_down is not None and is_shutting_down()) + def _can_restart() -> bool: + return not _shutting_down() and (can_restart is None or can_restart()) + try: drain_ready(handle, progress, poll=poll, on_meta=on_meta) except WorkerDied as exc: @@ -174,13 +238,48 @@ def _shutting_down() -> bool: on_ready() + attempts = 0 while True: dead = _first_dead(handle.processes) if dead is not None: if _shutting_down(): # Orderly stop in progress: the worker exit is expected — stay quiet. return - if on_failure is not None: - on_failure(f"backend worker {getattr(dead, 'name', '?')} exited") - return + reason, is_cuda_oom = _drain_pending_failure(lambda timeout: handle.ack_queue.get(timeout=timeout)) + message = reason or f"backend worker {getattr(dead, 'name', '?')} exited" + if not is_cuda_oom or restart_backend is None or attempts >= max_restarts: + if on_failure is not None: + on_failure(message) + return + if not _can_restart(): + return + attempts += 1 + if on_restarting is not None: + try: + on_restarting(message, attempts) + except Exception as exc: # noqa: BLE001 + if _can_restart() and on_failure is not None: + on_failure(f"OOM recovery transition failed: {exc}") + return + fresh = None + try: + _retire_oom_workers(handle.processes) + if not _can_restart(): + return + fresh = restart_backend(attempts) + drain_ready(fresh, progress, poll=poll, on_meta=on_meta, + timeout_s=restart_ready_timeout, is_shutting_down=lambda: not _can_restart()) + except Exception as exc: # noqa: BLE001 + failure = f"OOM recovery failed: {exc}" + if fresh is not None: + try: + _retire_oom_workers(fresh.processes) + except WorkerDied as cleanup: + failure += f"; {cleanup.message}" + if not _shutting_down() and on_failure is not None: + on_failure(failure) + return + handle = fresh + on_ready() + continue time.sleep(poll) diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index d3365bfecd..58f3722258 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -48,12 +48,13 @@ def _prompt_admitted_reply(msg: PromptAdmittedMsg) -> UserReply: finished=False, prompt_tokens_delta=msg.prompt_tokens, cached_tokens=msg.cached_tokens, + epoch=msg.epoch, ) def _error_reply(msg: ErrorReplyMsg) -> UserReply: return UserReply( - uid=msg.uid, incremental_output="", finished=True, error=msg.error, error_code=msg.code, + uid=msg.uid, incremental_output="", finished=True, error=msg.error, error_code=msg.code, epoch=msg.epoch, ) @@ -104,6 +105,7 @@ def _tokenize_requests( incremental_output="", finished=True, error=f"could not encode request: {exc}", + epoch=msg.epoch, ) ) continue @@ -116,6 +118,7 @@ def _tokenize_requests( incremental_output="", finished=True, error="prompt must contain at least one token", + epoch=msg.epoch, ) ) continue @@ -136,6 +139,7 @@ def tokenize_worker( model_source: str = "huggingface", ack_queue: mp.Queue[str] | None = None, mm: MultimodalConfig | None = None, + backend_epoch: int = 0, ) -> None: send_backend = ZmqPushQueue(backend_addr, create=False, encoder=BaseBackendMsg.encoder) send_frontend = ZmqPushQueue(frontend_addr, create=False, encoder=BaseFrontendMsg.encoder) @@ -162,6 +166,9 @@ def tokenize_worker( pending_msg = _unwrap_msg(recv_listener.get()) while len(pending_msg) < local_bs and not recv_listener.empty(): pending_msg.extend(_unwrap_msg(recv_listener.get())) + # The frontend transport stays up across an OOM restart. Drop anything buffered + # for the retired generation before it reaches tokenization or scheduler state. + pending_msg = [m for m in pending_msg if getattr(m, "epoch", backend_epoch) == backend_epoch] logger.debug(f"Received {len(pending_msg)} messages") @@ -182,6 +189,7 @@ def tokenize_worker( num_mamba_slots=m.num_mamba_slots, num_swa_pages=m.num_swa_pages, mode=m.mode, + epoch=m.epoch, ) ) elif isinstance(m, CacheRebuildResultMsg): @@ -194,6 +202,7 @@ def tokenize_worker( mamba_slots=m.mamba_slots, num_swa_pages=m.num_swa_pages, error=m.error, + epoch=m.epoch, ) ) n_control = sum( @@ -225,6 +234,7 @@ def tokenize_worker( swa_used_tokens=msg.swa_used_tokens, swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, + epoch=msg.epoch, ) for msg, reply in zip(detokenize_msg, replies, strict=True) ] @@ -256,7 +266,7 @@ def tokenize_worker( send_backend.put(backend[0] if len(backend) == 1 else BatchBackendMsg(data=backend)) if len(abort_msg) > 0: batch_output = BatchBackendMsg( - data=[AbortBackendMsg(uid=msg.uid) for msg in abort_msg] + data=[AbortBackendMsg(uid=msg.uid, epoch=msg.epoch) for msg in abort_msg] ) if len(batch_output.data) == 1: batch_output = batch_output.data[0] diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 772472573a..ba048a2d7d 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -92,11 +92,12 @@ def tokenize(self, msgs: List[TokenizeMsg]) -> List[UserMsg]: mm_items=mm.mm_items, mrope_positions=mm.mrope_positions, mrope_delta=mm.mrope_delta, + epoch=msg.epoch, ) ) else: results.append( - UserMsg(uid=msg.uid, input_ids=input_ids, sampling_params=msg.sampling_params) + UserMsg(uid=msg.uid, input_ids=input_ids, sampling_params=msg.sampling_params, epoch=msg.epoch) ) return results diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index 85c217859f..5e584ca1a4 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -60,7 +60,7 @@ def _setup(): finished_reqs=set(), eos_token_ids=set(), toolcall_anchor_id=None, - config=SimpleNamespace(page_size=1), + config=SimpleNamespace(page_size=1, backend_epoch=0), status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None), send_result=sent.extend, _kv_usage_pages=cm.page_usage, diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index 045010c6b4..8c121749dc 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -91,6 +91,7 @@ def test_prompt_admitted_signal_uses_existing_frontend_usage_channel(): def test_schedule_reports_admission_only_after_prepare_succeeds(): batch = SimpleNamespace(is_prefill=True, prompt_admissions=[(1, 12, 4), (2, 34, 0)]) scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=0) scheduler.prefill_budget = 99 scheduler.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: batch) scheduler.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None) @@ -115,6 +116,7 @@ def send(messages): def test_prepare_failure_emits_no_prompt_admission(): batch = SimpleNamespace(is_prefill=True, prompt_admissions=[(1, 12, 0)]) scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=0) scheduler.prefill_budget = 99 scheduler.prefill_manager = SimpleNamespace(schedule_next_batch=lambda budget: batch) scheduler.decode_manager = SimpleNamespace(schedule_next_batch=lambda: None) @@ -132,6 +134,7 @@ def fail_prepare(_batch): def test_scheduler_rejection_emits_error_but_no_admission(): scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=0) scheduler.engine = SimpleNamespace(max_seq_len=4) added = [] scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append) @@ -153,6 +156,7 @@ def test_scheduler_rejection_emits_error_but_no_admission(): def test_scheduler_always_emits_terminal_abort_ack_for_unknown_uid(): scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=0) scheduler.prefill_manager = SimpleNamespace(abort_req=lambda uid: None) scheduler.decode_manager = SimpleNamespace(abort_req=lambda uid: None) scheduler._pending_abort_acks = set() @@ -167,8 +171,36 @@ def test_scheduler_always_emits_terminal_abort_ack_for_unknown_uid(): assert sent[0].uid == 77 and sent[0].error == "request aborted" +def test_scheduler_epoch_filters_stale_input_and_stamps_admission_and_abort_replies(): + scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=3) + scheduler.engine = SimpleNamespace(max_seq_len=64) + added, sent = [], [] + scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append, abort_req=lambda uid: None) + scheduler.decode_manager = SimpleNamespace(abort_req=lambda uid: None) + scheduler._pending_abort_acks = set() + scheduler.send_result = sent.extend + + Scheduler._process_one_msg( + scheduler, + UserMsg(uid=1, input_ids=torch.arange(4, dtype=torch.int32), + sampling_params=SamplingParams(max_tokens=1), epoch=2), + ) + assert added == [] and sent == [] + + Scheduler._report_prompt_admissions( + scheduler, SimpleNamespace(is_prefill=True, prompt_admissions=[(2, 9, 3)]) + ) + scheduler._pending_abort_acks.add(3) + Scheduler._flush_abort_acks(scheduler) + assert [(type(message), message.epoch) for message in sent] == [ + (PromptAdmittedMsg, 3), (ErrorReplyMsg, 3), + ] + + def test_abort_before_cross_worker_user_message_cannot_resurrect_request(): scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=0) scheduler.engine = SimpleNamespace(max_seq_len=64) added = [] scheduler.prefill_manager = SimpleNamespace( @@ -200,6 +232,7 @@ def test_abort_before_cross_worker_user_message_cannot_resurrect_request(): def test_normal_loop_sends_prior_sample_before_abort_terminal(): scheduler = Scheduler.__new__(Scheduler) + scheduler.config = SimpleNamespace(backend_epoch=0) scheduler.prefill_manager = SimpleNamespace(runnable=False, abort_req=lambda uid: None) scheduler.decode_manager = SimpleNamespace(runnable=False, abort_req=lambda uid: None) scheduler._pending_abort_acks = set() diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py index f382093fc0..1ede1e060c 100644 --- a/tests/server/test_anthropic_api.py +++ b/tests/server/test_anthropic_api.py @@ -492,7 +492,7 @@ def __init__(self, outputs): def frontend_tokenizer(self): return self._count_manager - def new_user(self): + async def new_user(self): self._uid += 1 return self._uid diff --git a/tests/server/test_effort_dialect.py b/tests/server/test_effort_dialect.py index a25de8e728..b427352531 100644 --- a/tests/server/test_effort_dialect.py +++ b/tests/server/test_effort_dialect.py @@ -39,7 +39,7 @@ def __init__(self, reasoning_parser: str | None = None) -> None: ) self.sent: TokenizeMsg | None = None - def new_user(self) -> int: + async def new_user(self) -> int: return 42 async def send_one(self, msg): diff --git a/tests/server/test_generation_accounting.py b/tests/server/test_generation_accounting.py index a6d2bdb916..9bb5d5e70d 100644 --- a/tests/server/test_generation_accounting.py +++ b/tests/server/test_generation_accounting.py @@ -54,7 +54,7 @@ def __init__(self, replies: list[UserReply]) -> None: ) self._replies = replies - def new_user(self) -> int: + async def new_user(self) -> int: return 42 async def wait_for_ack(self, uid: int): diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py index fdc095921d..d0c3861002 100644 --- a/tests/server/test_message_wire.py +++ b/tests/server/test_message_wire.py @@ -41,40 +41,40 @@ def test_inline_system_policy_tokenizer_roundtrip(): def test_cache_rebuild_msg_roundtrip(): - msg = CacheRebuildMsg(request_id="abc", moe_cache_size=8, num_pages=1024, mode="if_idle") + msg = CacheRebuildMsg(request_id="abc", moe_cache_size=8, num_pages=1024, mode="if_idle", epoch=9) out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) assert isinstance(out, CacheRebuildMsg) - assert (out.request_id, out.moe_cache_size, out.num_pages, out.mode) == ("abc", 8, 1024, "if_idle") + assert (out.request_id, out.moe_cache_size, out.num_pages, out.mode, out.epoch) == ("abc", 8, 1024, "if_idle", 9) def test_cache_rebuild_backend_msg_roundtrip(): - msg = CacheRebuildBackendMsg(request_id="r1", moe_cache_size=None, num_pages=256, mode="drain") + msg = CacheRebuildBackendMsg(request_id="r1", moe_cache_size=None, num_pages=256, mode="drain", epoch=9) out = BaseBackendMsg.decoder(msg.encoder()) assert isinstance(out, CacheRebuildBackendMsg) - assert (out.request_id, out.moe_cache_size, out.num_pages, out.mode) == ("r1", None, 256, "drain") + assert (out.request_id, out.moe_cache_size, out.num_pages, out.mode, out.epoch) == ("r1", None, 256, "drain", 9) def test_cache_rebuild_result_msg_roundtrip(): - msg = CacheRebuildResultMsg(request_id="r2", status="ok", moe_cache_size=16, num_pages=512) + msg = CacheRebuildResultMsg(request_id="r2", status="ok", moe_cache_size=16, num_pages=512, epoch=9) out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) assert isinstance(out, CacheRebuildResultMsg) - assert (out.request_id, out.status, out.moe_cache_size, out.num_pages, out.error) == ( - "r2", "ok", 16, 512, None, + assert (out.request_id, out.status, out.moe_cache_size, out.num_pages, out.error, out.epoch) == ( + "r2", "ok", 16, 512, None, 9, ) def test_cache_rebuild_reply_roundtrip(): - msg = CacheRebuildReply(request_id="r3", status="failed", error="boom") + msg = CacheRebuildReply(request_id="r3", status="failed", error="boom", epoch=9) out = BaseFrontendMsg.decoder(BaseFrontendMsg.encoder(msg)) assert isinstance(out, CacheRebuildReply) - assert (out.request_id, out.status, out.error) == ("r3", "failed", "boom") + assert (out.request_id, out.status, out.error, out.epoch) == ("r3", "failed", "boom", 9) def test_prompt_admitted_msg_roundtrip(): - msg = PromptAdmittedMsg(uid=42, prompt_tokens=1234, cached_tokens=500) + msg = PromptAdmittedMsg(uid=42, prompt_tokens=1234, cached_tokens=500, epoch=9) out = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) assert isinstance(out, PromptAdmittedMsg) - assert (out.uid, out.prompt_tokens, out.cached_tokens) == (42, 1234, 500) + assert (out.uid, out.prompt_tokens, out.cached_tokens, out.epoch) == (42, 1234, 500, 9) def test_user_reply_token_deltas_round_trip(): @@ -88,6 +88,7 @@ def test_user_reply_token_deltas_round_trip(): kv_used_pages=40, kv_total_pages=512, gpu_mem_bytes=64 * (1 << 30), + epoch=9, ) decoded = BaseFrontendMsg.decoder(BaseFrontendMsg.encoder(msg)) @@ -102,6 +103,7 @@ def test_user_reply_token_deltas_round_trip(): assert decoded.kv_used_pages == 40 assert decoded.kv_total_pages == 512 assert decoded.gpu_mem_bytes == 64 * (1 << 30) + assert decoded.epoch == 9 def test_detokenize_msg_carries_kv_usage_round_trip(): @@ -110,12 +112,14 @@ def test_detokenize_msg_carries_kv_usage_round_trip(): kv_used_pages=10, kv_total_pages=256, gpu_mem_bytes=1 << 30, mamba_used_slots=7, mamba_total_slots=64, swa_used_tokens=8448, swa_total_tokens=76800, + epoch=9, ) decoded = BaseTokenizerMsg.decoder(BaseTokenizerMsg.encoder(msg)) assert isinstance(decoded, DetokenizeMsg) assert (decoded.kv_used_pages, decoded.kv_total_pages, decoded.gpu_mem_bytes) == (10, 256, 1 << 30) assert (decoded.mamba_used_slots, decoded.mamba_total_slots) == (7, 64) assert (decoded.swa_used_tokens, decoded.swa_total_tokens) == (8448, 76800) + assert decoded.epoch == 9 def test_client_dicts_with_the_wire_tag_key_survive_intact(): diff --git a/tests/server/test_oom_recovery.py b/tests/server/test_oom_recovery.py new file mode 100644 index 0000000000..9ff00d8bcb --- /dev/null +++ b/tests/server/test_oom_recovery.py @@ -0,0 +1,377 @@ +from __future__ import annotations + +import asyncio +import queue +import threading +from types import SimpleNamespace + +import pytest + +from freetoken.message import UserReply +from freetoken.server import api_server +from freetoken.server.accounting import AdmissionClosedError +from freetoken.server.api_server import FrontendManager + + +class _AsyncQueue: + def __init__(self) -> None: + self.q: asyncio.Queue = asyncio.Queue() + + async def get(self): + return await self.q.get() + + async def put(self, item) -> None: + await self.q.put(item) + + +def _state(*, maintenance_state: str = "serving") -> FrontendManager: + recv, send = _AsyncQueue(), _AsyncQueue() + return FrontendManager( + config=SimpleNamespace( + served_model_name="test", + model_path="/models/test", + tool_call_parser="llama3", + reasoning_parser=None, + mm=SimpleNamespace(text_model_only=False, disabled_encoders=frozenset()), + max_output_tokens=None, + oom_recovery_wait_queue=2, + oom_recovery_ready_timeout=0.02, + ), + send_tokenizer=send, + recv_tokenizer=recv, + maintenance_state=maintenance_state, + ) + + +def test_restarting_admission_is_bounded_cancellable_and_stamped_after_ready(monkeypatch): + async def scenario() -> None: + state = _state(maintenance_state="restarting") + state._loop = asyncio.get_running_loop() + state._ready_event = asyncio.Event() + + with pytest.raises(AdmissionClosedError, match="timed out"): + await state.new_user() + assert state._restart_waiters == 0 + + state._restart_waiters = state.config.oom_recovery_wait_queue + with pytest.raises(AdmissionClosedError, match="queue is full"): + await state.new_user() + state._restart_waiters = 0 + + cancelled = asyncio.create_task(state.new_user()) + await asyncio.sleep(0) + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + assert state._restart_waiters == 0 + + admitted = asyncio.create_task(state.new_user()) + await asyncio.sleep(0) + state.backend_epoch = 4 + state.maintenance_state = "serving" + state._ready_event.set() + uid = await admitted + assert state.request_epoch(uid) == 4 + assert state._restart_waiters == 0 + + state.maintenance_state = "restarting" + state._ready_event = asyncio.Event() + failed = asyncio.create_task(state.new_user()) + await asyncio.sleep(0) + state.maintenance_state = "failed" + state.wake_failed_restart_waiters() + with pytest.raises(AdmissionClosedError, match="failed"): + await failed + + asyncio.run(scenario()) + + +def test_oom_terminalizes_every_dispatched_stream_once_and_never_replays_it(): + async def scenario() -> None: + state = _state() + state._loop = asyncio.get_running_loop() + first, second = await state.new_user(), await state.new_user() + + state.fail_dispatched_for_oom("CUDA out of memory") + state.fail_dispatched_for_oom("second notification") + await asyncio.sleep(0) + + for uid in (first, second): + terminal = state.ack_map[uid] + assert len(terminal) == 1 + assert terminal[0].finished is True + assert terminal[0].error_code == "insufficient_workspace" + assert terminal[0].epoch == state.request_epoch(uid) + + asyncio.run(scenario()) + + +def test_listener_drops_stale_generation_before_delivery_or_global_stats(): + class _Stats: + def __init__(self) -> None: + self.observed: list[UserReply] = [] + + def on_new_user(self, uid: int) -> None: + pass + + def observe(self, reply: UserReply) -> None: + self.observed.append(reply) + + async def scenario() -> None: + state = _state() + state.stats = _Stats() + state._create_listener_once() + state.backend_epoch = 7 + uid = await state.new_user() + await state.recv_tokenizer.put(UserReply(uid=uid, incremental_output="old", finished=True, epoch=6)) + await asyncio.sleep(0) + assert state.ack_map[uid] == [] + assert state.stats.observed == [] + + fresh = UserReply(uid=uid, incremental_output="new", finished=True, epoch=7) + await state.recv_tokenizer.put(fresh) + await asyncio.sleep(0) + assert state.ack_map[uid] == [fresh] + assert state.stats.observed == [fresh] + + # An abort removes delivery state but cannot erase work already performed by the + # current backend; count it. A retired backend's unknown-uid reply is different and + # must remain invisible to global stats. + disconnected = UserReply(uid=999, incremental_output="late", finished=True, epoch=7) + stale_unknown = UserReply(uid=998, incremental_output="old late", finished=True, epoch=6) + await state.recv_tokenizer.put(disconnected) + await state.recv_tokenizer.put(stale_unknown) + await asyncio.sleep(0) + assert state.stats.observed == [fresh, disconnected] + + for task in asyncio.all_tasks(): + if task is not asyncio.current_task() and task.get_coro().__name__ == "listen": + task.cancel() + + asyncio.run(scenario()) + + +def test_recovery_defaults_remain_disabled_and_limits_are_validated(monkeypatch): + from freetoken.server.args import parse_args + + monkeypatch.delenv("FREETOKEN_API_KEY", raising=False) + defaults, _ = parse_args(["--model", "/models/test", "--dtype", "bfloat16"]) + assert defaults.oom_recovery is False + assert defaults.oom_recovery_max_restarts == 1 + assert defaults.oom_recovery_ready_timeout == 900.0 + assert defaults.oom_recovery_wait_queue == 128 + + configured, _ = parse_args([ + "--model", "/models/test", "--dtype", "bfloat16", "--oom-recovery", "--oom-recovery-max-restarts", "2", + "--oom-recovery-ready-timeout", "1.5", "--oom-recovery-wait-queue", "3", + ]) + assert (configured.oom_recovery, configured.oom_recovery_max_restarts, + configured.oom_recovery_ready_timeout, configured.oom_recovery_wait_queue) == (True, 2, 1.5, 3) + + with pytest.raises(SystemExit): + parse_args(["--model", "/models/test", "--dtype", "bfloat16", "--oom-recovery-max-restarts", "0"]) + with pytest.raises(SystemExit): + parse_args(["--model", "/models/test", "--dtype", "bfloat16", "--oom-recovery-wait-queue", "0"]) + for invalid in ("0", "-1", "inf", "nan"): + with pytest.raises(SystemExit): + parse_args(["--model", "/models/test", "--dtype", "bfloat16", "--oom-recovery-ready-timeout", invalid]) + + +def test_restart_gauge_reset_preserves_lifetime_request_and_token_accounting(): + from freetoken.server.stats import StatsTracker + + stats = StatsTracker() + stats.on_new_user(7) + stats.observe(UserReply( + uid=7, incremental_output="x", finished=True, prompt_tokens_delta=3, + completion_tokens_delta=5, kv_used_pages=2, kv_total_pages=8, + mamba_used_slots=1, mamba_total_slots=4, swa_used_tokens=10, + swa_total_tokens=20, gpu_mem_bytes=123, + )) + + stats.clear_runtime_gauges() + + assert stats.completed == 1 + assert stats.prompt_tokens_total == 3 + assert stats.completion_tokens_total == 5 + assert (stats.kv_used_pages, stats.kv_total_pages, stats.mamba_used_slots, + stats.mamba_total_slots, stats.swa_used_tokens, stats.swa_total_tokens, + stats.vram_bytes) == (0, 0, 0, 0, 0, 0, 0) + + +def test_rebuild_yield_racing_epoch_transition_cannot_register_or_reopen_after_restart(): + from freetoken.server.api_server import dispatch_rebuild + + async def scenario() -> None: + entered, release = asyncio.Event(), asyncio.Event() + sent = [] + + async def send_one(message): + sent.append(message) + entered.set() + await release.wait() + + state = SimpleNamespace( + backend_epoch=0, maintenance_state="serving", rebuild_futures={}, send_one=send_one + ) + task = asyncio.create_task(dispatch_rebuild(state, moe_cache_size=8, num_pages=None)) + await entered.wait() + assert len(state.rebuild_futures) == 1 + assert sent[0].epoch == 0 + + # This is the loop-serialized OOM transition while dispatch_rebuild is suspended in + # its transport await: it owns the gate and resolves/removes existing waiters first. + state.backend_epoch = 1 + state.maintenance_state = "restarting" + state.rebuild_futures.clear() + release.set() + result = await task + + assert result == {"status": "failed", "error": "cache rebuild superseded by engine lifecycle"} + assert state.rebuild_futures == {} + assert state.maintenance_state == "restarting" + assert [message.epoch for message in sent] == [0] + + asyncio.run(scenario()) + + +def test_post_restart_tokenize_error_is_terminal_in_replacement_epoch(): + from freetoken.core import SamplingParams + from freetoken.message import TokenizeMsg + from freetoken.tokenizer.server import _tokenize_requests + + class BrokenTokenizer: + def tokenize(self, messages): + raise ValueError("template rejected request") + + class Logger: + def warning(self, message): + pass + + _, errors = _tokenize_requests( + BrokenTokenizer(), + [TokenizeMsg(uid=12, text="prompt", sampling_params=SamplingParams(), epoch=1)], + Logger(), + ) + + assert len(errors) == 1 + assert errors[0].finished is True + assert errors[0].epoch == 1 + assert "could not encode request" in errors[0].error + + +def test_prepare_stop_while_restart_loading_keeps_admission_closed(): + async def scenario() -> None: + state = _state(maintenance_state="restarting") + state._loop = asyncio.get_running_loop() + state._ready_event = asyncio.Event() + waiter = asyncio.create_task(state.new_user()) + await asyncio.sleep(0) + + # prepare-stop owns the lifecycle state before replacement readiness races in. + state.maintenance_state = "stopping" + state._ready_event.set() + with pytest.raises(AdmissionClosedError, match="stopping"): + await waiter + with pytest.raises(AdmissionClosedError, match="stopping"): + await state.new_user() + assert state.uid_counter == 0 + + asyncio.run(scenario()) + + +def test_timed_out_queued_lifecycle_callback_cannot_mutate_after_loop_releases(): + """A supervisor timeout may cancel only an unclaimed loop callback. When that + callback finally runs it must observe cancellation and leave lifecycle state intact.""" + import concurrent.futures + + from freetoken.server.api_server import _run_on_loop + + class DelayedLoop: + callback = None + + def call_soon_threadsafe(self, callback): + self.callback = callback + + loop = DelayedLoop() + state = {"maintenance": "restarting"} + outcome: list[BaseException] = [] + + def caller() -> None: + try: + _run_on_loop(loop, lambda: state.__setitem__("maintenance", "serving"), timeout=0.01) + except BaseException as exc: + outcome.append(exc) + + thread = threading.Thread(target=caller) + thread.start() + thread.join(timeout=1) + assert not thread.is_alive() + assert isinstance(outcome[0], concurrent.futures.TimeoutError) + assert state["maintenance"] == "restarting" + + loop.callback() + assert state["maintenance"] == "restarting" + + +class _ThreadQueue: + """CPU-only bridge between a TestClient request thread and FrontendManager's loop.""" + def __init__(self) -> None: + self.q: queue.Queue = queue.Queue() + + async def get(self): + # Never delegate a permanent Queue.get() to the default executor: pytest's + # asyncio.run teardown waits for executor workers, which would hang this CPU test. + while True: + try: + return self.q.get_nowait() + except queue.Empty: + await asyncio.sleep(0.001) + + async def put(self, item) -> None: + self.q.put(item) + + +def test_four_protocol_routes_dispatch_after_restart_at_replacement_epoch(): + from fastapi import FastAPI + from fastapi.testclient import TestClient + from freetoken.server.anthropic_api import register_anthropic_routes + from freetoken.server.openai_api import register_openai_routes + from freetoken.server.responses_api import register_responses_routes + + requests = [ + ("/v1/chat/completions", {"model": "test", "messages": [{"role": "user", "content": "hi"}]}), + ("/v1/completions", {"model": "test", "prompt": "hi"}), + ("/v1/messages", {"model": "test", "max_tokens": 1, "messages": [{"role": "user", "content": "hi"}]}), + ("/v1/responses", {"model": "test", "input": "hi"}), + ] + for path, payload in requests: + recv, send = _ThreadQueue(), _ThreadQueue() + state = _state(maintenance_state="restarting") + state.recv_tokenizer = recv + state.send_tokenizer = send + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + register_anthropic_routes(app, lambda: state, lambda: {}) + register_responses_routes(app, lambda: state, lambda: {}) + with TestClient(app) as client: + result: list = [] + thread = threading.Thread(target=lambda: result.append(client.post(path, json=payload))) + thread.start() + for _ in range(100): + if state._ready_event is not None: + break + threading.Event().wait(0.01) + assert state._ready_event is not None, path + assert send.q.empty(), path + state.backend_epoch = 1 + state.maintenance_state = "serving" + # The waiting admission belongs to TestClient's event-loop thread. + # Cross-thread Event.set is not a valid asyncio synchronization primitive. + state._ready_event._loop.call_soon_threadsafe(state._ready_event.set) + tokenized = send.q.get(timeout=2) + assert tokenized.epoch == 1 + recv.q.put(UserReply(uid=tokenized.uid, incremental_output="ok", finished=True, epoch=1)) + thread.join(timeout=2) + assert not thread.is_alive(), path + assert result[0].status_code == 200, (path, result[0].text) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index e33018e195..54dd845275 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -39,7 +39,7 @@ def __init__( self.replies = replies self.sent: TokenizeMsg | None = None - def new_user(self) -> int: + async def new_user(self) -> int: return 42 async def send_one(self, msg): diff --git a/tests/server/test_prepare_stop_accounting.py b/tests/server/test_prepare_stop_accounting.py index 1e7d5664ae..1ebbc90ad9 100644 --- a/tests/server/test_prepare_stop_accounting.py +++ b/tests/server/test_prepare_stop_accounting.py @@ -125,6 +125,16 @@ def test_loading_engine_can_seal_zero_without_being_reopened(): assert state.maintenance_state == "stopping" +def test_prepare_stop_closes_a_restarting_engine_before_replacement_ready(): + state = _state(maintenance="restarting") + + result = asyncio.run(prepare_stop_accounting(state)) + + assert result["drain_complete"] is True + assert state.maintenance_state == "stopping" + assert state.stats.active == 0 + + def test_frontend_new_user_refuses_work_after_stop_gate_closes(): manager = FrontendManager( config=SimpleNamespace(served_model_name="model-a"), @@ -133,7 +143,7 @@ def test_frontend_new_user_refuses_work_after_stop_gate_closes(): maintenance_state="stopping", ) with pytest.raises(AdmissionClosedError, match="stopping"): - manager.new_user() + asyncio.run(manager.new_user()) assert manager.stats.active == 0 diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e56..a27280efc7 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -40,6 +40,7 @@ class _FakeState: to an injected impl (so a test can make the enqueue succeed or raise).""" def __init__(self, send_impl, *, maintenance_state="serving", fatal_error=None): + self.backend_epoch = 0 self.rebuild_futures: dict = {} self.maintenance_state = maintenance_state self.fatal_error = fatal_error @@ -274,7 +275,7 @@ async def send_one(msg): sent.append(msg) state = SimpleNamespace( - maintenance_state="serving", rebuild_futures={}, last_rebuild=None, send_one=send_one + backend_epoch=0, maintenance_state="serving", rebuild_futures={}, last_rebuild=None, send_one=send_one ) api_server._GLOBAL_STATE = state try: @@ -303,7 +304,7 @@ async def boom(msg): raise RuntimeError("zmq down") state = SimpleNamespace( - maintenance_state="serving", rebuild_futures={}, last_rebuild=None, send_one=boom + backend_epoch=0, maintenance_state="serving", rebuild_futures={}, last_rebuild=None, send_one=boom ) api_server._GLOBAL_STATE = state try: diff --git a/tests/server/test_responses_api.py b/tests/server/test_responses_api.py index 887e78e573..6eda006cb9 100644 --- a/tests/server/test_responses_api.py +++ b/tests/server/test_responses_api.py @@ -351,7 +351,7 @@ def __init__(self, outputs, finish_reason=None, cached_tokens=0): self._uid = 0 self.last_sent = None - def new_user(self): + async def new_user(self): self._uid += 1 return self._uid diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..5e311b0507 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -5,7 +5,15 @@ from queue import Empty as _Empty from freetoken.utils import progress -from freetoken.server.supervisor import BackendHandle, LoadProgress, drain_ready, phase_slug +from freetoken.server.supervisor import ( + BackendHandle, + LoadProgress, + WorkerDied, + _retire_oom_workers, + drain_ready, + phase_slug, + run_backend_supervisor, +) def test_drain_ready_counts_ready_acks_and_applies_progress(): @@ -118,6 +126,282 @@ def is_alive(self) -> bool: assert seen["failure"] == "ValueError: bad checkpoint config" +class _Process: + def __init__(self, name: str, alive: bool = True, survives_kill: bool = False): + self.name = name + self.alive = alive + self.survives_kill = survives_kill + self.calls: list[str] = [] + + def is_alive(self) -> bool: + return self.alive + + def terminate(self) -> None: + self.calls.append("terminate") + self.alive = False + + def join(self, timeout: float) -> None: + self.calls.append(f"join:{timeout}") + + def kill(self) -> None: + self.calls.append("kill") + if not self.survives_kill: + self.alive = False + + +def test_oom_worker_retirement_terminates_joins_and_kills_a_stubborn_worker(): + class Stubborn(_Process): + def terminate(self) -> None: + self.calls.append("terminate") + + process = Stubborn("scheduler") + + _retire_oom_workers([process]) + + assert process.calls == ["terminate", "join:10", "kill", "join:2"] + + +def test_oom_worker_retirement_denies_replacement_when_a_worker_cannot_stop(): + class Unstoppable(_Process): + def terminate(self) -> None: + self.calls.append("terminate") + + def kill(self) -> None: + self.calls.append("kill") + + process = Unstoppable("tokenizer") + + with __import__("pytest").raises(WorkerDied, match="could not be stopped: tokenizer"): + _retire_oom_workers([process]) + + +def test_post_ready_typed_cuda_oom_retires_then_replaces_once(): + class StopAfterReplacement(Exception): + pass + + initial_process = _Process("scheduler", alive=False) + initial_queue: "queue.Queue" = queue.Queue() + initial_queue.put("ready") + initial_queue.put(("cuda_oom", "CUDA out of memory in decode")) + initial = BackendHandle(initial_queue, [initial_process], expected_acks=1) + + replacement_process = _Process("replacement") + replacement_queue: "queue.Queue" = queue.Queue() + replacement_queue.put("ready") + replacement = BackendHandle(replacement_queue, [replacement_process], expected_acks=1) + events: list[tuple] = [] + + def ready() -> None: + events.append(("ready",)) + if len(events) == 3: + raise StopAfterReplacement() + + with __import__("pytest").raises(StopAfterReplacement): + run_backend_supervisor( + initial, + LoadProgress(), + on_ready=ready, + on_failure=lambda reason: events.append(("failure", reason)), + on_restarting=lambda reason, attempt: events.append(("restarting", reason, attempt)), + restart_backend=lambda attempt: replacement, + max_restarts=1, + poll=0, + ) + + assert initial_process.calls == ["join:10"] + assert events == [ + ("ready",), + ("restarting", "CUDA out of memory in decode", 1), + ("ready",), + ] + + +def test_non_oom_and_startup_failure_never_start_a_replacement(): + class DeadProc: + name = "scheduler" + + def is_alive(self) -> bool: + return False + + started: list[int] = [] + for receipt in (("error", "RuntimeError: bad checkpoint"), ("error", "CUDA out of memory during load")): + q: "queue.Queue" = queue.Queue() + q.put(receipt) + failures: list[str] = [] + run_backend_supervisor( + BackendHandle(q, [DeadProc()], expected_acks=1), + LoadProgress(), + on_ready=lambda: None, + on_failure=failures.append, + restart_backend=lambda attempt: started.append(attempt), + max_restarts=1, + poll=0, + ) + assert failures == [receipt[1]] + assert started == [] + + +def test_typed_cuda_oom_during_startup_is_a_terminal_startup_failure_not_ready(): + """Recovery is deliberately post-ready only: a load-time CUDA OOM has no healthy + request generation to preserve and must not be mistaken for a readiness ack.""" + q: "queue.Queue" = queue.Queue() + q.put(("cuda_oom", "CUDA out of memory while loading weights")) + failures: list[str] = [] + starts: list[int] = [] + + run_backend_supervisor( + BackendHandle(q, [], expected_acks=1), + LoadProgress(), + on_ready=lambda: (_ for _ in ()).throw(AssertionError("startup OOM became ready")), + on_failure=failures.append, + restart_backend=lambda attempt: starts.append(attempt), + max_restarts=1, + poll=0, + ) + + assert failures == ["CUDA out of memory while loading weights"] + assert starts == [] + + +def test_second_post_ready_oom_exhausts_the_bounded_recovery_budget(): + initial_process = _Process("initial", alive=False) + initial_queue: "queue.Queue" = queue.Queue() + initial_queue.put("ready") + initial_queue.put(("cuda_oom", "first OOM")) + + replacement_process = _Process("replacement", alive=False) + replacement_queue: "queue.Queue" = queue.Queue() + replacement_queue.put("ready") + replacement_queue.put(("cuda_oom", "second OOM")) + starts: list[int] = [] + failures: list[str] = [] + ready: list[bool] = [] + + run_backend_supervisor( + BackendHandle(initial_queue, [initial_process], expected_acks=1), + LoadProgress(), + on_ready=lambda: ready.append(True), + on_failure=failures.append, + restart_backend=lambda attempt: starts.append(attempt) or BackendHandle( + replacement_queue, [replacement_process], expected_acks=1 + ), + max_restarts=1, + poll=0, + ) + + assert starts == [1] + assert ready == [True, True] + assert failures == ["second OOM"] + + +def test_replacement_readiness_deadline_fails_and_retires_the_partial_factory_handle(): + initial_process = _Process("initial", alive=False) + initial_queue: "queue.Queue" = queue.Queue() + initial_queue.put("ready") + initial_queue.put(("cuda_oom", "decode OOM")) + fresh_process = _Process("partial", alive=True) + fresh_queue: "queue.Queue" = queue.Queue() + failures: list[str] = [] + + run_backend_supervisor( + BackendHandle(initial_queue, [initial_process], expected_acks=1), + LoadProgress(), + on_ready=lambda: None, + on_failure=failures.append, + restart_backend=lambda _: BackendHandle(fresh_queue, [fresh_process], expected_acks=1), + restart_ready_timeout=0, + max_restarts=1, + poll=0, + ) + + assert failures == ["OOM recovery failed: backend readiness timed out"] + assert fresh_process.calls == ["terminate", "join:10"] + + +def test_stop_barrier_prevents_launch_after_oom_cleanup(): + class StopOnTerminate(_Process): + def __init__(self): + super().__init__("initial", alive=False) + self.stopped = False + + def join(self, timeout: float) -> None: + super().join(timeout) + self.stopped = True + + process = StopOnTerminate() + q: "queue.Queue" = queue.Queue() + q.put("ready") + q.put(("cuda_oom", "decode OOM")) + launches: list[int] = [] + + run_backend_supervisor( + BackendHandle(q, [process], expected_acks=1), + LoadProgress(), + on_ready=lambda: None, + restart_backend=lambda attempt: launches.append(attempt), + can_restart=lambda: not process.stopped, + max_restarts=1, + poll=0, + ) + + assert process.calls == ["join:10"] + assert launches == [] + + +def test_preparation_stop_barrier_wins_before_recovery_starts(): + q: "queue.Queue" = queue.Queue() + q.put("ready") + q.put(("cuda_oom", "decode OOM")) + process = _Process("initial", alive=False) + launches: list[int] = [] + restarts: list[tuple[str, int]] = [] + + run_backend_supervisor( + BackendHandle(q, [process], expected_acks=1), + LoadProgress(), + on_ready=lambda: None, + restart_backend=lambda attempt: launches.append(attempt), + on_restarting=lambda message, attempt: restarts.append((message, attempt)), + can_restart=lambda: False, + max_restarts=1, + poll=0, + ) + + assert launches == [] + assert restarts == [] + + +def test_stop_during_replacement_readiness_cleans_the_current_handle_and_never_reopens(): + initial_q: "queue.Queue" = queue.Queue() + initial_q.put("ready") + initial_q.put(("cuda_oom", "decode OOM")) + initial = _Process("initial", alive=False) + fresh = _Process("fresh") + stopped = {"value": False} + + class StopWhileWaiting: + def get(self, timeout: float): + stopped["value"] = True + raise _Empty() + + failures: list[str] = [] + ready: list[bool] = [] + run_backend_supervisor( + BackendHandle(initial_q, [initial], expected_acks=1), + LoadProgress(), + on_ready=lambda: ready.append(True), + on_failure=failures.append, + restart_backend=lambda _: BackendHandle(StopWhileWaiting(), [fresh], expected_acks=1), + can_restart=lambda: not stopped["value"], + max_restarts=1, + poll=0, + ) + + assert ready == [True] + assert failures == ["OOM recovery failed: backend startup cancelled by shutdown"] + assert fresh.calls == ["terminate", "join:10"] + + def test_supervisor_reports_failure_on_startup_death(): import queue From eb9e70f30ed7a714ce9918ff905e86d2628bc126 Mon Sep 17 00:00:00 2001 From: shangfeng-liu Date: Wed, 14 Oct 2026 13:32:59 +0800 Subject: [PATCH 2/4] fix(server): retire oom request accounting before backend replacement --- python/freetoken/server/api_server.py | 57 ++++++++++++++++++++------- tests/server/test_oom_recovery.py | 24 +++++++++++ 2 files changed, 66 insertions(+), 15 deletions(-) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 0d6082735c..bc2afd397d 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -384,15 +384,40 @@ def fail_dispatched_for_oom(self, message: str) -> None: """Terminally fail live requests once; no partial work is replayed into a new epoch.""" if self._loop is None: return + def _fail() -> None: - for uid, replies in list(self.ack_map.items()): - if any(reply.finished for reply in replies): - continue - replies.append(UserReply(uid=uid, incremental_output="", finished=True, - error=message, error_code="insufficient_workspace", epoch=self.request_epochs[uid])) - self.event_map[uid].set() + self.terminalize_oom_inflight(message) + self._loop.call_soon_threadsafe(_fail) + def terminalize_oom_inflight(self, message: str) -> None: + """Finish every admitted request exactly once on the API event loop. + + ``StatsTracker`` deliberately retains an aborted request until a terminal scheduler + acknowledgement arrives. A CUDA OOM retires that scheduler, so synthesize its terminal + acknowledgement for *all* inflight UIDs, including ones whose client has already removed + its delivery queue. Observe before delivery: that is the sole completion/accounting + chokepoint and leaves the zero-token OOM reply out of token totals. + """ + for uid in self.stats.inflight_uids: + replies = self.ack_map.get(uid) + if replies is not None and any(reply.finished for reply in replies): + continue + reply = UserReply( + uid=uid, + incremental_output="", + finished=True, + error=message, + error_code="insufficient_workspace", + epoch=self.request_epochs.get(uid, self.backend_epoch), + ) + self.stats.observe(reply) + if replies is not None: + replies.append(reply) + event = self.event_map.get(uid) + if event is not None: + event.set() + def _create_listener_once(self): if not self.initialized: self._loop = asyncio.get_running_loop() @@ -1164,12 +1189,9 @@ def _apply() -> None: _GLOBAL_STATE._ready_event.clear() # Resolve all loop-owned state atomically with the epoch transition, before old # workers are retired. No cache rebuild coroutine can re-open this gate afterward. - for uid, replies in list(_GLOBAL_STATE.ack_map.items()): - if not any(reply.finished for reply in replies): - replies.append(UserReply(uid=uid, incremental_output="", finished=True, - error="worker_restarting_after_cuda_oom: " + message, - error_code="insufficient_workspace", epoch=_GLOBAL_STATE.request_epochs[uid])) - _GLOBAL_STATE.event_map[uid].set() + _GLOBAL_STATE.terminalize_oom_inflight( + "worker_restarting_after_cuda_oom: " + message + ) for request_id, future in list(_GLOBAL_STATE.rebuild_futures.items()): _GLOBAL_STATE.rebuild_futures.pop(request_id, None) if not future.done(): @@ -1188,14 +1210,19 @@ def _apply() -> None: future.set_result({"status": "failed", "error": message}) if _GLOBAL_STATE._ready_event is not None: _GLOBAL_STATE._ready_event.set() - _on_loop(_apply) logger.error("Backend supervisor: %s", message) # No CacheRebuildReply will ever arrive from a dead backend, so wake any caller blocked # in dispatch_rebuild's await now — otherwise it strands until the full rebuild timeout. # Then take the whole serve down (see _exit_after_backend_death). Shell mode is excluded: # a person is sitting at that TUI, the API is theirs alone, and its stop path is ^C. - if not run_shell: - _exit_after_backend_death(BACKEND_DEATH_EXIT_GRACE_S) + # Publishing the failure is loop-serialized, but a wedged loop may reject that callback. + # The outer daemon must still receive a finite terminal signal; otherwise it could retain + # a ``serving`` frontend backed by workers that are already dead. + try: + _on_loop(_apply) + finally: + if not run_shell: + _exit_after_backend_death(BACKEND_DEATH_EXIT_GRACE_S) def _on_meta(meta: dict) -> None: # Per-unit cache VRAM costs + the free-VRAM seed + per-pool floors + the actual pool diff --git a/tests/server/test_oom_recovery.py b/tests/server/test_oom_recovery.py index 9ff00d8bcb..b9bb7aae9b 100644 --- a/tests/server/test_oom_recovery.py +++ b/tests/server/test_oom_recovery.py @@ -102,6 +102,30 @@ async def scenario() -> None: assert terminal[0].finished is True assert terminal[0].error_code == "insufficient_workspace" assert terminal[0].epoch == state.request_epoch(uid) + assert state.stats.active == 0 + # An OOM terminal finishes an admitted request; it is not a client cancellation. + assert state.stats.completed == 2 + + asyncio.run(scenario()) + + +def test_oom_terminalizes_an_already_aborted_delivery_uid_in_global_accounting(): + async def scenario() -> None: + state = _state() + state._loop = asyncio.get_running_loop() + uid = await state.new_user() + # abort_user removes delivery maps before a scheduler terminal acknowledgement; OOM + # must synthesize that terminal acknowledgement for StatsTracker nevertheless. + state.ack_map.pop(uid) + state.event_map.pop(uid) + state.request_epochs.pop(uid) + state.stats.on_abort(uid) + + state.fail_dispatched_for_oom("CUDA out of memory") + await asyncio.sleep(0) + + assert state.stats.active == 0 + assert state.stats.completed == 0 asyncio.run(scenario()) From 66c1b8889d4919d715ffaf9060f26ea70f8304ff Mon Sep 17 00:00:00 2001 From: shangfeng-liu Date: Wed, 14 Oct 2026 15:06:12 +0800 Subject: [PATCH 3/4] fix(api): preserve backend error codes in raw completions --- python/freetoken/server/openai_api.py | 7 +++- tests/server/test_openai_api.py | 55 +++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index d53e0d01ca..958b54d203 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -449,7 +449,7 @@ async def handle_completion( finish_reason = "stop" async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): - return create_error_response(ack.error) + return create_error_response(ack.error, code=getattr(ack, "error_code", None)) prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens @@ -476,7 +476,10 @@ async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) finish_reason = "stop" async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): - yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}}) + yield _sse( + {"error": {"message": ack.error, "type": "invalid_request_error", + "code": getattr(ack, "error_code", None)}} + ) yield b"data: [DONE]\n\n" return prompt_tokens += ack.prompt_tokens_delta diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 54dd845275..3b1fdada81 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -377,6 +377,61 @@ def test_completion_accepts_text_prompt(): assert response["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} +def _completion_client(replies: list[UserReply]) -> TestClient: + app = FastAPI() + register_openai_routes(app, lambda: FakeState(replies), lambda: {}) + return TestClient(app) + + +def test_raw_completion_non_stream_preserves_terminal_error_code(): + with _completion_client([ + UserReply( + uid=42, incremental_output="", finished=True, + error="worker_restarting_after_cuda_oom", error_code="insufficient_workspace", + ) + ]) as client: + response = client.post("/v1/completions", json={"model": "client-model", "prompt": "hi"}) + + assert response.status_code == 400 + assert response.json()["error"] == { + "message": "worker_restarting_after_cuda_oom", + "type": "invalid_request_error", + "param": None, + "code": "insufficient_workspace", + } + + +def test_raw_completion_stream_keeps_partial_text_then_emits_one_coded_oom_terminal(): + with _completion_client([ + UserReply(uid=42, incremental_output="partial", finished=False), + UserReply( + uid=42, incremental_output="", finished=True, + error="worker_restarting_after_cuda_oom", error_code="insufficient_workspace", + ), + ]) as client: + response = client.post( + "/v1/completions", + json={"model": "client-model", "prompt": "hi", "stream": True}, + ) + + assert response.status_code == 200 + events = parse_sse([response.content]) + assert events.count("[DONE]") == 1 + assert events[-1] == "[DONE]" + assert [event["choices"][0]["text"] for event in events if isinstance(event, dict) and event.get("choices")] == ["partial"] + error = next(event["error"] for event in events if isinstance(event, dict) and "error" in event) + assert error == { + "message": "worker_restarting_after_cuda_oom", + "type": "invalid_request_error", + "code": "insufficient_workspace", + } + assert not any( + choice.get("finish_reason") == "length" + for event in events if isinstance(event, dict) + for choice in event.get("choices", []) + ) + + def test_completion_forwards_length_finish_reason(): state = FakeState( [UserReply(uid=42, incremental_output="hello", finished=True, finish_reason="length")] From a4ebbfda742fbad1df240621647014a39c70acd6 Mon Sep 17 00:00:00 2001 From: shangfeng-liu Date: Wed, 14 Oct 2026 15:09:38 +0800 Subject: [PATCH 4/4] test(api): exercise raw streaming error response over http --- tests/server/test_openai_api.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 3b1fdada81..0a76780b8f 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -50,6 +50,11 @@ async def wait_for_ack(self, uid: int): for reply in self.replies: yield reply + async def stream_with_cancellation(self, gen, request, uid: int): + # Route serialization tests exercise its wire output, not disconnect handling. + async for chunk in gen: + yield chunk + def tool_schema(): return [