fix(server): opt-in backend recovery after runtime CUDA OOM

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)

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.

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). ---

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

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

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

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)

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

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

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:

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

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))