"""Single-process, text-only OpenAI HTTP adapter for an already loaded engine. Run the returned app with Uvicorn in one process (no multiple workers). Admission is bounded by ``queue_size`` INCLUDING the active job. Tokenization, chat formatting, prefill, and generation all run on one dedicated thread. Serving decode (Stream S1 -- iteration-level continuous batching): - Lone single-stream jobs run the dense ``engine.stream`` (``Session``) path, bitwise-identical to before, so single-stream tokens/s are unregressed by construction. The native paged-KV path (``create_paged_session``/block-16, bitwise-identical math) is the default single-stream path (``CISM_SERVE_PAGED=0`` opts back out to dense): cooled A/B measured 0.999x of dense single-stream (shared BlockPool, zero per-session alloc; Stream S2) with ~26x prefill savings on shared prefixes via copy-on-write fork. ``spec_k > 0``, Surjo, and FWKV jobs always keep the dense ``engine.stream`` path. - Stream S3 automatic prefix reuse: the paged single path above matches each arrival against a per-engine radix prefix index (``paging.PrefixCache``) and fork-shares the longest cached token prefix (suffix-only prefill). Shared-prefix decode is token-identical to full prefill (native fork re-checks token equality); stale anchors (S2 eviction / pool reset) invalidate on first failed fork with exact full-prefill fallback. Hit/miss counters via ``prefix_cache_stats``; ``CISM_SERVE_PREFIXCACHE=0`` disables. The dense ``BatchSession`` epoch path has no fork entry point, so it keeps full-prefill rebuilds (documented v0 limit). - Eligible co-queued greedy jobs (non-streaming, ``temperature == 0``, ``spec_k == 0``, no A2 opt-ins) on a dense engine with a steppable native ``BatchSession`` run in an iteration-level scheduler epoch: per-step ADD (new arrivals join mid-decode via session rebuild over survivors' ``prompt + generated`` prefixes -- exact for greedy decoding), per-step EVICT (finished/errored/cancelled leave immediately with one terminal publish), oldest-first waiting queue with a starvation guard, and chunked prefill (at most ``prefill_chunk_tokens`` new prompt tokens per rebuild; env ``CISM_SERVE_PREFILL_CHUNK``). Sampled jobs (same exclusions, ``temperature > 0``) form same-parameter epochs: identical (``temperature``, ``top_p``, ``top_k``) jobs share one native session with each job's own serial seed per sequence (token-identical to serial); sampled epochs form only on a drain (no mid-decode joins) and never rotate (a rebuild would restart RNG streams), so they always run to drain. Streaming, speculative, constrained/tool/logprob, Surjo/FWKV, pool/batch, and test-double jobs always run serially (remaining v0 limits below). - Kill switches: ``CISM_SERVE_PAGED=0`` (back to dense single), ``CISM_SERVE_MICROBATCH=0`` (disable batching entirely), ``CISM_SERVE_CONTINUOUS=0`` (fall back to the A1 FIFO micro-batch path), ``CISM_SERVE_PREFIXCACHE=0`` (disable S3 automatic prefix reuse; ``CISM_SERVE_PREFIXCACHE_ANCHORS`` bounds retained anchors, default 8). Rotation tunables: ``CISM_SERVE_EPOCH_STEPS`` (default 1024), ``CISM_SERVE_MAX_WAIT_S`` (default 30), ``CISM_SERVE_JOIN_DEFER_TOKENS`` (default 2: skip a join rebuild when the live batch is that close to draining). Stream S1 v0 limits (documented, not faked): - Batches are heterogeneous in prompt/limit/stop; greedy epochs are seed-independent (mid-decode joins allowed), sampled epochs are parameter-homogeneous form-on-drain groups (no joins, no rotation, so per-seq RNG streams stay serial-identical). Per-job truncation keeps ``"length"``/``"stop"`` exact. - No preemption with KV preservation: epoch rotation recomputes survivors' prefixes (one prefill per rebuild); no swap-out, no priorities. - No intra-prompt prefill chunking: one long prompt prefills atomically inside a single native call (it is deferred while an epoch is live so it never stalls decodes, then runs alone). - No cross-engine fairness beyond FIFO drain order on a shared worker. Stream A2 API surface (exact-math opt-in paths; never faked): - Token logprobs (``logprobs`` on ``/v1/completions``, ``logprobs`` + ``top_logprobs`` on ``/v1/chat/completions``, non-streaming only): after generation, each emitted token is rescored with the real engine ``logits()`` over its true prefix (the same native logits the decoder used -- bitwise vs the reference in ``test_logits_and_cached_generation``), with an exact float64 log-softmax. Cross-check: the mean negative logprob equals ``engine.nll()`` over the same IDs within fp32 noise. - JSON structured output (``response_format`` ``json_object``/``json_schema`` on chat, greedy only): a Python-side greedy loop calls the real engine ``logits()`` per step and takes the argmax over tokens whose full-sequence decode keeps the output a viable JSON (resp. schema) prefix. Only a documented schema subset is supported (object/array/string/integer/number/ boolean/null, properties, required, additionalProperties, items, enum, const); anything else is a 400, never silently ignored. - Tool calling (``tools`` + ``tool_choice`` on chat, non-streaming only): the assistant's text is parsed for a real ``{"name", "arguments"}`` call matching a declared tool, arguments are validated against the tool's schema, the call is executed by a real local registry function (currently ``calc``: exact arithmetic, no imports/calls), and the result is fed back as a follow-up turn. Bounded (at most ``_TOOL_MAX_CALLS`` executions), so the loop always terminates. Tool names outside the registry are a 400: the server cannot execute arbitrary client code and will not pretend to. - Everything without an exact path stays 400/404 (embeddings, audio/image, ``n > 1``, penalties, ``logit_bias``, streaming+logprobs/tools/JSON, batch+JSON/tools, pool engines+logprobs, unknown tools). Constrained, tool, and logprob jobs always run serially -- they never join a micro-batch -- so the single-stream gate is unregressed by construction. """ from __future__ import annotations import asyncio from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field import ast import hmac import json import math import os import queue import statistics import threading import time from typing import Annotated, Any, Iterator, Literal, Protocol import uuid from fastapi import APIRouter, Depends, FastAPI, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse from pydantic import BaseModel, ConfigDict, Field, model_validator from starlette.exceptions import HTTPException class _Engine(Protocol): model_id: str context_length: int def encode(self, prompt: str | list[int]) -> list[int]: ... def format_chat(self, messages: list[dict]) -> list[int]: ... def stream(self, prompt: str | list[int], **kwargs: Any) -> Iterator[Any]: ... class _StrictModel(BaseModel): model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False) class _StreamOptions(_StrictModel): include_usage: bool = False class _TextFormat(_StrictModel): type: Literal["text"] class _JsonObjectFormat(_StrictModel): type: Literal["json_object"] class _JsonSchemaBody(_StrictModel): model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False, populate_by_name=True) name: _NonemptyString schema_: dict[str, Any] = Field(alias="schema") strict: bool | None = None @property def schema(self) -> dict[str, Any]: return self.schema_ class _JsonSchemaFormat(_StrictModel): type: Literal["json_schema"] json_schema: _JsonSchemaBody _ResponseFormat = _TextFormat | _JsonObjectFormat | _JsonSchemaFormat class _FunctionDef(_StrictModel): name: _NonemptyString description: str | None = None parameters: dict[str, Any] | None = None strict: bool | None = None class _ToolDef(_StrictModel): type: Literal["function"] function: _FunctionDef class _NamedToolFunction(_StrictModel): name: _NonemptyString class _NamedToolChoice(_StrictModel): type: Literal["function"] function: _NamedToolFunction _ToolChoice = Literal["auto", "none", "required"] | _NamedToolChoice _NonemptyString = Annotated[str, Field(min_length=1)] _Token = Annotated[int, Field(ge=0)] class _GenerationRequest(_StrictModel): model: _NonemptyString max_tokens: int | None = Field(default=None, ge=1) max_completion_tokens: int | None = Field(default=None, ge=1) temperature: float = Field(default=0.0, ge=0.0, le=2.0) top_p: float = Field(default=1.0, gt=0.0, le=1.0) top_k: int = Field(default=0, ge=0) seed: int = Field(default=0, ge=0) stop: _NonemptyString | Annotated[list[_NonemptyString], Field(max_length=4)] | None = None stream: bool = False n: int = Field(default=1, ge=1, le=1) stream_options: _StreamOptions | None = None # CISM extension: greedy prompt-lookup speculation (0 = off). spec_k: int = Field(default=0, ge=0, le=16) @model_validator(mode="after") def check_options(self) -> _GenerationRequest: if self.max_tokens is not None and self.max_completion_tokens is not None: raise ValueError("Specify only one token limit.") if self.stream_options is not None and not self.stream: raise ValueError("stream_options requires stream=true.") return self @property def token_limit(self) -> int: if self.max_completion_tokens is not None: return self.max_completion_tokens return self.max_tokens if self.max_tokens is not None else 128 class _CompletionRequest(_GenerationRequest): prompt: str | Annotated[list[_Token], Field(min_length=1)] # Exact-math token logprobs (top-N per emitted token, non-streaming only). logprobs: int | None = Field(default=None, ge=0, le=20) class _Message(_StrictModel): role: Literal["system", "developer", "user", "assistant"] content: str class _ChatRequest(_GenerationRequest): messages: Annotated[list[_Message], Field(min_length=1)] response_format: _ResponseFormat | None = None # Exact-math token logprobs (OpenAI chat shape, non-streaming only). logprobs: bool = False top_logprobs: int | None = Field(default=None, ge=0, le=20) # Server-side tool loop (non-streaming only; registry tools only). tools: Annotated[list[_ToolDef], Field(max_length=8)] | None = None tool_choice: _ToolChoice | None = None @model_validator(mode="after") def check_a2_options(self) -> _ChatRequest: if self.top_logprobs is not None and not self.logprobs: raise ValueError("top_logprobs requires logprobs=true.") if self.tool_choice is not None and not self.tools: raise ValueError("tool_choice requires tools.") return self class _BatchRequest(_StrictModel): """Phase2 batched completions: list of prompts -> order-preserving results. Single ``model`` for the whole batch (same-model same-quant fast path via native ``BatchSession`` when available). ``prompts`` items mirror ``_CompletionRequest.prompt`` (text or token IDs). Sampling fields mirror ``_GenerationRequest`` except streaming (batch is non-streaming only); ``spec_k`` must be 0 because ``BatchSession`` has no speculation path. """ model: _NonemptyString prompts: Annotated[list[str | Annotated[list[_Token], Field(min_length=1)]], Field(min_length=1, max_length=64)] max_tokens: int | None = Field(default=None, ge=1) max_completion_tokens: int | None = Field(default=None, ge=1) temperature: float = Field(default=0.0, ge=0.0, le=2.0) top_p: float = Field(default=1.0, gt=0.0, le=1.0) top_k: int = Field(default=0, ge=0) seed: int = Field(default=0, ge=0) stop: _NonemptyString | Annotated[list[_NonemptyString], Field(max_length=4)] | None = None spec_k: int = Field(default=0, ge=0, le=16) @model_validator(mode="after") def check_options(self) -> _BatchRequest: if self.max_tokens is not None and self.max_completion_tokens is not None: raise ValueError("Specify only one token limit.") return self @property def token_limit(self) -> int: if self.max_completion_tokens is not None: return self.max_completion_tokens return self.max_tokens if self.max_tokens is not None else 128 class _LoadRequest(_StrictModel): model: _NonemptyString precision: Literal["fp32", "int8", "hybrid-int4", "hybrid-fp4"] = "int8" act_precision: Literal["fp32", "int8"] | None = None revision: str | None = None threads: int = Field(default=2, ge=1, le=64) trust_remote_code: bool = False class _UnloadRequest(_StrictModel): model: _NonemptyString class _CacheRequest(_StrictModel): model: _NonemptyString decode_tokens: int = Field(default=64, ge=1, le=2048) scrub_mb: int = Field(default=96, ge=1, le=1024) class _TokenizeRequest(_StrictModel): model: _NonemptyString text: _NonemptyString class _NllRequest(_StrictModel): model: _NonemptyString text: _NonemptyString window: int = Field(default=128, ge=2, le=2048) class _RuntimeRequest(_StrictModel): model: _NonemptyString lock_pages: bool | None = None keep_warm: float | None = None class _AutoBenchRequest(_StrictModel): model: _NonemptyString precision: Literal["fp32", "int8", "hybrid-int4", "hybrid-fp4"] = "int8" threads: int = Field(default=2, ge=1, le=64) max_tokens: int = Field(default=512, ge=1, le=2048) thread_sweep: bool = False spec_sweep: bool = False bench_act_int8: bool = False include_torch: bool = True compile: bool = False class _APIError(Exception): def __init__( self, status: int, message: str, code: str, param: str | None = None ) -> None: super().__init__(message) self.status = status error_type = { 401: "authentication_error", 429: "rate_limit_error", 503: "server_error", }.get(status, "invalid_request_error") self.body = {"error": { "message": message, "type": error_type, "param": param, "code": code, }} @dataclass class _Chunk: text: str finish_reason: str | None prompt_tokens: int completion_tokens: int spec_proposed: int = 0 spec_accepted: int = 0 # Full decoded token IDs behind this chunk (empty when the backend does # not report them). Used only by opt-in exact-math paths (logprobs). token_ids: list[int] = field(default_factory=list) # Accumulated OpenAI-style tool calls for tool-loop jobs. tool_calls: list[dict] = field(default_factory=list) # Stream S3 automatic prefix-cache outcome for this job's prefill. prefix_hit: bool = False prefix_len: int = 0 @dataclass(eq=False) class _Job: body: _CompletionRequest | _ChatRequest engine: _Engine | None = None output: queue.Queue = field(default_factory=lambda: queue.Queue(maxsize=8)) cancel: threading.Event = field(default_factory=threading.Event) finished: bool = False consumed: bool = False output_ready: asyncio.Event = field(default_factory=asyncio.Event) notification_lock: threading.Lock = field(default_factory=threading.Lock) notification_pending: bool = False notified: bool = False # Worker-side arrival time (monotonic) for the S1 starvation guard. enqueued_at: float = 0.0 def notify(self, loop: asyncio.AbstractEventLoop) -> None: with self.notification_lock: if self.notification_pending or self.notified: return self.notification_pending = True try: loop.call_soon_threadsafe(self._wake) except RuntimeError: # Shutdown can close the loop before a native call returns. self.notification_pending = False self.cancel.set() def _wake(self) -> None: with self.notification_lock: self.notification_pending = False self.notified = True self.output_ready.set() @dataclass class _ContState: """Per-job mutable state for one S1 continuous epoch (worker thread only).""" job: _Job engine: Any prompt_ids: list[int] stop: list[str] | None eff: int # effective generation budget: min(token_limit, ctx - plen) context: int arrival: float admittable: bool # S1 epoch kind: "greedy" (mid-decode joins allowed -- seed-independent), # "sampled" (form-only epochs -- per-seq RNG state cannot survive a # rebuild), or "serial" (never batched). kind: str = "serial" gen: list[int] = field(default_factory=list) @property def plen(self) -> int: return len(self.prompt_ids) @property def remaining(self) -> int: return self.eff - len(self.gen) class _Worker: def __init__(self, engine: _Engine, capacity: int) -> None: self.engine = engine self.loop = asyncio.get_running_loop() self.capacity = capacity self.pending: queue.Queue[_Job] = queue.Queue(maxsize=capacity) self.lock = threading.Lock() self.jobs: set[_Job] = set() self.stopping = threading.Event() # A non-cooperative native call cannot be killed safely by Python. self.thread = threading.Thread(target=self._run, name="cism-inference", daemon=True) @property def ready(self) -> bool: return self.thread.is_alive() and not self.stopping.is_set() def submit(self, body: _CompletionRequest | _ChatRequest, engine: _Engine | None = None) -> _Job: with self.lock: if not self.ready: raise _APIError(503, "Server is not ready.", "server_unavailable") if len(self.jobs) >= self.capacity: raise _APIError(429, "Generation queue is full.", "queue_full") job = _Job(body, engine=engine) job.enqueued_at = time.monotonic() self.jobs.add(job) self.pending.put_nowait(job) return job def release(self, job: _Job) -> None: job.cancel.set() job.notify(self.loop) with self.lock: job.consumed = True if job.finished: self.jobs.discard(job) def stop(self) -> None: with self.lock: self.stopping.set() for job in self.jobs: job.cancel.set() job.notify(self.loop) self.jobs.clear() def _put(self, job: _Job, value: _Chunk | _APIError) -> bool: while not job.cancel.is_set() and not self.stopping.is_set(): try: job.output.put(value, timeout=0.05) job.notify(self.loop) return True except queue.Full: pass return False def _run(self) -> None: while True: try: first = self.pending.get(timeout=0.05) except queue.Empty: if self.stopping.is_set(): return continue takes = 1 try: if first.cancel.is_set() or self.stopping.is_set(): self._finish(first) elif self._use_continuous(first): takes = self._run_continuous(first) else: batch = self._collect_batch(first) takes = len(batch) if len(batch) > 1: self._run_micro_batch(batch) else: self._run_single(first) finally: for _ in range(takes): self.pending.task_done() def _use_continuous(self, job: _Job) -> bool: """True when ``job`` heads an S1 iteration-level epoch. Requires both kill switches on (``CISM_SERVE_CONTINUOUS``, ``CISM_SERVE_MICROBATCH``), a batch-eligible body (non-streaming greedy, or non-streaming sampled with valid params -- same exclusions otherwise: no spec/A2 opt-ins), and a dense engine with a steppable native ``BatchSession``. Everything else keeps the exact A1 FIFO path, so streaming/speculative/Surjo/test-double jobs are bit-for-bit unchanged. Non-admittable heads inside the loop still run on the serial ``_run_single`` path. """ if not _serve_continuous_enabled() or not _serve_microbatch_enabled(): return False if not (_micro_batch_eligible(job.body) or _sampled_batch_eligible(job.body)): return False engine = job.engine or self.engine if engine is None or _is_pool_or_batch_engine(engine): return False return _can_native_batch(engine) def _collect_batch(self, first: _Job) -> list[_Job]: """Drain co-queued jobs without waiting (zero added latency). A lone job returns alone and runs immediately on the single path, so single-stream latency/tokens-s are unregressed by construction. Only when jobs are already co-queued under concurrent load is a micro-batch formed, capped at ``_MICRO_BATCH_MAX``. """ batch = [first] if not _serve_microbatch_enabled(): return batch while len(batch) < _MICRO_BATCH_MAX: try: batch.append(self.pending.get_nowait()) except queue.Empty: break return batch def _finish(self, job: _Job) -> None: with self.lock: job.finished = True if job.consumed: self.jobs.discard(job) def _run_single(self, job: _Job) -> None: iterator = None engine = job.engine or self.engine try: if job.cancel.is_set() or self.stopping.is_set(): return body = job.body try: if isinstance(body, _ChatRequest): prompt = engine.format_chat( [message.model_dump() for message in body.messages] ) tokens = engine.encode(prompt) else: tokens = engine.encode(body.prompt) except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None if len(tokens) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") if job.cancel.is_set() or self.stopping.is_set(): return stop = [body.stop] if isinstance(body.stop, str) else body.stop # Stream A2 opt-in jobs run serially on exact-math paths (never # micro-batched, never paged): constrained JSON over engine.logits # and multi-turn tool loops over engine.stream. if isinstance(body, _ChatRequest) and _needs_tools(body): self._run_tool_loop(job, engine, tokens, body, stop) return if _needs_constrained(body): assert isinstance(body, _ChatRequest) self._run_constrained_job(job, engine, tokens, body) return if body.spec_k == 0 and _serve_paged_enabled() and _can_paged(engine): # Paged-KV single path: bitwise-identical to dense Session. iterator = iter(_stream_paged( engine, tokens, max_new_tokens=body.token_limit, temperature=body.temperature, top_p=body.top_p, top_k=body.top_k, seed=body.seed, stop=stop, cancel=job.cancel, chunk_size=16, )) else: iterator = iter(engine.stream( tokens, max_new_tokens=body.token_limit, temperature=body.temperature, top_p=body.top_p, top_k=body.top_k, seed=body.seed, stop=stop, cancel=job.cancel, chunk_size=16, spec_k=body.spec_k, )) terminal = False for chunk in iterator: if not self._put(job, _Chunk( chunk.text, chunk.finish_reason, chunk.prompt_tokens, chunk.completion_tokens, getattr(chunk, "spec_proposed", 0) or 0, getattr(chunk, "spec_accepted", 0) or 0, token_ids=list(getattr(chunk, "token_ids", None) or []), prefix_hit=bool(getattr(chunk, "prefix_hit", False)), prefix_len=int(getattr(chunk, "prefix_len", 0) or 0), )): break if chunk.finish_reason is not None: terminal = True break if not terminal and not job.cancel.is_set(): raise _APIError(503, "Generation did not complete.", "generation_failed") except _APIError as error: self._put(job, error) except Exception: self._put(job, _APIError(503, "Generation failed.", "generation_failed")) finally: if iterator is not None: with suppress(Exception): close = getattr(iterator, "close", None) if close is not None: close() self._finish(job) def _run_constrained_job(self, job: _Job, engine: Any, tokens: list[int], body: _ChatRequest) -> None: """One greedy constrained-JSON turn over real engine logits.""" try: if job.cancel.is_set() or self.stopping.is_set(): return target = _constrained_target(body) text, ids, finish, prompt_n, gen_n = _run_constrained_sync( engine, tokens, token_limit=body.token_limit, target=target, cancel=job.cancel) self._put(job, _Chunk(text, finish, prompt_n, gen_n, token_ids=list(ids))) except _APIError as error: self._put(job, error) except Exception: self._put(job, _APIError(503, "Generation failed.", "generation_failed")) finally: self._finish(job) def _collect_turn(self, job: _Job, engine: Any, tokens: list[int], body: _ChatRequest, stop: list[str] | None) -> tuple[str, list[int], str | None, int, int] | None: """Collect one dense ``engine.stream`` turn fully in memory. Tool-loop turns (intermediate and final) never publish incrementally: the caller publishes exactly one synthetic terminal chunk, so ``generate()`` observes the normal single-terminal protocol with the accumulated ``tool_calls`` attached. Returns (text, ids, finish, prompt_n, gen_n), or None when cancelled. """ iterator = None try: iterator = iter(engine.stream( tokens, max_new_tokens=body.token_limit, temperature=body.temperature, top_p=body.top_p, top_k=body.top_k, seed=body.seed, stop=stop, cancel=job.cancel, chunk_size=16, spec_k=body.spec_k, )) parts: list[str] = [] ids: list[int] = [] last: Any = None terminal = False for chunk in iterator: if job.cancel.is_set() or self.stopping.is_set(): return None ids.extend(list(getattr(chunk, "token_ids", None) or [])) parts.append(chunk.text) last = chunk if chunk.finish_reason is not None: terminal = True break if not terminal or last is None: raise _APIError(503, "Generation did not complete.", "generation_failed") return "".join(parts), ids, last.finish_reason, last.prompt_tokens, last.completion_tokens finally: if iterator is not None: with suppress(Exception): close = getattr(iterator, "close", None) if close is not None: close() def _run_tool_loop(self, job: _Job, engine: Any, tokens: list[int], body: _ChatRequest, stop: list[str] | None) -> None: """Bounded server-side tool loop with the real engine. Each turn generates assistant text with ``engine.stream``; text that parses as a ``{"name", "arguments"}`` call for a declared registry tool is validated and executed for real, and the JSON result is fed back as a follow-up user turn. At most ``_TOOL_MAX_CALLS`` executions, so the loop always terminates. Every turn is collected in memory and exactly one synthetic terminal chunk is published (with the accumulated ``tool_calls`` attached), so ``generate()`` observes the normal single-terminal protocol. """ try: assert body.tools is not None declared = {tool.function.name: tool.function for tool in body.tools} choice = body.tool_choice required_name = (choice.function.name if isinstance(choice, _NamedToolChoice) else None) require_call = choice == "required" or required_name is not None messages = [message.model_dump() for message in body.messages] tool_calls: list[dict] = [] final: tuple[str, list[int], str | None, int, int] | None = None pending_call: tuple[str, dict] | None = None current_tokens = list(tokens) for turn in range(_TOOL_MAX_CALLS + 1): if job.cancel.is_set() or self.stopping.is_set(): return last_turn = turn == _TOOL_MAX_CALLS turn_result = self._collect_turn( job, engine, current_tokens, body, stop) if turn_result is None: return text, _ids, finish, _pn, _gn = turn_result parsed = _parse_tool_call(text) if parsed is not None: name, arguments = parsed if name not in declared or (required_name is not None and name != required_name): parsed = None if parsed is None or last_turn: if parsed is not None: pending_call = parsed final = turn_result break name, arguments = parsed function = declared[name] params_schema = (function.parameters if function.parameters is not None else _TOOL_REGISTRY[name][1]) if _schema_support_error(params_schema, "tools.parameters") is not None: result: dict[str, Any] = { "error": "unsupported parameters schema for this tool"} elif not _mini_validate(arguments, params_schema): result = {"error": "arguments do not match the tool parameters schema"} else: try: result = dict(_TOOL_REGISTRY[name][0](arguments)) except ValueError as exc: result = {"error": str(exc)} except Exception: result = {"error": "tool execution failed"} tool_calls.append({ "id": f"call_{len(tool_calls)}", "type": "function", "function": {"name": name, "arguments": json.dumps( arguments, ensure_ascii=True, allow_nan=False, separators=(",", ":"))}, }) messages = [*messages, {"role": "assistant", "content": text}, {"role": "user", "content": ( f"Tool {name} returned: {json.dumps(result, ensure_ascii=True, allow_nan=False)}")}] try: current_tokens = engine.encode( engine.format_chat(messages)) except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None if len(current_tokens) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") assert final is not None text, ids, finish, prompt_n, gen_n = final if pending_call is not None and require_call: # Budget spent while the model still wants tools: report the # pending call honestly instead of looping forever. name, arguments = pending_call tool_calls.append({ "id": f"call_{len(tool_calls)}", "type": "function", "function": {"name": name, "arguments": json.dumps( arguments, ensure_ascii=True, allow_nan=False, separators=(",", ":"))}, }) finish = "tool_calls" self._put(job, _Chunk(text, finish, prompt_n, gen_n, token_ids=list(ids), tool_calls=tool_calls)) except _APIError as error: self._put(job, error) except Exception: self._put(job, _APIError(503, "Generation failed.", "generation_failed")) finally: self._finish(job) def _run_micro_batch(self, batch: list[_Job]) -> None: """FIFO continuous-batching step for co-queued greedy jobs (v0). Policy: strict FIFO; the head job must be micro-batch eligible (non-streaming, greedy, ``spec_k == 0``) on a dense engine with a native single-weight-stream batch path. The eligible FIFO prefix with the head's ``(stop, token_limit)`` key runs as one ``BatchSession`` micro-batch (chunked further by ``ContinuousSchedulerV0`` budgets when configured on the worker); every other job runs serially in order. Greedy decode is seed-independent, so coalescing is bitwise-identical to serial execution. Per-job validation errors never poison the batch. """ try: head = batch[0] engine = head.engine or self.engine key = _micro_batch_key(head.body) if _micro_batch_eligible(head.body) else None group: list[_Job] = [] rest: list[_Job] = [] if (key is not None and not _is_pool_or_batch_engine(engine) and _can_native_batch(engine)): for job in batch: if ((job.engine or self.engine) is engine and _micro_batch_eligible(job.body) and _micro_batch_key(job.body) == key): group.append(job) else: rest.append(job) else: rest = list(batch) if len(group) < 2: rest = batch group = [] for job in group: if job.cancel.is_set() or self.stopping.is_set(): self._finish(job) group = [job for job in group if not job.cancel.is_set() and not self.stopping.is_set()] if group: self._run_batch_group(engine, group) for job in rest: if job.cancel.is_set() or self.stopping.is_set(): self._finish(job) else: self._run_single(job) finally: # Jobs finished inside _run_single/_run_batch_group are already # marked; mark any leftovers (cancelled heads) here. for job in batch: with self.lock: if not job.finished: job.finished = True if job.consumed: self.jobs.discard(job) def _run_batch_group(self, engine: Any, group: list[_Job]) -> None: """Run one compatible FIFO group as a single-weight-stream batch.""" body0 = group[0].body stop = [body0.stop] if isinstance(body0.stop, str) else (list(body0.stop) if body0.stop else None) token_limit = body0.token_limit cfg = getattr(self, "scheduler_config", None) or {} # Per-job encode/validation first: one bad prompt must not fail others. valid: list[tuple[_Job, Any]] = [] for job in group: body = job.body try: try: if isinstance(body, _ChatRequest): prompt = engine.format_chat( [message.model_dump() for message in body.messages]) prompt_arg: Any = engine.encode(prompt) else: prompt_arg = body.prompt probe = engine.encode(prompt_arg) if len(probe) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") if isinstance(prompt_arg, list) and len(prompt_arg) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") except _APIError: raise except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None valid.append((job, prompt_arg)) except _APIError as error: self._put(job, error) self._finish(job) except Exception: self._put(job, _APIError(503, "Generation failed.", "generation_failed")) self._finish(job) if not valid: return try: from .paging import ContinuousSchedulerV0 sched = ContinuousSchedulerV0( num_blocks=int(cfg.get("scheduler_blocks", 512)), block_size=16, max_batch_tokens=int(cfg.get("max_batch_tokens", 4096)), prefill_chunk_tokens=int(cfg.get("prefill_chunk_tokens", 512))) lens: list[int] = [] for job, prompt_arg in valid: try: lens.append(len(engine.encode(prompt_arg))) except (ValueError, TypeError): lens.append(token_limit) for idx, length in enumerate(lens): sched.add(idx, length, token_limit, stop=stop) chunks = sched.plan_chunks() except Exception: chunks = [[{"req_id": i} for i in range(len(valid))]] sched = None # type: ignore[assignment] ordered: dict[int, tuple] = {} failed: dict[int, _APIError] = {} prompts = [prompt_arg for _, prompt_arg in valid] try: for chunk in chunks: indices = [entry["req_id"] for entry in chunk] sub = [prompts[i] for i in indices] try: rows = _run_batch_chunk_sync( engine, sub, token_limit, 0.0, 1.0, 0, body0.seed, stop) except _APIError as error: for i in indices: failed[i] = error continue except (ValueError, TypeError) as exc: msg = str(exc) error = _APIError( 400, "Prompt exceeds the model context length.", "context_length_exceeded") if ( "context" in msg.lower() or "Prompt length" in msg) else _APIError( 400, "Invalid prompt or chat messages.", "invalid_prompt") for i in indices: failed[i] = error continue except Exception: error = _APIError(503, "Generation failed.", "generation_failed") for i in indices: failed[i] = error continue if len(rows) != len(indices): error = _APIError(503, "Generation did not complete.", "generation_failed") for i in indices: failed[i] = error continue for i, row in zip(indices, rows): ordered[i] = row if sched is not None: for entry in chunk: try: sched.free(entry["req_id"]) except Exception: pass finally: for pos, (job, _) in enumerate(valid): try: if pos in failed: self._put(job, failed[pos]) else: row = ordered.get(pos) if row is None: self._put(job, _APIError(503, "Generation did not complete.", "generation_failed")) else: text, _ids, finish, prompt_n, gen_n = row if finish is None: self._put(job, _APIError(503, "Generation did not complete.", "generation_failed")) else: self._put(job, _Chunk(text, finish, prompt_n, gen_n, token_ids=list(_ids))) finally: self._finish(job) def _run_continuous(self, first: _Job) -> int: """Iteration-level scheduler loop for S1 (vLLM-style, Python-side). ADD: each iteration drains new arrivals into an :class:`paging.IterationScheduler` waiting queue (oldest-first) and admits group-matching jobs into the active batch -- including mid-decode joins, which rebuild the native ``BatchSession`` with survivors' ``prompt + generated`` prefixes (greedy continuation is an exact function of the prefix, so joined decoding is bitwise-identical to an uninterrupted run; per-job remaining budgets truncate with ``"length"`` exactly as a lone run would). EVICT: one shared-session ``BatchSession`` weight-stream pass per iteration (``next_tokens(1)`` while the batch is shared or waiters are present, so finished/errored/cancelled leave immediately; ``min(16, remaining)`` chunks for a lone idle job, mirroring ``Engine.stream`` -- this keeps single-stream tokens/s at parity by avoiding a Python GIL round-trip per token when nobody else needs the batch; join latency is bounded by one chunk). Per-job remaining budgets truncate with ``"length"`` exactly as a lone run would. CHUNKED PREFILL: admissions per rebuild are capped at ``prefill_chunk_tokens`` new prompt tokens (env ``CISM_SERVE_PREFILL_CHUNK``); an oversized single prompt is deferred while an epoch is live and runs alone once it drains (intra-prompt chunking still needs native support -- documented limit). Epochs rotate after ``CISM_SERVE_EPOCH_STEPS`` steps or when a same-engine waiter exceeds ``CISM_SERVE_MAX_WAIT_S``. Non-admissible jobs (streaming/sampled/spec/A2, other engines) keep global FIFO order: when the batch is empty and the head is not admittable it runs on the serial ``_run_single`` path. Returns the number of ``pending`` gets consumed (``first`` + drains). """ from .paging import ContinuousSchedulerV0, IterationScheduler cfg = getattr(self, "scheduler_config", None) or {} chunk = _continuous_prefill_chunk( int(cfg.get("prefill_chunk_tokens", 512))) budget = int(cfg.get("max_batch_tokens", 4096)) sched = IterationScheduler( max_batch=_MICRO_BATCH_MAX, prefill_chunk_tokens=chunk, max_wait_s=_continuous_max_wait_s(), max_batch_tokens=budget) max_wait = sched.max_wait_s join_defer = _continuous_join_defer_tokens() states: dict[int, _ContState] = {} takes = 0 def arrive(job: _Job) -> None: if job.cancel.is_set() or self.stopping.is_set(): self._finish(job) return engine = job.engine or self.engine if engine is None: self._put(job, _APIError(503, "Server is not ready.", "server_unavailable")) self._finish(job) return try: prompt_ids, stop = _cont_encode(engine, job.body) except _APIError as error: self._put(job, error) self._finish(job) return except Exception: self._put(job, _APIError(503, "Generation failed.", "generation_failed")) self._finish(job) return limit = job.body.token_limit context = int(engine.context_length) eff = min(limit, context - len(prompt_ids)) if eff <= 0: # Prompt fills the context: zero new tokens, terminal # "length" with empty text (mirrors native max_new=0). self._put(job, _Chunk("", "length", len(prompt_ids), 0, token_ids=[])) self._finish(job) return admittable = _micro_batch_eligible(job.body) if admittable: kind = "greedy" group = (id(engine), "greedy") elif _sampled_batch_eligible(job.body): kind = "sampled" admittable = True group = (id(engine), "sampled", job.body.temperature, job.body.top_p, job.body.top_k) else: kind = "serial" group = (id(engine), "serial") states[id(job)] = _ContState( job, engine, prompt_ids, stop, eff, context, arrival=(job.enqueued_at or time.monotonic()), admittable=admittable, kind=kind) sched.add(id(job), len(prompt_ids), group, states[id(job)].arrival, admittable=admittable) def fits(rid: int, _admitted: list) -> bool: # Context-window compatibility: with maxR = max remaining over # the would-be active set, every member's extended prompt must # still fit (native requires prompt + max_new <= context), and # the token budget must hold (lone head exempt). active = [states[r] for r in sched.running_ids] cand = states[rid] # Sampled RNG exactness: a sampled job never joins a live epoch # (a rebuild would restart its RNG stream), and nothing joins a # live sampled epoch. Sampled epochs form only on a drain. if active and (cand.kind == "sampled" or sched.group_of(sched.running_ids[0])[1] == "sampled"): return False if cand.plen + cand.eff > cand.context: return False pool = active + [cand] max_r = max(s.remaining for s in pool) if any(s.plen + len(s.gen) + max_r > s.context for s in pool): return False if len(pool) > 1 and sum(s.plen + s.eff for s in pool) > budget: return False if active and join_defer > 0: starved = (time.monotonic() - cand.arrival) > max_wait if not starved and all(s.remaining <= join_defer for s in active): return False # drain is imminent; join would only # force a pointless recompute-prefill return True arrive(first) takes += 1 session = None slot_rids: list[int] = [] fresh = False # first step after a (re)build stays granular epoch_steps = 0 epoch_cap = _continuous_epoch_steps() try: while len(sched) > 0: if self.stopping.is_set(): if session is not None: try: session.cancel() except Exception: pass session = None break while True: try: job = self.pending.get_nowait() except queue.Empty: break takes += 1 arrive(job) now = time.monotonic() for rid in sched.waiting_ids: state = states.get(rid) if state is not None and state.job.cancel.is_set(): states.pop(rid, None) try: sched.drop(rid, now) except Exception: pass self._finish(state.job) if not sched.running_ids: session = None slot_rids = [] head = sched.head_id() if head is None: break state = states[head] if not state.admittable: states.pop(head, None) try: sched.drop(head, now) except Exception: pass self._run_single(state.job) continue admitted, _ = sched.admit(now, fits=fits) if not admitted: # Head admittable but unfittable: fail it alone so # one bad job never wedges the queue (per-job # isolation, as in the A1 batch group). states.pop(head, None) try: sched.drop(head, now) except Exception: pass self._put(state.job, _APIError( 503, "Generation failed.", "generation_failed")) self._finish(state.job) continue session, slot_rids = self._cont_build(states, sched) fresh = True epoch_steps = 0 continue admitted, rotate = sched.admit(now, fits=fits) if admitted: session, slot_rids = self._cont_build( states, sched, session=session) fresh = True epoch_steps = 0 if rotate or epoch_steps >= epoch_cap: running_kind = (sched.group_of(sched.running_ids[0])[1] if sched.running_ids else "greedy") if running_kind == "sampled": # Never rotate a live sampled epoch: rotation # requeues + rebuilds, restarting per-seq RNG # streams and diverging from serial. Reset the # counter and fall through to stepping -- sampled # epochs always run to drain; waiters (even # starved same-group ones) join the next epoch. # (A bare `continue` here would livelock: rotation # changes no state, so the live epoch would never # step again.) epoch_steps = 0 else: if session is not None: try: session.cancel() except Exception: pass session = None slot_rids = [] fresh = False sched.requeue_running() epoch_steps = 0 continue # Step size: granular (1 token) while the batch is shared # or waiters are present (prompt evict/publish + fast # joins); chunked (16, mirroring Engine.stream) for a lone # idle job, so a lone request follows the same native call # pattern as the dense path (avoids a GIL round-trip per # token at parity). if len(sched.running_ids) > 1 or sched.waiting_ids or fresh: count = 1 else: count = min(16, states[sched.running_ids[0]].remaining) if self._cont_step(states, sched, session, slot_rids, count): session = None # all drained; loop re-admits slot_rids = [] else: epoch_steps += 1 fresh = False finally: for rid in list(states): state = states.pop(rid) try: self._put(state.job, _APIError( 503, "Generation failed.", "generation_failed")) except Exception: pass try: self._finish(state.job) except Exception: pass return takes def _cont_build(self, states: dict[int, _ContState], sched: Any, session: Any = None) -> tuple[Any, list[int]]: """(Re)build the epoch's shared native session for mid-decode joins. Prompts are survivors' ``prompt + generated`` prefixes plus new arrivals' prompts; ``max_new`` is the max remaining budget (per-job truncation keeps ``"length"`` exact). Returns ``(session, slot_rids)`` where ``slot_rids`` is the build-time member order (native slot ``i`` forever maps to ``slot_rids[i]`` until the next rebuild). Postcondition: a ``None`` session implies empty running state (every member failed, so the caller simply re-admits). """ if session is not None: try: session.cancel() except Exception: pass engine = states[sched.running_ids[0]].engine def fail_running(error: _APIError) -> tuple[Any, list[int]]: for rid in sched.running_ids: state = states.pop(rid, None) if state is None: continue self._put(state.job, error) self._finish(state.job) for rid in list(sched.running_ids): try: sched.finish(rid, time.monotonic()) except Exception: pass return None, [] members = [rid for rid in sched.running_ids if states[rid].remaining > 0] if not members: return fail_running(_APIError(503, "Generation did not complete.", "generation_failed")) prompts = [states[r].prompt_ids + states[r].gen for r in members] max_r = max(states[r].remaining for r in members) body0 = states[members[0]].job.body eos = _engine_eos(engine) # Epoch-homogeneous sampling params (group key): greedy epochs run # temperature=0 (seed-independent, join-safe); sampled epochs carry # the group's (temperature, top_p, top_k) with each job's own serial # seed per sequence, so batched sampled decode is token-identical # to serial execution. if states[members[0]].kind == "sampled": temperature, top_p, top_k = (body0.temperature, body0.top_p, body0.top_k) seq_seeds = [states[r].job.body.seed for r in members] else: temperature, top_p, top_k = 0.0, 1.0, 0 seq_seeds = [] seed = states[members[0]].job.body.seed lock = getattr(engine, "_generation_lock", None) if lock is not None: lock.acquire() try: try: fresh = engine._native.create_batch_session( prompts, max_new_tokens=max_r, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=eos, seq_seeds=seq_seeds) except (ValueError, TypeError) as exc: message = str(exc) if "context" in message.lower() or "KV cache" in message: error = _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") else: error = _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") return fail_running(error) except Exception: return fail_running(_APIError(503, "Generation failed.", "generation_failed")) return fresh, list(members) finally: if lock is not None: lock.release() def _cont_step(self, states: dict[int, _ContState], sched: Any, session: Any, slot_rids: list[int], count: int) -> bool: """Run one decode pass (``count`` tokens) over the active batch. Native slot ``i`` maps to ``slot_rids[i]`` (build-time order, stable across evictions since finished seqs are ``cancel_seq``-masked, not removed). Finished / errored / cancelled sequences leave immediately with exactly one terminal publish. Returns True when no active member remains (epoch drained). """ from .paging import ContinuousSchedulerV0 alive = [states[r] for r in slot_rids if r in states] if not alive: return True engine = alive[0].engine lock = getattr(engine, "_generation_lock", None) def drop(rid: int, value: _Chunk | _APIError | None) -> None: state = states.pop(rid, None) try: sched.finish(rid, time.monotonic()) except Exception: pass if state is not None: if value is not None: self._put(state.job, value) self._finish(state.job) if lock is not None: lock.acquire() try: try: out = session.next_tokens(max(1, count)) reasons = list(session.finish_reasons) except Exception: for rid in slot_rids: drop(rid, _APIError(503, "Generation failed.", "generation_failed")) return True finally: if lock is not None: lock.release() if len(out) != len(slot_rids) or len(reasons) != len(slot_rids): for rid in slot_rids: drop(rid, _APIError(503, "Generation did not complete.", "generation_failed")) return True for index, rid in enumerate(slot_rids): state = states.get(rid) if state is None: continue # evicted earlier; slot masked via cancel_seq toks, reason = out[index], reasons[index] if state.job.cancel.is_set() or self.stopping.is_set(): try: session.cancel_seq(index) except Exception: pass drop(rid, None) continue if reason == "cancelled": drop(rid, None) continue state.gen.extend(list(toks or [])) finish = reason if reason in ("stop", "length") else None if finish is None and state.remaining <= 0: # Per-job budget spent while the shared session runs on: # terminal "length", exactly as a lone run would report. finish = "length" if finish is None: continue try: text = engine.tokenizer.decode( list(state.gen), skip_special_tokens=True, clean_up_tokenization_spaces=False) except Exception: try: session.cancel_seq(index) except Exception: pass drop(rid, _APIError(503, "Generation failed.", "generation_failed")) continue if state.stop: text, stop_finish = ContinuousSchedulerV0.apply_stop( text, state.stop) finish = stop_finish or finish try: session.cancel_seq(index) except Exception: pass drop(rid, _Chunk(text, finish, state.plen, len(state.gen), token_ids=list(state.gen))) return not sched.running_ids async def receive( self, job: _Job, request: Request, *, streaming: bool = False ) -> _Chunk | _APIError: # Before ASGI 2.4 StreamingResponse owns the disconnect listener. spec = request.scope.get("asgi", {}).get("spec_version", "2.0") poll_disconnect = not streaming or tuple(map(int, spec.split("."))) >= (2, 4) while True: if self.stopping.is_set(): raise _APIError(503, "Server is shutting down.", "server_unavailable") if job.cancel.is_set(): raise _APIError(503, "Request disconnected.", "request_cancelled") try: return job.output.get_nowait() except queue.Empty: if not self.thread.is_alive(): raise _APIError(503, "Inference worker is unavailable.", "server_unavailable") # Re-arm only when drained. Keep any scheduled callback marked as # pending until it runs, even if the consumer drains a new burst. # Thus neither a fast producer nor a slow client can queue callbacks # without bound. Rechecking after clear prevents a lost wakeup. with job.notification_lock: job.notified = False job.output_ready.clear() try: return job.output.get_nowait() except queue.Empty: pass try: async with asyncio.timeout(0.05): await job.output_ready.wait() except TimeoutError: if poll_disconnect and await request.is_disconnected(): job.cancel.set() raise _APIError(503, "Request disconnected.", "request_cancelled") class _JobResponse(StreamingResponse): def __init__(self, content, worker: _Worker, job: _Job) -> None: super().__init__(content, media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }) self.worker = worker self.job = job async def __call__(self, scope, receive, send) -> None: # A send failure can leave the generator suspended at yield, so its # finally block alone cannot reliably cancel inference on disconnect. try: await super().__call__(scope, receive, send) finally: self.worker.release(self.job) await self.body_iterator.aclose() def _usage(chunk: _Chunk) -> dict[str, int]: proposed, accepted = chunk.spec_proposed, chunk.spec_accepted return { "prompt_tokens": chunk.prompt_tokens, "completion_tokens": chunk.completion_tokens, "total_tokens": chunk.prompt_tokens + chunk.completion_tokens, "spec_proposed": proposed, "spec_accepted": accepted, "spec_accept_rate": round(accepted / proposed, 4) if proposed else None, } async def _rescore_logprobs(engine: Any, body: _CompletionRequest | _ChatRequest, gen_ids: list[int], text: str) -> dict[str, Any]: """Rescore emitted IDs with real engine logits (thread-offloaded). The prompt is re-encoded on the tokenizer-only path (identical to the worker's encoding); scoring itself runs in a worker thread since native ``logits()`` calls block. Raises _APIError on any failure -- logprobs are never synthesized. """ top_n = _logprobs_top_n(body) chat = isinstance(body, _ChatRequest) if text and not gen_ids: raise _APIError(503, "Engine did not report token IDs; logprobs unavailable.", "generation_failed") def _run() -> dict[str, Any]: try: if chat: assert isinstance(body, _ChatRequest) prompt_ids = engine.encode(engine.format_chat( [message.model_dump() for message in body.messages])) else: assert isinstance(body, _CompletionRequest) prompt_ids = engine.encode(body.prompt) except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None if gen_ids and not all(isinstance(t, int) for t in gen_ids): raise _APIError(503, "Engine reported invalid token IDs.", "generation_failed") surfaces, chosen, tops = _score_token_logprobs(engine, prompt_ids, gen_ids, top_n) if chat: return {"content": [ {"token": token, "logprob": logprob, "top_logprobs": [{"token": t, "logprob": lp} for t, lp in top]} for token, logprob, top in zip(surfaces, chosen, tops)]} return {"tokens": surfaces, "token_logprobs": chosen, "top_logprobs": [{token: logprob for token, logprob in top} for top in tops]} return await asyncio.to_thread(_run) # -- Stream A2: exact-math opt-in paths (logprobs / JSON / tools) ------------- #: Max tool executions per chat request; the loop always terminates. _TOOL_MAX_CALLS = 4 #: Documented JSON-schema subset for constrained decoding and tool args. _SCHEMA_TYPES = frozenset({"object", "array", "string", "integer", "number", "boolean", "null"}) _SCHEMA_VALUE_KEYS = frozenset({"type", "properties", "required", "additionalProperties", "items", "enum", "const"}) def _wants_logprobs(body: _CompletionRequest | _ChatRequest) -> bool: if isinstance(body, _ChatRequest): return bool(body.logprobs) return body.logprobs is not None def _logprobs_top_n(body: _CompletionRequest | _ChatRequest) -> int: if isinstance(body, _ChatRequest): return body.top_logprobs or 0 return body.logprobs or 0 def _needs_constrained(body: _CompletionRequest | _ChatRequest) -> bool: fmt = getattr(body, "response_format", None) return fmt is not None and getattr(fmt, "type", "text") != "text" def _constrained_target(body: _CompletionRequest | _ChatRequest) -> Any: """None (plain text), True (any JSON), or a schema dict.""" fmt = getattr(body, "response_format", None) if fmt is None or getattr(fmt, "type", "text") == "text": return None if getattr(fmt, "type", None) == "json_object": return True return fmt.json_schema.schema # type: ignore[union-attr] def _needs_tools(body: _CompletionRequest | _ChatRequest) -> bool: tools = getattr(body, "tools", None) if not tools: return False return getattr(body, "tool_choice", None) != "none" # -- minimal JSON-schema subset (validate + support check) -------------------- def _schema_support_error(schema: Any, where: str = "schema") -> str | None: """None when ``schema`` is inside the supported subset, else a reason. Supported: type (object/array/string/integer/number/boolean/null, or a list thereof), properties, required, additionalProperties (bool), items, enum (JSON scalars), const. Everything else ($ref, allOf/anyOf/oneOf, pattern, format, length/range keywords, ...) is rejected so the server never silently ignores a constraint. """ if not isinstance(schema, dict): return f"{where} must be an object" for key in schema: if key not in _SCHEMA_VALUE_KEYS: return f"unsupported keyword {key!r} at {where}" kind = schema.get("type", None) if kind is not None: kinds = kind if isinstance(kind, list) else [kind] if not kinds or any(k not in _SCHEMA_TYPES for k in kinds): return f"unsupported type at {where}" if "properties" in schema: props = schema["properties"] if not isinstance(props, dict): return f"properties must be an object at {where}" for name, sub in props.items(): reason = _schema_support_error(sub, f"{where}.properties.{name}") if reason is not None: return reason if "required" in schema: required = schema["required"] if (not isinstance(required, list) or not required or any(not isinstance(name, str) for name in required)): return f"required must be a nonempty string list at {where}" if "additionalProperties" in schema and not isinstance( schema["additionalProperties"], bool): return f"additionalProperties must be a boolean at {where}" if "items" in schema: reason = _schema_support_error(schema["items"], f"{where}.items") if reason is not None: return reason if "enum" in schema: enum = schema["enum"] if (not isinstance(enum, list) or not enum or any(isinstance(v, (dict, list)) for v in enum)): return f"enum must be a nonempty scalar list at {where}" return None def _mini_validate(value: Any, schema: Any) -> bool: """Validate ``value`` against the supported schema subset.""" if not isinstance(schema, dict): return False if "const" in schema and value != schema["const"]: return False if "enum" in schema and value not in schema["enum"]: return False kind = schema.get("type", None) if kind is not None: kinds = set(kind) if isinstance(kind, list) else {kind} if not any(_mini_type_ok(value, k) for k in kinds): return False if len(kinds) == 1: kind = next(iter(kinds)) else: kind = None if kind == "object" or (kind is None and "properties" in schema): if not isinstance(value, dict): return kind is None props = schema.get("properties", {}) for name, sub in props.items(): if name in value and not _mini_validate(value[name], sub): return False for name in schema.get("required", []): if name not in value: return False if schema.get("additionalProperties") is False: if any(name not in props for name in value): return False return True if kind == "array" or (kind is None and "items" in schema): if not isinstance(value, list): return kind is None items = schema.get("items") return items is None or all(_mini_validate(v, items) for v in value) return True def _mini_type_ok(value: Any, kind: str) -> bool: if kind == "null": return value is None if kind == "boolean": return isinstance(value, bool) if kind == "integer": return isinstance(value, int) and not isinstance(value, bool) if kind == "number": return isinstance(value, (int, float)) and not isinstance(value, bool) if kind == "string": return isinstance(value, str) if kind == "array": return isinstance(value, list) if kind == "object": return isinstance(value, dict) return False # -- JSON prefix viability (structural + top-level key allowlist) ------------- def _skip_ws(text: str, i: int) -> int: while i < len(text) and text[i] in " \t\n\r": i += 1 return i def _parse_string(text: str, i: int) -> tuple[str, int]: """Parse text[i] == '\"'. Returns (ok|partial|bad, next index).""" n = len(text) j = i + 1 while j < n: c = text[j] if c == '"': return "ok", j + 1 if c == "\\": j += 1 if j >= n: return "partial", j e = text[j] if e not in "\"\\/bfnrtu": return "bad", j if e == "u": hex_digits = text[j + 1:j + 5] if len(hex_digits) < 4: rest = text[j + 1:] if all(ch in "0123456789abcdefABCDEF" for ch in rest): return "partial", n return "bad", j if any(ch not in "0123456789abcdefABCDEF" for ch in hex_digits): return "bad", j j += 4 j += 1 continue if ord(c) < 0x20: return "bad", j j += 1 return "partial", n def _parse_literal(text: str, i: int, word: str) -> tuple[str, int]: rest = text[i:i + len(word)] if rest == word: return "ok", i + len(word) if word.startswith(rest) and i + len(rest) >= len(text): return "partial", len(text) return "bad", i def _parse_number(text: str, i: int) -> tuple[str, int]: n = len(text) j = i if j < n and text[j] == "-": j += 1 digits = 0 while j < n and text[j].isdigit(): j += 1 digits += 1 if digits == 0: # Lone "-" (or empty) is a partial number, anything else is bad. return ("partial", j) if j > i else ("bad", i) if j < n and text[j] == ".": j += 1 frac = 0 while j < n and text[j].isdigit(): j += 1 frac += 1 if frac == 0: return "partial", j if j < n and text[j] in "eE": j += 1 if j < n and text[j] in "+-": j += 1 exp = 0 while j < n and text[j].isdigit(): j += 1 exp += 1 if exp == 0: return "partial", j return "ok", j def _parse_value(text: str, i: int) -> tuple[str, int]: """Returns (ok|partial|bad, next index) for one JSON value at i.""" n = len(text) i = _skip_ws(text, i) if i >= n: return "partial", i c = text[i] if c == "{": i += 1 i = _skip_ws(text, i) if i >= n: return "partial", i if text[i] == "}": return "ok", i + 1 while True: i = _skip_ws(text, i) if i >= n: return "partial", i if text[i] != '"': return "bad", i status, i = _parse_string(text, i) if status != "ok": return status, i i = _skip_ws(text, i) if i >= n or text[i] != ":": return ("partial", i) if i >= n else ("bad", i) i += 1 status, i = _parse_value(text, i) if status != "ok": return status, i i = _skip_ws(text, i) if i >= n: return "partial", i if text[i] == ",": i += 1 continue if text[i] == "}": return "ok", i + 1 return "bad", i if c == "[": i += 1 i = _skip_ws(text, i) if i >= n: return "partial", i if text[i] == "]": return "ok", i + 1 while True: status, i = _parse_value(text, i) if status != "ok": return status, i i = _skip_ws(text, i) if i >= n: return "partial", i if text[i] == ",": i += 1 continue if text[i] == "]": return "ok", i + 1 return "bad", i if c == '"': return _parse_string(text, i) if c == "t": return _parse_literal(text, i, "true") if c == "f": return _parse_literal(text, i, "false") if c == "n": return _parse_literal(text, i, "null") if c == "-" or c.isdigit(): return _parse_number(text, i) return "bad", i def _json_prefix_viable(text: str) -> bool: """True when ``text`` is a strict prefix of (or equal to) some valid JSON.""" status, i = _parse_value(text, 0) if status == "bad": return False if status == "partial": return True return _skip_ws(text, i) >= len(text) def _top_level_keys(text: str) -> list[str]: """Complete top-level object keys visible in a (possibly partial) prefix.""" keys: list[str] = [] i = _skip_ws(text, 0) n = len(text) if i >= n or text[i] != "{": return keys depth = 0 while i < n: c = text[i] if c in "{[": depth += 1 i += 1 elif c in "}]": depth -= 1 i += 1 if depth <= 0: break elif c == '"': status, j = _parse_string(text, i) if status != "ok": break if depth == 1: k = _skip_ws(text, j) if k < n and text[k] == ":": try: keys.append(json.loads(text[i:j])) except ValueError: pass i = j else: i += 1 return keys def _schema_prefix_ok(text: str, schema: Any) -> bool: """Top-level container/key guidance for the supported schema subset.""" if not isinstance(schema, dict): return True kind = schema.get("type", None) expected: str | None = None if kind == "object" or (kind is None and "properties" in schema): expected = "object" elif kind == "array" or (kind is None and "items" in schema): expected = "array" stripped = text.lstrip() if not stripped: return True if expected == "object": if not stripped.startswith("{"): return False if schema.get("additionalProperties") is False and isinstance( schema.get("properties"), dict): allowed = set(schema["properties"]) if any(key not in allowed for key in _top_level_keys(text)): return False elif expected == "array": if not stripped.startswith("["): return False return True def _prefix_viable(text: str, schema: Any | None) -> bool: if not _json_prefix_viable(text): return False return schema is None or _schema_prefix_ok(text, schema) def _complete_and_valid(text: str, target: Any) -> bool: """True only for a complete document that validates (schema or any JSON).""" try: value = json.loads(text) except ValueError: return False if target is True or target is None: return True return _mini_validate(value, target) # -- exact log-softmax rescoring over real engine logits ---------------------- def _log_softmax_rows(logits: Any) -> list[float]: flat = list(logits) if not flat: raise ValueError("empty logits") peak = max(float(v) for v in flat) if not math.isfinite(peak): raise ValueError("non-finite logits") total = math.fsum(math.exp(float(v) - peak) for v in flat) offset = peak + math.log(total) return [float(v) - offset for v in flat] def _argsort_desc(values: Any) -> list[int]: try: import numpy as _np # hard dependency (pyproject); exact float64 path except ImportError: # pragma: no cover - portable fallback _np = None # type: ignore[assignment] if _np is not None: arr = _np.asarray(list(values), dtype=_np.float64) return [int(i) for i in _np.argsort(arr, kind="stable")[::-1]] order = list(range(len(list(values)))) order.sort(key=lambda i: float(list(values)[i]), reverse=True) return order def _score_token_logprobs(engine: Any, prompt_ids: list[int], gen_ids: list[int], top_n: int) -> tuple[list[str], list[float], list[list[tuple[str, float]]]]: """Exact per-token logprobs: real ``engine.logits`` + float64 log-softmax. ``prompt_ids`` is the exact encoded prompt; ``gen_ids`` the emitted IDs. Returns (token surfaces, chosen logprobs, top-N lists) in order. """ logits_fn = getattr(engine, "logits", None) if not callable(logits_fn): raise _APIError(400, "This engine does not expose logits; logprobs are unsupported.", "unsupported", "logprobs") decode = engine.tokenizer.decode surfaces: list[str] = [] chosen: list[float] = [] tops: list[list[tuple[str, float]]] = [] for pos, tid in enumerate(gen_ids): try: row = _log_softmax_rows(logits_fn(list(prompt_ids) + gen_ids[:pos])) except _APIError: raise except (ValueError, TypeError, RuntimeError) as exc: raise _APIError(400, f"Logprob rescoring failed: {exc}.", "scoring_failed") if tid < 0 or tid >= len(row): raise _APIError(503, "Engine reported an out-of-vocabulary token.", "generation_failed") order = _argsort_desc(row) keep = order[:max(top_n, 1)] if top_n else [tid] # Decode surfaces lazily but exactly (single-token decode each). table: dict[int, str] = {} for cand in dict.fromkeys([tid, *keep]): try: table[cand] = decode([cand], skip_special_tokens=True, clean_up_tokenization_spaces=False) except Exception: raise _APIError(503, "Engine tokenizer failed during rescoring.", "generation_failed") from None surfaces.append(table[tid]) chosen.append(row[tid]) tops.append([(table[c], row[c]) for c in keep[:top_n]] if top_n else []) return surfaces, chosen, tops # -- greedy constrained JSON decoding over real engine logits ----------------- def _run_constrained_sync(engine: Any, prompt_ids: list[int], *, token_limit: int, target: Any, cancel: threading.Event) -> tuple[str, list[int], str, int, int]: """Greedy argmax over schema-viable tokens using real engine logits. Returns (text, token_ids, finish_reason, prompt_len, gen_len). ``finish_reason`` is ``"stop"`` only for a complete, validating document; otherwise ``"length"`` (budget/context spent). """ schema = None if target is True else target config = getattr(engine, "config", None) vocab = config.get("vocab_size") if isinstance(config, dict) else None if not isinstance(vocab, int) or vocab <= 0: raise _APIError(400, "This engine does not expose a vocabulary; structured output is unsupported.", "unsupported", "response_format") logits_fn = getattr(engine, "logits", None) if not callable(logits_fn): raise _APIError(400, "This engine does not expose logits; structured output is unsupported.", "unsupported", "response_format") try: context = int(engine.context_length) except (TypeError, ValueError): raise _APIError(400, "This engine does not expose a context length.", "unsupported", "response_format") from None decode = engine.tokenizer.decode eos = set(_engine_eos(engine)) budget = min(token_limit, max(context - len(prompt_ids), 0)) prefix = list(prompt_ids) generated: list[int] = [] text = "" if _complete_and_valid(text, target): return text, generated, "stop", len(prompt_ids), 0 finish = "length" for _ in range(budget): if cancel.is_set(): raise _APIError(503, "Request disconnected.", "request_cancelled") if len(prefix) >= context: break try: row = _log_softmax_rows(logits_fn(list(prefix))) except _APIError: raise except (ValueError, TypeError, RuntimeError) as exc: raise _APIError(400, f"Constrained decoding failed: {exc}.", "generation_failed") if len(row) != vocab: raise _APIError(503, "Engine logits do not match the model vocabulary.", "generation_failed") picked: tuple[int, str] | None = None for tid in _argsort_desc(row): try: full = decode(generated + [tid], skip_special_tokens=True, clean_up_tokenization_spaces=False) except Exception: continue if tid in eos and not _complete_and_valid(full, target): continue if _prefix_viable(full, schema): picked = (tid, full) break if picked is None: raise _APIError(503, "Constrained decoding found no valid continuation.", "generation_failed") tid, full = picked generated.append(tid) prefix.append(tid) text = full if _complete_and_valid(text, target): finish = "stop" break return text, generated, finish, len(prompt_ids), len(generated) # -- tool calling: real parse + real local execution + bounded loop ----------- def _eval_arith(node: ast.AST, depth: int = 0) -> Any: if depth > 64: raise ValueError("expression too deeply nested") if isinstance(node, ast.Expression): return _eval_arith(node.body, depth + 1) if isinstance(node, ast.Constant): if isinstance(node.value, bool) or not isinstance(node.value, (int, float)): raise ValueError("only plain numbers are supported") if isinstance(node.value, float) and not math.isfinite(node.value): raise ValueError("non-finite literal") return node.value if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)): value = _eval_arith(node.operand, depth + 1) return value if isinstance(node.op, ast.UAdd) else -value if isinstance(node, ast.BinOp) and isinstance( node.op, (ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, ast.Mod, ast.Pow)): left = _eval_arith(node.left, depth + 1) right = _eval_arith(node.right, depth + 1) try: if isinstance(node.op, ast.Add): return left + right if isinstance(node.op, ast.Sub): return left - right if isinstance(node.op, ast.Mult): return left * right if isinstance(node.op, ast.Div): return left / right if isinstance(node.op, ast.FloorDiv): return left // right if isinstance(node.op, ast.Mod): return left % right if isinstance(right, int) and right < 0 and isinstance(left, int): raise ValueError("negative integer exponents are unsupported") if isinstance(right, (int, float)) and abs(right) > 10**6: raise ValueError("exponent too large") return left ** right except ZeroDivisionError: raise ValueError("division by zero") from None except OverflowError: raise ValueError("arithmetic overflow") from None raise ValueError("only + - * / // % **, parentheses, and numbers are supported") def _tool_calc(arguments: Any) -> dict: """Exact local arithmetic: {"expression": "6*7"} -> {"result": 42}.""" if not isinstance(arguments, dict): raise ValueError("arguments must be an object") expression = arguments.get("expression", None) if not isinstance(expression, str) or not expression.strip(): raise ValueError("expression must be a nonempty string") if len(expression) > 4096: raise ValueError("expression too long") try: tree = ast.parse(expression, mode="eval") except SyntaxError as exc: raise ValueError(f"invalid expression: {exc.msg}") from None result = _eval_arith(tree) if isinstance(result, float) and not math.isfinite(result): raise ValueError("non-finite result") return {"result": result} _TOOL_REGISTRY: dict[str, tuple[Any, dict, str]] = { "calc": (_tool_calc, {"type": "object", "properties": {"expression": {"type": "string"}}, "required": ["expression"], "additionalProperties": False}, "Evaluate a decimal arithmetic expression (+ - * / // % **)."), } def _parse_tool_call(text: str) -> tuple[str, dict] | None: """Parse a ``{"name", "arguments"}`` tool call from assistant text. Accepts raw JSON or a single ```json fenced block; returns None when the text is plain assistant content (not a tool call). """ stripped = text.strip() if stripped.startswith("```"): lines = stripped.split("\n") try: end = next(i for i, line in enumerate(lines[1:], 1) if line.strip().startswith("```")) except StopIteration: return None stripped = "\n".join(lines[1:end]).strip() if stripped.lower().startswith("json"): stripped = stripped[4:].strip() candidate: Any = None if stripped.startswith("{"): try: candidate = json.loads(stripped) except ValueError: candidate = None if candidate is None: start, end = stripped.find("{"), stripped.rfind("}") if 0 <= start < end: try: candidate = json.loads(stripped[start:end + 1]) except ValueError: candidate = None if not isinstance(candidate, dict): return None name = candidate.get("name") arguments = candidate.get("arguments", {}) if not isinstance(name, str) or not name: return None if isinstance(arguments, str): try: arguments = json.loads(arguments) except ValueError: return None if not isinstance(arguments, dict): return None return name, arguments def _validate_a2_request(body: _CompletionRequest | _ChatRequest, engine: Any) -> None: """Fail-fast 400s for Stream A2 options (before queue admission).""" chat = isinstance(body, _ChatRequest) if _wants_logprobs(body): if body.stream: raise _APIError(400, "logprobs are not supported with stream=true; use non-streaming.", "invalid_parameter", "logprobs") if _needs_constrained(body) or _needs_tools(body): raise _APIError(400, "logprobs cannot be combined with response_format tools yet.", "invalid_parameter", "logprobs") if not callable(getattr(engine, "logits", None)): raise _APIError(400, "This engine does not expose logits; logprobs are unsupported.", "unsupported", "logprobs") target = _constrained_target(body) if chat else None if target is not None: if not chat: # pragma: no cover - completions carry no response_format raise _APIError(400, "response_format is only supported on chat completions.", "invalid_parameter", "response_format") assert isinstance(body, _ChatRequest) if body.stream: raise _APIError(400, "response_format json is not supported with stream=true; use non-streaming.", "invalid_parameter", "response_format") if body.temperature != 0 or body.top_p != 1.0 or body.top_k != 0 or body.spec_k != 0: raise _APIError(400, "response_format json requires greedy decoding " "(temperature=0, top_p=1, top_k=0, spec_k=0).", "invalid_parameter", "response_format") if body.stop: raise _APIError(400, "response_format json cannot be combined with stop strings.", "invalid_parameter", "response_format") if _is_pool_or_batch_engine(engine): raise _APIError(400, "Structured output is not supported for pool/batch engines.", "unsupported", "response_format") if target is not True: reason = _schema_support_error(target) if reason is not None: raise _APIError(400, f"Unsupported response_format schema: {reason}.", "unsupported_schema", "response_format") if not callable(getattr(engine, "logits", None)): raise _APIError(400, "This engine does not expose logits; structured output is unsupported.", "unsupported", "response_format") if chat and _needs_tools(body): assert isinstance(body, _ChatRequest) if body.stream: raise _APIError(400, "tools are not supported with stream=true; use non-streaming.", "invalid_parameter", "tools") if _is_pool_or_batch_engine(engine): raise _APIError(400, "Tool calling is not supported for pool/batch engines.", "unsupported", "tools") assert body.tools is not None for tool in body.tools: name = tool.function.name if name not in _TOOL_REGISTRY: raise _APIError(400, f"Unknown tool {name!r}: the server executes only its built-in " "registry and will not fake external tools.", "unknown_tool", "tools") declared = tool.function.parameters if declared is not None: reason = _schema_support_error(declared) if reason is not None: raise _APIError(400, f"Unsupported tool parameters schema: {reason}.", "unsupported_schema", "tools") choice = body.tool_choice if isinstance(choice, _NamedToolChoice): wanted = choice.function.name if all(tool.function.name != wanted for tool in body.tools): raise _APIError(400, f"tool_choice names {wanted!r} which is not in tools.", "invalid_parameter", "tool_choice") # -- Stream A1: paged-KV single-stream + FIFO micro-batching ------------------- #: Dense-only paged decode; Surjo/FWKV keep their own session types. _PAGED_EXCLUDED_MODEL_TYPES = frozenset({"surjo", "fwkv"}) #: Max co-queued jobs coalesced into one BatchSession micro-batch. _MICRO_BATCH_MAX = 8 def _serve_paged_enabled() -> bool: # Default-on since Stream S2 (shared BlockPool, 0.999x dense single-stream): return os.environ.get("CISM_SERVE_PAGED", "1") != "0" def _serve_prefixcache_enabled() -> bool: """Kill switch: ``CISM_SERVE_PREFIXCACHE=0`` disables automatic reuse. Default-on (Stream S3): the paged single-stream path matches each arrival against a per-engine radix prefix index and fork-shares the longest cached prefix via S2 copy-on-write, prefilling only the suffix. Misses and stale anchors fall back to full prefill (exact). """ return os.environ.get("CISM_SERVE_PREFIXCACHE", "1") != "0" def _prefix_cache_max_anchors(default: int = 8) -> int: """Retained prefix-anchor sessions per engine (pool-pressure bound).""" try: value = int(os.environ.get("CISM_SERVE_PREFIXCACHE_ANCHORS", str(default))) except ValueError: return default return value if value > 0 else default #: Priority demotion for retained prefix anchors: under pool pressure S2's #: ``evict_one`` frees the lowest-priority entry first, so idle anchors #: (never decoded again, only forked from) are evicted before active work. _PREFIX_ANCHOR_PRIORITY = -8 def _prefix_cache_for(engine: Any) -> Any | None: """Per-engine S3 radix index, or None when the path is ineligible. The cache lives on the engine object itself (freed with it -- no global registry to leak). Requires the S3 kill switch on, the paged path eligible (``_can_paged``), and a native fork entry point. """ if not _serve_prefixcache_enabled() or not _can_paged(engine): return None native = getattr(engine, "_native", None) if native is None or not callable(getattr(native, "create_paged_fork_session", None)): return None cache = getattr(engine, "_prefix_cache", None) if cache is None: try: from .paging import PrefixCache cache = PrefixCache(max_anchors=_prefix_cache_max_anchors()) engine._prefix_cache = cache except (AttributeError, TypeError, ValueError): return None return cache def prefix_cache_stats(engine: Any | None = None) -> dict[str, Any]: """Hit/miss counters for the S3 automatic prefix cache (ops/tests). With an engine returns its radix-index stats (``enabled: False`` when the kill switch or eligibility excludes it); with None reports whether the switch is on. Never raises: stats must not break serving. """ try: if engine is None: return {"enabled": _serve_prefixcache_enabled()} cache = getattr(engine, "_prefix_cache", None) if cache is None: return {"enabled": _serve_prefixcache_enabled() and _can_paged(engine), "hits": 0, "misses": 0, "fork_failures": 0, "inserts": 0, "evictions": 0, "anchors": 0, "hit_rate": None} stats = dict(cache.stats()) stats["enabled"] = _serve_prefixcache_enabled() return stats except Exception: return {"enabled": False} def _serve_microbatch_enabled() -> bool: return os.environ.get("CISM_SERVE_MICROBATCH", "1") != "0" def _serve_continuous_enabled() -> bool: """Kill switch: ``CISM_SERVE_CONTINUOUS=0`` falls back to A1 FIFO.""" return os.environ.get("CISM_SERVE_CONTINUOUS", "1") != "0" def _continuous_prefill_chunk(default: int) -> int: """New prompt tokens admitted per epoch rebuild (chunked prefill).""" try: value = int(os.environ.get("CISM_SERVE_PREFILL_CHUNK", str(default))) except ValueError: return default return value if value > 0 else default def _continuous_epoch_steps() -> int: """Max decode steps before an epoch rotates (starvation backstop).""" try: value = int(os.environ.get("CISM_SERVE_EPOCH_STEPS", "1024")) except ValueError: return 1024 return value if value > 0 else 1024 def _continuous_max_wait_s() -> float: """Starvation guard: bound on an admittable waiter's queue wait.""" try: value = float(os.environ.get("CISM_SERVE_MAX_WAIT_S", "30")) except ValueError: return 30.0 return value if value > 0 else 30.0 def _continuous_join_defer_tokens() -> int: """Skip a join rebuild when the live batch is this close to draining. A newcomer arriving with <= N tokens left in every active member waits for the drain instead of forcing a recompute-prefill that saves nothing (starved waiters bypass the deferral via the max-wait guard). """ try: value = int(os.environ.get("CISM_SERVE_JOIN_DEFER_TOKENS", "2")) except ValueError: return 2 return value if value >= 0 else 2 def _cont_encode(engine: Any, body: _CompletionRequest | _ChatRequest, ) -> tuple[list[int], list[str] | None]: """Validate + encode one epoch arrival (per-job isolation). Raises ``_APIError`` (400 invalid/context) without touching any batch. """ try: if isinstance(body, _ChatRequest): prompt = engine.format_chat( [message.model_dump() for message in body.messages]) prompt_ids = engine.encode(prompt) else: prompt_ids = engine.encode(body.prompt) except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None if len(prompt_ids) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") stop = ([body.stop] if isinstance(body.stop, str) else (list(body.stop) if body.stop else None)) return prompt_ids, stop def _can_paged(engine: Any) -> bool: """True when ``engine`` can serve a job on the native paged-KV path. Dense ``Engine`` only: needs ``_native.create_paged_session``, ``tokenizer``, ``config`` (model_type not surjo/fwkv), ``_eos``, and ``_generation_lock``. Pool/batch engines, test doubles, and Surjo/FWKV return False and keep their existing paths bitwise identical. """ native = getattr(engine, "_native", None) if native is None or not callable(getattr(native, "create_paged_session", None)): return False if getattr(engine, "tokenizer", None) is None: return False if getattr(engine, "_generation_lock", None) is None: return False config = getattr(engine, "config", None) if isinstance(config, dict) and config.get("model_type") in _PAGED_EXCLUDED_MODEL_TYPES: return False return isinstance(getattr(engine, "_eos", None), list) def _stream_paged(engine: Any, tokens: list[int], *, max_new_tokens: int = 128, temperature: float = 0.0, top_p: float = 1.0, top_k: int = 0, seed: int = 0, stop: list[str] | None = None, cancel: threading.Event | None = None, chunk_size: int = 16) -> Iterator[_Chunk]: """Yield :class:`_Chunk` deltas over a native ``PagedSession``. Exact-math mirror of ``Engine.stream`` for ``spec_k == 0`` (same validation, same context clamp, same lock/cancel/monitor discipline, same stop-string truncation with partial-prefix holdback, same incremental decode check). Native ``test_paging_native.py`` proves paged blocked-SDPA is bitwise-identical to the dense ``Session`` path, so greedy and sampled outputs match ``engine.stream`` token-for-token; spec tallies report 0 (no speculation on the paged path -- callers route ``spec_k > 0`` to the dense path). """ if type(chunk_size) is not int or not 1 <= chunk_size <= 1024: raise ValueError("chunk_size must be an integer in [1, 1024]") if type(max_new_tokens) is not int or max_new_tokens < 0: raise ValueError("max_new_tokens must be a nonnegative integer") if type(seed) is not int or not 0 <= seed < 2**64: raise ValueError("seed must be an unsigned 64-bit integer") if type(top_k) is not int or top_k < 0: raise ValueError("top_k must be a nonnegative integer") for name, value in (("temperature", temperature), ("top_p", top_p)): if type(value) not in (float, int) or not math.isfinite(value): raise ValueError(f"{name} must be finite") if temperature < 0 or not 0 < top_p <= 1: raise ValueError("temperature must be nonnegative and top_p in (0, 1]") if stop is not None and (not isinstance(stop, list) or any(not isinstance(s, str) or not s for s in stop)): raise ValueError("stop must be a list of nonempty strings") stops = list(stop or []) tokens = list(tokens) vocab = engine.config.get("vocab_size") if isinstance(engine.config, dict) else None for token in tokens: if type(token) is not int or (vocab is not None and not 0 <= token < vocab): raise ValueError("Token IDs must be integers within the model vocabulary") if len(tokens) > engine.context_length: raise ValueError("Prompt length exceeds the model context length") max_new_tokens = min(max_new_tokens, engine.context_length - len(tokens)) cancel = cancel if cancel is not None else threading.Event() lock = engine._generation_lock while not lock.acquire(timeout=0.05): if cancel.is_set(): return session = None monitor_done = threading.Event() monitor = None monitor_started = False released = False def cleanup() -> None: nonlocal released if released: return try: monitor_done.set() if session is not None: session.cancel() if monitor_started: monitor.join() finally: released = True lock.release() try: if cancel.is_set(): return # Stream S3: longest cached token prefix is fork-shared via S2 # copy-on-write (child prefills only the suffix). Sampling config # is irrelevant to reuse (KV is sampling-independent); the fork # re-checks token equality, so a hit is exact by construction. # Stale anchors (S2 evict_one / pool reset) raise here -> invalidate # the branch and fall back to full prefill (still exact). cache = _prefix_cache_for(engine) prefix_hit_len = 0 if cache is not None: anchor, hit_len = cache.match(tokens) if anchor is not None: try: session = engine._native.create_paged_fork_session( tokens, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=list(engine._eos), src=anchor, prefix_len=hit_len) prefix_hit_len = int(hit_len) cache.note_hit() except (ValueError, RuntimeError): cache.invalidate_session(anchor) cache.note_fork_failed() cache.note_miss() session = None else: cache.note_miss() if session is None: session = engine._native.create_paged_session( tokens, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=list(engine._eos)) def watch_cancel() -> None: while not monitor_done.wait(0.01): if cancel.is_set(): session.cancel() return monitor = threading.Thread(target=watch_cancel, name="cism-cancel", daemon=True) monitor.start() monitor_started = True generated: list[int] = [] emitted = "" first = True while True: if cancel.is_set(): session.cancel() return batch = session.next_tokens(1 if first or stops else chunk_size) first = False generated.extend(batch) reason = session.finish_reason or None if reason == "cancelled": return text = engine.tokenizer.decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False) if reason is None: text = text.rstrip("\ufffd") matches = [text.find(s) for s in stops if s in text] if matches: text = text[:min(matches)] reason = "stop" session.cancel() elif reason is None: hold = 0 for sequence in stops: for size in range(1, min(len(sequence), len(text) + 1)): if text.endswith(sequence[:size]): hold = max(hold, size) if hold: text = text[:-hold] if not text.startswith(emitted): raise ValueError("Tokenizer decoding rewrote emitted text; incremental streaming is unsupported") delta = text[len(emitted):] emitted = text if reason is not None: # Retain this finished session as an S3 fork anchor (pool KV # stays resident; demoted priority so pressure evicts idle # anchors before active work). Cancelled streams never anchor. if (cache is not None and reason != "cancelled" and not cancel.is_set()): try: session.set_priority(_PREFIX_ANCHOR_PRIORITY) cache.insert(tokens, session) except (ValueError, RuntimeError): pass cleanup() yield _Chunk(delta, reason, len(tokens), len(generated), token_ids=list(batch), prefix_hit=bool(prefix_hit_len), prefix_len=prefix_hit_len) if reason is not None: break if not batch: raise RuntimeError("Native decoder made no progress") finally: cleanup() def _micro_batch_eligible(body: _CompletionRequest | _ChatRequest) -> bool: """Only non-streaming greedy jobs may join a micro-batch. Greedy (``temperature == 0``) decode is seed-independent, so coalescing jobs with different seeds/sampling trivia stays bitwise-identical to serial execution. ``spec_k`` has no batch path; streaming jobs keep the incremental serial path (v0 limit). Stream A2 opt-in jobs (logprobs rescoring, constrained JSON, tool loops) always run serially so the shared fast path is untouched by construction. """ return (not body.stream and body.spec_k == 0 and body.temperature == 0 and not _wants_logprobs(body) and not _needs_constrained(body) and not _needs_tools(body)) def _sampled_batch_eligible(body: _CompletionRequest | _ChatRequest) -> bool: """Only non-streaming sampled jobs may join a sampled micro-batch epoch. Same serial-path exclusions as greedy (streaming, ``spec_k``, A2 opt-ins). ``temperature > 0`` (``== 0`` is the greedy path). Sampling params must be in their valid ranges; per-sequence seeds ride along in the native session, so coalescing same-param jobs stays token-identical to serial execution. Sampled epochs never take mid-decode joins and never rotate (a rebuild would restart per-seq RNG streams) -- enforced in ``fits`` and the rotation gate, not here. """ try: temperature = float(body.temperature) top_p = float(body.top_p) top_k = int(body.top_k) except (TypeError, ValueError): return False return (not body.stream and body.spec_k == 0 and temperature > 0 and temperature == temperature and 0 < top_p <= 1.0 and top_k >= 0 and not _wants_logprobs(body) and not _needs_constrained(body) and not _needs_tools(body)) def _micro_batch_key(body: _CompletionRequest | _ChatRequest) -> tuple: stop = [body.stop] if isinstance(body.stop, str) else tuple(body.stop or ()) return (tuple(stop), body.token_limit) def _is_pool_or_batch_engine(engine: Any) -> bool: """True for EnginePool/BatchEngine (generate_batch, no stream). Plain Engine/FakeEngine expose ``stream``; pool/batch engines expose ``generate_batch`` without ``stream``. This keeps the single-stream ``_Worker`` path bitwise identical for plain engines. """ return (hasattr(engine, "generate_batch") and callable(getattr(engine, "generate_batch")) and not (hasattr(engine, "stream") and callable(getattr(engine, "stream")))) def _engine_eos(engine: Any) -> list[int]: eos = getattr(engine, "_eos", None) if isinstance(eos, list): return list(eos) config = getattr(engine, "config", None) tokenizer = getattr(engine, "tokenizer", None) raw = None if isinstance(config, dict): raw = config.get("eos_token_id", None) if raw is None: raw = getattr(tokenizer, "eos_token_id", None) if raw is None: return [] return raw if isinstance(raw, list) else [raw] def _can_native_batch(engine: Any) -> bool: """Plain dense Engine with a native BatchSession path (no Surjo).""" native = getattr(engine, "_native", None) if native is None or not hasattr(native, "create_batch_session"): return False config = getattr(engine, "config", None) if isinstance(config, dict) and config.get("model_type") == "surjo": return False return True def _run_batch_chunk_sync(engine: Any, chunk_prompts: list, token_limit: int, temperature: float, top_p: float, top_k: int, seed: int, stop_list: list[str] | None): """Run one scheduler chunk synchronously, preserving chunk order. Dispatch: - BatchEngine/EnginePool (``_is_pool_or_batch_engine``): use its ``generate_batch`` (pool = N weight streams in parallel threads; batch = single weight stream via native ``BatchSession``). - Plain dense Engine with native batch support: single-weight-stream ``_native.create_batch_session`` under ``_generation_lock`` so concurrent batches for *unlike* models (different engines/locks) run in parallel while same-engine batches serialize. - Otherwise (Surjo/no native batch): sequential ``engine.generate`` loop (order-preserving fallback). Returns a list of ``(text, token_ids, finish_reason, prompt_len, gen_len)``. """ stops = list(stop_list or []) if _is_pool_or_batch_engine(engine): try: results = engine.generate_batch( chunk_prompts, max_new_tokens=token_limit, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, **({"stop": stops} if stops else {})) except TypeError: # Older pool mocks without stop= support. results = engine.generate_batch( chunk_prompts, max_new_tokens=token_limit, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed) if stops: from .paging import ContinuousSchedulerV0 as _S fixed = [] for r in results: text, reason = _S.apply_stop(r.text, stops) fixed.append(type(r)(text, r.token_ids, reason or r.finish_reason, r.prompt_tokens, r.completion_tokens, getattr(r, "spec_proposed", 0), getattr(r, "spec_accepted", 0))) results = fixed return [(r.text, list(r.token_ids), r.finish_reason, r.prompt_tokens, r.completion_tokens) for r in results] if _can_native_batch(engine): # Encode via the engine (BOS fallback + vocab checks identical to # the single-stream worker). encoded = [engine.encode(p) for p in chunk_prompts] context = int(engine.context_length) effective = min([token_limit] + [context - len(t) for t in encoded]) eos = _engine_eos(engine) lock = getattr(engine, "_generation_lock", None) acquire = lock.acquire if lock is not None else None # Blocking acquire is fine inside asyncio.to_thread. if acquire is not None: acquire() try: session = engine._native.create_batch_session( encoded, max_new_tokens=effective, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, eos_token_ids=eos) batched = session.next_tokens(effective) reasons = list(session.finish_reasons) finally: if lock is not None: lock.release() decode = engine.tokenizer.decode out = [] for tokens, token_ids, reason in zip(encoded, batched, reasons): text = decode(list(token_ids), skip_special_tokens=True, clean_up_tokenization_spaces=False) if stops: hits = [text.find(s) for s in stops if s in text] if hits: text = text[:min(hits)] reason = "stop" out.append((text, list(token_ids), reason, len(tokens), len(token_ids))) return out # Fallback: sequential generate/stream (Surjo / custom / test doubles). # Prefer Engine.generate when available; otherwise collect Engine.stream # chunks (FakeEngine in tests only implements stream). if hasattr(engine, "generate") and callable(getattr(engine, "generate")): out = [] for prompt in chunk_prompts: try: r = engine.generate(prompt, max_new_tokens=token_limit, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, stop=stops or None) except TypeError: r = engine.generate(prompt, max_new_tokens=token_limit, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed) if stops: from .paging import ContinuousSchedulerV0 as _S2 text, reason = _S2.apply_stop(r.text, stops) r = type(r)(text, r.token_ids, reason or r.finish_reason, r.prompt_tokens, r.completion_tokens, getattr(r, "spec_proposed", 0), getattr(r, "spec_accepted", 0)) out.append((r.text, list(r.token_ids), r.finish_reason, r.prompt_tokens, r.completion_tokens)) return out # Stream-only backends (e.g. FakeEngine test double). out = [] for prompt in chunk_prompts: parts: list[str] = [] ids: list[int] = [] finish = None prompt_n = 0 gen_n = 0 for chunk in engine.stream(prompt, max_new_tokens=token_limit, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, stop=stops or None, cancel=threading.Event(), chunk_size=16, spec_k=0): parts.append(chunk.text) ids.extend(chunk.token_ids) prompt_n = chunk.prompt_tokens gen_n = chunk.completion_tokens if chunk.finish_reason is not None: finish = chunk.finish_reason break if finish is None: # Mirror the FIFO worker: a missing terminal chunk is a failure. raise RuntimeError("Generation did not complete") out.append(("".join(parts), ids, finish, prompt_n, gen_n)) return out def create_app(engine: _Engine | dict[str, _Engine], api_key: str | None = None, queue_size: int = 8, *, allow_remote_code_load: bool = False, load_lock_pages: bool = False, load_keep_warm: float = 0.0, load_threads_default: int = 2, load_local_files_only: bool = False, load_act_precision: str = "fp32", max_batch_tokens: int = 4096, scheduler_blocks: int = 512, prefill_chunk_tokens: int = 512) -> FastAPI: """Create an app around engine(s) owned by the caller. A mapping of model ID to engine serves several loaded models; requests route by ``model`` field through one shared FIFO worker, so only one generation runs at a time regardless of how many models are loaded. Passing a single engine keeps the original behavior. Phase2 batching (``POST /v1/batch``):: {"model": "", "prompts": ["a", "b", [1, 2]], "max_tokens": 64} returns order-preserving ``choices`` (``choices[i]`` is ``prompts[i]``). Dispatch: same-model same-quant batches use one native ``BatchSession`` weight stream (``BatchEngine`` or plain dense ``Engine._native``); pool engines use ``EnginePool.generate_batch`` (N weight streams); concurrent batches for *unlike* models overlap in ``asyncio.to_thread`` (separate locks), while same-engine batches serialize. Single-stream ``/v1/completions`` + ``/v1/chat/completions`` for plain engines still go through the shared ``_Worker`` on the dense ``engine.stream`` path (unregressed by construction); the bitwise-identical native paged-KV path (``create_paged_session``) is opt-in via ``CISM_SERVE_PAGED=1``. Eligible co-queued greedy jobs on dense engines with a steppable native ``BatchSession`` run in an S1 iteration-level epoch (per-step add/evict, oldest-first + starvation guard, chunked prefill; ``CISM_SERVE_CONTINUOUS=0`` falls back to the A1 FIFO micro-batch); streaming/sampled/speculative/A2 jobs always run serially. Pool/batch engines serve single-stream via the same batch machinery with ``B=1``. ``/v1/batch`` chunks FIFO waiting->running requests by ``max_batch_tokens`` (alias ``max_batched_tokens``, ~512 operating point) plus ``prefill_chunk_tokens`` (128-512) via ``ContinuousSchedulerV0`` (block16, per-seq EOS/stop, BlockManager eviction, recompute-requeue); see its docstring for v0 limits (no preempt-swap, no cross-seq sharing). Runtime model management: ``POST /v1/models/load`` downloads and loads a HF model on a background thread (one load at a time). The load applies ``load_lock_pages``/``load_keep_warm``/``load_threads_default`` to newly loaded models. ``trust_remote_code=True`` in a load request is rejected unless ``allow_remote_code_load`` is enabled server-side, because loading executes downloaded tokenizer code with this process's permissions. ``POST /v1/models/unload`` releases idle models (409 if generating). ``None`` reads CISM_API_KEY once at creation; an empty string disables auth. Health/readiness probes are public. ``queue_size`` bounds all admitted jobs, including active inference and responses still being consumed. Cancellation is cooperative; shutdown waits at most one second for the daemon worker. """ if type(queue_size) is not int or queue_size < 1: raise ValueError("queue_size must be a positive integer") if type(max_batch_tokens) is not int or max_batch_tokens < 1: raise ValueError("max_batch_tokens must be a positive integer") if type(scheduler_blocks) is not int or scheduler_blocks < 1: raise ValueError("scheduler_blocks must be a positive integer") if type(prefill_chunk_tokens) is not int or prefill_chunk_tokens < 1: raise ValueError("prefill_chunk_tokens must be a positive integer") if isinstance(engine, dict): model_map = dict(engine) if not model_map: raise ValueError("At least one model is required") else: model_map = {engine.model_id: engine} registry_lock = threading.Lock() load_status: dict[str, str] = {} load_in_flight = threading.Event() bench_state: dict[str, Any] = {"status": "idle", "progress": "", "result": None, "error": None} bench_in_flight = threading.Event() if api_key is None: api_key = os.environ.get("CISM_API_KEY") if api_key is not None and not isinstance(api_key, str): raise TypeError("api_key must be a string or None") secret = api_key.encode("utf-8") if api_key else None @asynccontextmanager async def lifespan(app: FastAPI): previous = app.state.worker if previous is not None and previous.thread.is_alive(): raise RuntimeError("The previous inference worker is still running") # Pool/Batch engines never use the FIFO worker's default engine (they # bypass it via the batch machinery), so keep it None for them to # avoid ever calling Pool.generate_batch as Engine.stream. default_engine = next(iter(model_map.values())) if len(model_map) == 1 else None if default_engine is not None and _is_pool_or_batch_engine(default_engine): default_engine = None worker = _Worker(default_engine, queue_size) worker.scheduler_config = { "scheduler_blocks": scheduler_blocks, "max_batch_tokens": max_batch_tokens, "prefill_chunk_tokens": prefill_chunk_tokens, } app.state.worker = worker worker.thread.start() try: yield finally: worker.stop() deadline = time.monotonic() + 1.0 while worker.thread.is_alive() and time.monotonic() < deadline: await asyncio.sleep(0.01) app = FastAPI( title="CISM Text Generation API", lifespan=lifespan, description="Text completions with exact-math logprobs, constrained JSON, and registry tool calls.", docs_url=None, redoc_url=None, openapi_url=None, ) app.state.worker = None @app.exception_handler(_APIError) async def api_error_handler(request: Request, error: _APIError): headers = {"WWW-Authenticate": "Bearer"} if error.status == 401 else None if error.status == 429: headers = {"Retry-After": "1"} return JSONResponse(error.body, status_code=error.status, headers=headers) @app.exception_handler(RequestValidationError) async def validation_error_handler(request: Request, error: RequestValidationError): # Pydantic errors include raw inputs and exception context. Never expose them. return JSONResponse(_APIError( 400, "Invalid request body or unsupported parameter. Only text generation is supported.", "invalid_request", ).body, status_code=400) @app.exception_handler(HTTPException) async def http_error_handler(request: Request, error: HTTPException): message = "Not found." if error.status_code == 404 else "HTTP request failed." return JSONResponse(_APIError(error.status_code, message, "http_error").body, status_code=error.status_code) @app.exception_handler(Exception) async def unexpected_error_handler(request: Request, error: Exception): return JSONResponse(_APIError(503, "Server unavailable.", "server_unavailable").body, status_code=503) async def authorize(request: Request) -> None: if secret is None: return scheme, _, credential = request.headers.get("authorization", "").partition(" ") matches = hmac.compare_digest(credential.encode("utf-8"), secret) if scheme.lower() != "bearer" or not matches: raise _APIError(401, "Invalid API key.", "invalid_api_key") router = APIRouter(prefix="/v1", dependencies=[Depends(authorize)]) @app.get("/health") async def health(): return {"status": "ok"} @app.get("/ready") async def ready(): worker = app.state.worker if worker is None or not worker.ready: raise _APIError(503, "Server is not ready.", "server_unavailable") return {"status": "ready"} @app.get("/", response_class=HTMLResponse) async def index(): # Static, dependency-free tester page. It never contains the API key; # visitors paste their own credential, like any other API client. # no-cache keeps model dropdowns fresh across server restarts. return HTMLResponse(_WEBUI_HTML, headers={"Cache-Control": "no-cache"}) @router.get("/models") async def models(): with registry_lock: items = list(model_map.items()) data = [] for model_id, model in items: entry = {"id": model_id, "object": "model", "created": 0, "owned_by": "cism", "context_length": model.context_length} # Optional engine introspection (real engines provide info). info = getattr(model, "info", None) if callable(info): info = info() if isinstance(info, dict): for key in ("precision", "threads", "kernel", "architecture", "quantization", "act_precision"): if key in info: entry[key] = info[key] if isinstance(info.get("weight_bytes"), (int, float)): entry["weight_mb"] = round(info["weight_bytes"] / (1024 * 1024), 2) data.append(entry) return {"object": "list", "data": data} @router.get("/models/loading") async def models_loading(): with registry_lock: return {"loading": dict(load_status)} def _apply_load_options(loaded_engine) -> None: if load_lock_pages and loaded_engine.lock_pages() == 0: load_status[loaded_engine.model_id] = "loaded (page locking failed)" if load_keep_warm: loaded_engine.start_keep_warm(load_keep_warm) def _register_loaded_engine(loaded_engine) -> None: with registry_lock: model_map[loaded_engine.model_id] = loaded_engine load_status.pop(loaded_engine.model_id, None) _apply_load_options(loaded_engine) def _load_via_loader(model: str, precision: str, threads: int, revision: str | None = None, trust_remote_code: bool = False, act_precision: str | None = None, reserved: bool = False): """Load through the single-slot loader; returns None if it is busy. ``reserved=True`` means the caller already holds the load slot (the ``/models/load`` endpoint reserves it under the registry lock before starting its worker); without it the flag would read as someone else's load and the worker would refuse its own reservation. """ with registry_lock: existing = model_map.get(model) existing_info = getattr(existing, "info", {}) if existing is not None else {} if callable(existing_info): existing_info = existing_info() if not isinstance(existing_info, dict): existing_info = {} want_act = act_precision or load_act_precision if (existing is not None and existing_info.get("precision") == precision and existing_info.get("act_precision", "fp32") == want_act and existing_info.get("threads", threads) == threads): if reserved: load_in_flight.clear() # the endpoint's reservation is done return existing if load_in_flight.is_set() and not reserved: return None load_in_flight.set() load_status[model] = "loading" try: from cism.engine import Engine loaded_engine = Engine.from_pretrained( model, precision=precision, revision=revision, threads=threads, trust_remote_code=trust_remote_code, local_files_only=load_local_files_only, act_precision=want_act) except Exception as error: with registry_lock: load_status[model] = f"failed: {type(error).__name__}: {error}" raise finally: # The loader releases its own slot on every exit path, whoever # called it (the /models/load worker or the bench harness). load_in_flight.clear() with registry_lock: if load_status.get(model) == "loading": load_status.pop(model, None) _register_loaded_engine(loaded_engine) return loaded_engine def _load_worker(request_body: _LoadRequest) -> None: try: _load_via_loader(request_body.model, request_body.precision, request_body.threads, revision=request_body.revision, trust_remote_code=request_body.trust_remote_code, act_precision=request_body.act_precision, reserved=True) except Exception: pass # failure is recorded in load_status for the pollers finally: load_in_flight.clear() def _bench_worker(request_body: _AutoBenchRequest) -> None: from . import autobench def report_progress(message: str) -> None: with registry_lock: bench_state["progress"] = message try: result = autobench.run_benchmark( request_body.model, precision=request_body.precision, threads=request_body.threads, max_new_tokens=request_body.max_tokens, local_files_only=load_local_files_only, thread_sweep=request_body.thread_sweep, spec_ks=(0, 2, 4, 8, 12, 16) if request_body.spec_sweep else (0,), act_modes=("fp32", "int8") if request_body.bench_act_int8 else ("fp32",), include_torch=request_body.include_torch, include_compile=request_body.compile, load_engine=lambda requested, threads=None, act_precision=None: _load_via_loader( request_body.model, requested or request_body.precision, threads or request_body.threads, act_precision=act_precision), progress=report_progress) with registry_lock: bench_state.update(status="done", result=result, progress="", error=None) except Exception as error: with registry_lock: bench_state.update(status="failed", error=f"{type(error).__name__}: {error}", progress="", result=None) finally: bench_in_flight.clear() @router.get("/bench/auto") async def bench_auto_status(): with registry_lock: return dict(bench_state) @router.post("/bench/auto", status_code=202) async def bench_auto(body: _AutoBenchRequest): with registry_lock: if bench_in_flight.is_set(): raise _APIError(409, "A benchmark is already running.", "bench_in_progress") if body.model not in model_map and load_in_flight.is_set(): raise _APIError(409, "A model load is already in progress.", "load_in_progress") bench_in_flight.set() bench_state.update(status="running", progress="starting", result=None, error=None) threading.Thread(target=_bench_worker, args=(body,), name="cism-bench", daemon=True).start() return {"status": "running", "model": body.model, "precision": body.precision} @router.post("/models/load", status_code=202) async def models_load(body: _LoadRequest): with registry_lock: if body.model in model_map: # Same weights + activation mode: no-op. A different precision # (the shootout's load-bench cycle) must fall through and reload. existing = model_map[body.model] existing_info = getattr(existing, "info", {}) if callable(existing_info): existing_info = existing_info() if not isinstance(existing_info, dict): existing_info = {} want_act = body.act_precision or load_act_precision if (existing_info.get("precision") == body.precision and existing_info.get("act_precision", "fp32") == want_act): return {"model": body.model, "status": "ready"} if body.trust_remote_code and not allow_remote_code_load: raise _APIError(403, "trust_remote_code is disabled on this server; restart with " "--allow-remote-code-load to permit downloaded tokenizer code.", "remote_code_disabled") if load_in_flight.is_set(): raise _APIError(429, "Another model load is already in progress.", "load_in_progress") load_in_flight.set() # Stale failures from earlier attempts would otherwise resurface in # every later poll; the UI has already shown them once. for failed_id in [m for m, s in load_status.items() if s.startswith("failed")]: del load_status[failed_id] load_status[body.model] = "loading" threading.Thread(target=_load_worker, args=(body,), name="cism-loader", daemon=True).start() return {"model": body.model, "status": "loading"} @router.post("/models/unload") async def models_unload(body: _UnloadRequest): with registry_lock: loaded_engine = model_map.pop(body.model, None) if loaded_engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") load_status.pop(body.model, None) # Idle check: an active generation holds the engine's generation lock. generation_lock = getattr(loaded_engine, "_generation_lock", None) if generation_lock is not None and not generation_lock.acquire(blocking=False): with registry_lock: model_map[body.model] = loaded_engine # keep serving in-flight work raise _APIError(409, "Model is generating; unload after it finishes.", "model_busy") try: info = loaded_engine.info if hasattr(loaded_engine, "info") else {} if callable(info): info = info() if isinstance(info, dict) and info.get("pages_locked"): loaded_engine.unlock_pages() if getattr(loaded_engine, "keep_warm_active", False): loaded_engine.stop_keep_warm() except Exception: with registry_lock: model_map[body.model] = loaded_engine raise _APIError(503, "Model unload failed; it remains available.", "unload_failed") if generation_lock is not None: generation_lock.release() return {"model": body.model, "status": "unloaded"} @router.post("/cache/report") async def cache_report(body: _CacheRequest): engine = model_map.get(body.model) if engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") def _run() -> dict: import numpy as np info = getattr(engine, "info", None) if callable(info): info = info() weight_bytes = info.get("weight_bytes") if isinstance(info, dict) else None if not isinstance(weight_bytes, (int, float)) or weight_bytes <= 0: raise _APIError(400, "This engine does not expose weight bytes.", "invalid_engine", "model") try: prompt = engine.encode("cache report probe") def timed_scan(repeats: int, evictions: int) -> float: # Unique 2x-L3 writes evict residency; the scan is timed alone. scrub = np.empty(body.scrub_mb * 1024 * 1024, dtype=np.uint8) samples = [] for _ in range(repeats): for _ in range(evictions): scrub[:] = 1 start = time.perf_counter() engine.scan() samples.append(time.perf_counter() - start) return statistics.median(samples) warm_seconds = timed_scan(20, evictions=0) cold_seconds = timed_scan(5, evictions=2) session = engine._native.create_session( prompt, max_new_tokens=body.decode_tokens, eos_token_ids=[]) session.next_tokens(1) start = time.perf_counter() generated = session.next_tokens(body.decode_tokens - 1) decode_seconds = time.perf_counter() - start except _APIError: raise except (ValueError, RuntimeError) as exc: raise _APIError(400, f"Cache probe failed: {exc}.", "cache_probe_failed") tokens_per_second = len(generated) / decode_seconds if generated and decode_seconds > 0 else 0 effective_gbps = weight_bytes * tokens_per_second / 1e9 warm_gbps = weight_bytes / warm_seconds / 1e9 if warm_seconds > 0 else 0 cold_gbps = weight_bytes / cold_seconds / 1e9 if cold_seconds > 0 else 0 ratio = cold_seconds / warm_seconds if warm_seconds > 0 else 0 if ratio >= 2: residency = "cache-resident when warm (cold scan is much slower)" elif ratio >= 1.3: residency = "partially cache-resident" else: residency = "no measurable residency benefit from this probe" if not tokens_per_second: bottleneck = "decode produced no tokens" elif effective_gbps <= cold_gbps: bottleneck = ("compute-bound: decode bandwidth is below even the cold DRAM scan; " "the cache is not the bottleneck, the kernels are") elif effective_gbps < warm_gbps * 0.8: bottleneck = "DRAM-bandwidth-limited: weights are effectively streaming from RAM" else: bottleneck = "bandwidth-limited at cache-resident rates" return {"model": engine.model_id, "weight_bytes": weight_bytes, "warm_scan_gbps": round(warm_gbps, 2), "cold_scan_gbps": round(cold_gbps, 2), "warm_to_cold_ratio": round(ratio, 2), "decode_tokens_per_second": round(tokens_per_second, 1), "decode_effective_gbps": round(effective_gbps, 2), "residency": residency, "decode_bottleneck": bottleneck} return await asyncio.to_thread(_run) async def _generate_single_via_batch(engine: Any, body: _CompletionRequest | _ChatRequest, request: Request): """Serve one completions/chat request via the batch machinery (B=1). Used only for pool/batch engines. Returns the same JSON/SSE shape as the FIFO worker path (non-streaming choices + usage; streaming SSE emits the full text as one delta for v0). Raises _APIError on validation/context failures with sanitized messages. """ chat = isinstance(body, _ChatRequest) stop = [body.stop] if isinstance(body.stop, str) else (list(body.stop) if body.stop else None) # Resolve the prompt without holding any lock (tokenizer-only). try: if chat: assert isinstance(body, _ChatRequest) prompt_tokens = engine.format_chat( [message.model_dump() for message in body.messages]) # format_chat may return tokens already; encode validates them. prompt: Any = engine.encode(prompt_tokens) else: assert isinstance(body, _CompletionRequest) prompt = body.prompt # Early context check for string prompts (encode to measure). if isinstance(prompt, str): prompt_tokens_probe = engine.encode(prompt) if len(prompt_tokens_probe) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") prompt = body.prompt except _APIError: raise except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None token_limit = body.token_limit def _run_one(): try: rows = _run_batch_chunk_sync( engine, [prompt], token_limit, body.temperature, body.top_p, body.top_k, body.seed, stop) except (ValueError, TypeError) as exc: msg = str(exc) if "context" in msg.lower() or "Prompt length" in msg: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") from None raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None if len(rows) != 1: raise _APIError(503, "Generation did not complete.", "generation_failed") return rows[0] try: text, _ids, finish, prompt_n, gen_n = await asyncio.to_thread(_run_one) except _APIError: raise except Exception: raise _APIError(503, "Generation failed.", "generation_failed") from None if finish is None: raise _APIError(503, "Generation did not complete.", "generation_failed") base = { "id": ("chatcmpl-" if chat else "cmpl-") + uuid.uuid4().hex, "object": "chat.completion" if (chat and not body.stream) else ( "chat.completion.chunk" if chat else "text_completion"), "created": int(time.time()), "model": engine.model_id, } usage = {"prompt_tokens": prompt_n, "completion_tokens": gen_n, "total_tokens": prompt_n + gen_n, "spec_proposed": 0, "spec_accepted": 0, "spec_accept_rate": None} if not body.stream: choice: dict[str, Any] = {"index": 0, "finish_reason": finish, "logprobs": None} if chat: choice["message"] = {"role": "assistant", "content": text} else: choice["text"] = text return {**base, "choices": [choice], "usage": usage} # Streaming via batch path (v0): single delta + terminal + DONE. include_usage = body.stream_options is not None and body.stream_options.include_usage def event(choices: list[dict], usage_payload: dict | None = None, completion_tokens: int | None = None) -> str: payload: dict[str, Any] = {**base, "choices": choices} if include_usage: payload["usage"] = usage_payload if completion_tokens is not None: payload["cism_completion_tokens"] = completion_tokens return "data: " + json.dumps(payload, ensure_ascii=True, allow_nan=False) + "\n\n" async def events(): try: if chat: yield event([{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], None, gen_n) if text: content = {"delta": {"content": text}} if chat else {"text": text, "logprobs": None} yield event([{"index": 0, **content, "finish_reason": None}], None, gen_n) content = {"delta": {}} if chat else {"text": "", "logprobs": None} yield event([{"index": 0, **content, "finish_reason": finish}], None, gen_n) if include_usage: yield event([], usage) yield "data: [DONE]\n\n" except Exception: error = _APIError(503, "Generation failed.", "generation_failed") yield "data: " + json.dumps(error.body) + "\n\n" yield "data: [DONE]\n\n" async def _gen_wrapper(): async for chunk in events(): yield chunk # Plain StreamingResponse (no _Worker job to release). return StreamingResponse(_gen_wrapper(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }) async def generate(body: _CompletionRequest | _ChatRequest, request: Request): engine = model_map.get(body.model) if engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") # SDK engines expose config; lightweight engine implementations need not. vocab_size = getattr(engine, "config", {}).get("vocab_size") if vocab_size is not None and body.top_k > vocab_size: raise _APIError(400, "top_k exceeds the model vocabulary size.", "invalid_parameter", "top_k") known_length = len(body.prompt) if isinstance(body, _CompletionRequest) and isinstance(body.prompt, list) else 0 if known_length > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") if body.spec_k >= 2 and body.temperature != 0: raise _APIError(400, "spec_k requires temperature=0 (greedy speculation).", "invalid_parameter", "spec_k") if not 0 <= body.seed < 2**64: raise _APIError(400, "seed must be an unsigned 64-bit integer.", "invalid_parameter", "seed") # Stream A2: fail-fast 400s before queue admission (never faked). _validate_a2_request(body, engine) # Phase2: pool/batch engines bypass the shared FIFO worker and use the # batch machinery with B=1 (single weight stream for BatchEngine / # native path, N streams for EnginePool). Plain engines keep the exact # FIFO _Worker path below (bitwise identical). if _is_pool_or_batch_engine(engine): if body.spec_k != 0: raise _APIError(400, "spec_k is not supported for pool/batch engines.", "invalid_parameter", "spec_k") worker = app.state.worker if worker is None or not worker.ready: raise _APIError(503, "Server is not ready.", "server_unavailable") with worker.lock: if len(worker.jobs) >= worker.capacity: raise _APIError(429, "Generation queue is full.", "queue_full") return await _generate_single_via_batch(engine, body, request) worker = app.state.worker if worker is None: raise _APIError(503, "Server is not ready.", "server_unavailable") job = worker.submit(body, engine) handed_off = False chat = isinstance(body, _ChatRequest) base = { "id": ("chatcmpl-" if chat else "cmpl-") + uuid.uuid4().hex, "object": "chat.completion.chunk" if chat else "text_completion", "created": int(time.time()), "model": engine.model_id, } try: first = await worker.receive(job, request) if isinstance(first, _APIError): raise first if not body.stream: parts = [first.text] ids = list(first.token_ids) tool_calls = list(first.tool_calls) chunk = first while chunk.finish_reason is None: chunk = await worker.receive(job, request) if isinstance(chunk, _APIError): raise chunk parts.append(chunk.text) ids.extend(chunk.token_ids) tool_calls.extend(chunk.tool_calls) text = "".join(parts) logprobs_payload: Any = None if _wants_logprobs(body): logprobs_payload = await _rescore_logprobs(engine, body, ids, text) choice = {"index": 0, "finish_reason": chunk.finish_reason, "logprobs": logprobs_payload} if chat: assert isinstance(body, _ChatRequest) message: dict[str, Any] = {"role": "assistant", "content": text} if tool_calls: message["tool_calls"] = tool_calls choice["message"] = message base["object"] = "chat.completion" else: choice["text"] = text return {**base, "choices": [choice], "usage": _usage(chunk)} include_usage = body.stream_options is not None and body.stream_options.include_usage def event(choices: list[dict], usage: dict | None = None, completion_tokens: int | None = None) -> str: payload = {**base, "choices": choices} if include_usage: payload["usage"] = usage if completion_tokens is not None: # Nonstandard live progress field; OpenAI clients ignore it. payload["cism_completion_tokens"] = completion_tokens return "data: " + json.dumps(payload, ensure_ascii=True, allow_nan=False) + "\n\n" async def events(): try: if chat: yield event([{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], None, first.completion_tokens) chunk = first while True: if chunk.text: content = {"delta": {"content": chunk.text}} if chat else {"text": chunk.text, "logprobs": None} yield event([{"index": 0, **content, "finish_reason": None}], None, chunk.completion_tokens) if chunk.finish_reason is not None: content = {"delta": {}} if chat else {"text": "", "logprobs": None} yield event([{"index": 0, **content, "finish_reason": chunk.finish_reason}], None, chunk.completion_tokens) if include_usage: yield event([], _usage(chunk)) break chunk = await worker.receive(job, request, streaming=True) if isinstance(chunk, _APIError): raise chunk yield "data: [DONE]\n\n" except _APIError as error: yield "data: " + json.dumps(error.body) + "\n\n" yield "data: [DONE]\n\n" except Exception: error = _APIError(503, "Generation failed.", "generation_failed") yield "data: " + json.dumps(error.body) + "\n\n" yield "data: [DONE]\n\n" response = _JobResponse(events(), worker, job) handed_off = True return response finally: if not handed_off: worker.release(job) @router.post("/tokenize") async def tokenize(body: _TokenizeRequest): """Prompt accounting: token count, context fit, KV footprint.""" engine = model_map.get(body.model) if engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") try: tokens = engine.encode(body.text) except (ValueError, TypeError) as exc: raise _APIError(400, f"Invalid prompt: {exc}.", "invalid_prompt") context = getattr(engine, "context_length", None) fits = len(tokens) <= context if isinstance(context, int) else None kv = _kv_cache_info(engine) if kv is not None: kv = {**kv, "bytes_at_prompt": kv["bytes_per_token"] * len(tokens)} return {"model": body.model, "tokens": tokens, "count": len(tokens), "context_length": context, "fits_context": fits, "kv_cache": kv} @router.post("/nll") async def score(body: _NllRequest): """Standalone perplexity scoring over pasted text.""" engine = model_map.get(body.model) if engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") try: result = await asyncio.to_thread(engine.nll, body.text, body.window) except ValueError as exc: raise _APIError(400, f"Scoring failed: {exc}.", "scoring_failed") return {"model": body.model, **result} @router.post("/models/runtime") async def models_runtime(body: _RuntimeRequest): """Page locking and keep-warm control for a loaded model.""" engine = model_map.get(body.model) if engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") status: dict[str, Any] = {} try: if body.lock_pages is True: if not hasattr(engine, "lock_pages"): raise _APIError(400, "This engine does not support page locking.", "unsupported", "lock_pages") status["page_lock_bytes"] = engine.lock_pages() elif body.lock_pages is False: if not hasattr(engine, "unlock_pages"): raise _APIError(400, "This engine does not support page locking.", "unsupported", "lock_pages") status["page_unlock_bytes"] = engine.unlock_pages() if body.keep_warm is not None: if not hasattr(engine, "start_keep_warm"): raise _APIError(400, "This engine does not support keep-warm.", "unsupported", "keep_warm") if body.keep_warm == 0: engine.stop_keep_warm() status["keep_warm"] = "stopped" else: engine.start_keep_warm(body.keep_warm) status["keep_warm"] = body.keep_warm except _APIError: raise except ValueError as exc: raise _APIError(400, f"Invalid runtime option: {exc}.", "invalid_parameter") except RuntimeError as exc: raise _APIError(409, f"Runtime change refused: {exc}.", "runtime_busy") status["keep_warm_active"] = bool(getattr(engine, "keep_warm_active", False)) return {"model": body.model, **status} @router.post("/batch") async def batch(body: _BatchRequest, request: Request): """Phase2 batched completions (order-preserving, non-streaming). Request ``{"model", "prompts": [...], ...sampling}`` returns ``{"object": "batch.completion", "choices": [{index, text, finish_reason}], "usage": {...}}`` with ``choices[i]`` for ``prompts[i]``. Same-model batches run as one native ``BatchSession`` weight stream when available (``BatchEngine`` or dense ``Engine``); pool engines fan out over ``EnginePool`` slots; concurrent batches for unlike models overlap (separate locks), same-engine batches serialize. Large batches are split FIFO waiting->running by ``max_batch_tokens`` plus ``prefill_chunk_tokens`` via ``ContinuousSchedulerV0`` (block16, per-seq EOS/stop, BlockManager eviction, recompute-requeue: evicted seqs restart from scratch). Limits: no preempt-swap yet, no cross-seq KV sharing. ``spec_k`` must be 0 (no speculation in batch). """ engine = model_map.get(body.model) if engine is None: raise _APIError(404, "Requested model is not available.", "model_not_found", "model") vocab_size = getattr(engine, "config", {}).get("vocab_size") if vocab_size is not None and body.top_k > vocab_size: raise _APIError(400, "top_k exceeds the model vocabulary size.", "invalid_parameter", "top_k") if body.spec_k != 0: raise _APIError(400, "spec_k is not supported for batch requests.", "invalid_parameter", "spec_k") if not 0 <= body.seed < 2**64: raise _APIError(400, "seed must be an unsigned 64-bit integer.", "invalid_parameter", "seed") # Early context check for token-list prompts with known lengths. for prompt in body.prompts: if isinstance(prompt, list) and len(prompt) > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") worker = app.state.worker if worker is None or not worker.ready: raise _APIError(503, "Server is not ready.", "server_unavailable") with worker.lock: queued = len(worker.jobs) # Coarse admission: a batch counts as one job; oversized batches and a # full FIFO both return 429 with Retry-After (same as single-stream). if len(body.prompts) > worker.capacity: raise _APIError(429, "Batch is larger than the generation queue.", "queue_full") if queued >= worker.capacity: raise _APIError(429, "Generation queue is full.", "queue_full") if await request.is_disconnected(): raise _APIError(503, "Request disconnected.", "request_cancelled") stop = [body.stop] if isinstance(body.stop, str) else (list(body.stop) if body.stop else None) token_limit = body.token_limit def _run_all(): from .paging import ContinuousSchedulerV0 # Measure prompt lens via engine.encode (validates + BOS fallback). try: encoded_lens: list[int] = [] for p in body.prompts: encoded_lens.append(len(engine.encode(p))) except (ValueError, TypeError): raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") for length in encoded_lens: if length > engine.context_length: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") # FIFO scheduler: waiting->running budget + prefill chunk + # block16 eviction with recompute-requeue, trivial prefix-reuse # accounting (no KV sharing, no preempt-swap in v0). sched = ContinuousSchedulerV0( num_blocks=scheduler_blocks, block_size=16, max_batch_tokens=max_batch_tokens, prefill_chunk_tokens=prefill_chunk_tokens) for idx, length in enumerate(encoded_lens): sched.add(idx, length, token_limit, stop=stop) chunks = sched.plan_chunks() num_chunks = len(chunks) ordered: list | None = [None] * len(body.prompts) prompt_list = list(body.prompts) for chunk in chunks: indices = [entry["req_id"] for entry in chunk] sub_prompts = [prompt_list[i] for i in indices] # Seed offset preserves per-seq determinism across chunks for # the single-weight-stream path (seq i uses seed+i globally). base_seed = (body.seed + min(indices)) % (2**64) if indices else body.seed if _is_pool_or_batch_engine(engine): # Pool uses one seed for all seqs (existing semantics); # keep the request seed for every chunk. chunk_seed = body.seed else: chunk_seed = base_seed try: rows = _run_batch_chunk_sync( engine, sub_prompts, token_limit, body.temperature, body.top_p, body.top_k, chunk_seed, stop) except (ValueError, TypeError) as exc: msg = str(exc) if "context" in msg.lower() or "Prompt length" in msg: raise _APIError(400, "Prompt exceeds the model context length.", "context_length_exceeded") from None raise _APIError(400, "Invalid prompt or chat messages.", "invalid_prompt") from None if len(rows) != len(indices): raise _APIError(503, "Generation did not complete.", "generation_failed") for global_idx, row in zip(indices, rows): text, _ids, finish, _pn, _gn = row if finish is None: raise _APIError(503, "Generation did not complete.", "generation_failed") assert ordered is not None ordered[global_idx] = (text, _ids, finish, _pn, _gn) # Free the planner's per-chunk blocks (accounting only; the # native sessions already released their state). for entry in chunk: try: sched.free(entry["req_id"]) except Exception: pass assert ordered is not None and all(r is not None for r in ordered) return ordered, num_chunks try: ordered, num_chunks = await asyncio.to_thread(_run_all) except _APIError: raise except Exception: raise _APIError(503, "Generation failed.", "generation_failed") from None choices = [] prompt_tokens_total = 0 completion_tokens_total = 0 for index, row in enumerate(ordered): assert row is not None text, _ids, finish, prompt_n, gen_n = row prompt_tokens_total += prompt_n completion_tokens_total += gen_n choices.append({"index": index, "text": text, "finish_reason": finish, "logprobs": None}) return {"id": "batch-" + uuid.uuid4().hex, "object": "batch.completion", "created": int(time.time()), "model": engine.model_id, "choices": choices, "usage": {"prompt_tokens": prompt_tokens_total, "completion_tokens": completion_tokens_total, "total_tokens": prompt_tokens_total + completion_tokens_total}, "scheduler": {"max_batch_tokens": max_batch_tokens, "max_batched_tokens": max_batch_tokens, "prefill_chunk_tokens": prefill_chunk_tokens, "block_size": 16, "num_blocks": scheduler_blocks, "chunks": num_chunks}} @router.post("/completions") async def completions(body: _CompletionRequest, request: Request): return await generate(body, request) @router.post("/chat/completions") async def chat_completions(body: _ChatRequest, request: Request): return await generate(body, request) app.include_router(router) return app def _kv_cache_info(engine: Any) -> dict | None: """Per-token and full-context KV cache bytes, or None without config.""" config = getattr(engine, "config", None) context = getattr(engine, "context_length", None) if not isinstance(config, dict) or not isinstance(context, int): return None try: per_token = (int(config["num_hidden_layers"]) * int(config["num_key_value_heads"]) * int(config["head_dim"]) * 2 * 4) except (KeyError, TypeError, ValueError): return None if per_token <= 0: return None return {"bytes_per_token": per_token, "bytes_at_full_context": per_token * context} _WEBUI_HTML = """ CISM tester CISM tester

CISM tester

native CPU inference console

health: unknown

Models fetch one, pin it down, compare it

Active model

Request say it plain, tune what matters

Measure numbers, not adjectives

Precision shootout

Perplexity scoring

scores the prompt above on the selected model

Cache residency

uses the selected model

Output


  

  

  

    
  









"""