From 4a61cd630ae9f5097d502115c98003dbde5d984c Mon Sep 17 00:00:00 2001 From: Hendrik Date: Thu, 17 Sep 2026 17:54:40 +0200 Subject: [PATCH 1/7] serve: model registry with hot-swap and per-request model logging Add a swappable model registry to so the model behind a stable route name can change without reconfiguring clients. - POST /v1/models/load swaps the loaded model; GET /v1/models lists the registry with loaded flags; generation requests are validated: unknown model is a 404, registered-but-not-loaded is a 409 pointing at the load endpoint instead of silently serving whatever is resident. - CLI: --models PATH[=ID] registers swappable containers (repeatable), --keep-previous keeps a replaced model resident instead of unloading it. - Swap opens the new container before closing the old one, so a failed load leaves the previous model serving; swap happens under the current engine's lock so a generation in flight finishes first. - Request log lines now name the model: the one that served a success, the one refused on a 409/404, the one a load switched to. New --no-log-requests silences request logging. - Tests over real sockets cover registry listing, swap, keep-previous, validation statuses, and the model names in the log; docs updated. --- docs/SERVE.md | 72 +++++++- serve/__main__.py | 70 +++++++- serve/api.py | 9 +- serve/server.py | 271 ++++++++++++++++++++++++++++--- tests/serve/fake_engine.py | 5 + tests/serve/test_server.py | 325 ++++++++++++++++++++++++++++++++++++- 6 files changed, 725 insertions(+), 27 deletions(-) diff --git a/docs/SERVE.md b/docs/SERVE.md index e5f478195..dcf850eaf 100644 --- a/docs/SERVE.md +++ b/docs/SERVE.md @@ -268,7 +268,8 @@ transcribed in this repo. That is what is left of | endpoint | notes | |---|---| | `GET /health` | liveness; never requires the API key | -| `GET /v1/models`, `GET /v1/models/{id}` | reports the container's real shape under a `waste` key | +| `GET /v1/models`, `GET /v1/models/{id}` | the registry: the loaded model (its real shape under a `waste` key), plus any registered-but-not-loaded containers with `"loaded": false` | +| `POST /v1/models/load` | swap models; see "Swapping models" below | | `POST /v1/chat/completions` | streaming and not, tools, images | | `POST /v1/completions` | raw continuation, no chat template | @@ -318,6 +319,55 @@ one lock, and requests queue. On a model streaming experts off an SSD at a few tokens a second, the wait for the lock is small next to the wait for the answer. +### Swapping models + +By default the process serves the one container it was started with, and a +request naming any other model is a 404 — before the registry existed any +name was silently served by the loaded model, which made `model` a +decorative string. + +`--models PATH[=ID]` (repeatable) registers additional containers a client +may switch to: + + python3 -m serve ~/models/k3.waste \ + --models ~/models/glm53.waste --models ~/models/deepseek41.waste=ds41 + + curl localhost:8000/v1/models # the registry, with + # "loaded" per entry + curl localhost:8000/v1/models/load \ + -d '{"model":"glm53"}' # swap, then answer + +What a swap does: + +- It waits for the engine lock, so a generation in flight finishes before + the model moves under it. +- It opens the new container, re-derives everything the container decides — + reply format, markers, stop tokens, the thinking default, `/v1/models`'s + shape — and only then **unloads the previous model** (`waste_close`). + One model resident at a time is the default, and the flag to change it + is `--keep-previous`. +- The new container is opened before the old one is unloaded, so a swap + that fails — a truncated container, an `--exclusive-open` conflict — + leaves the server serving what it was serving and answers 500 with the + engine's own reason. The cost of that guarantee is a moment where both + containers are resident; size `--budget` so that moment fits. +- Generation always serves the current model. A request naming a + registered-but-not-loaded model is a 409, telling the client to + `POST /v1/models/load` first, rather than an unnoticed multi-gigabyte + swap in the middle of a conversation. A request with no model, or + naming the current one, is untouched — which is every client that does + not know about the registry. +- A swap discards the previous model's KV state. That is inherent: + `waste_close` frees the context. An agent harness should treat a swap + as rare and expensive, not as a per-turn choice. + +`--keep-previous` keeps the replaced model resident instead of unloading +it. Switching back to it is then a slot move rather than a reopen — its +`waste_ctx` and the state it holds are still there. The RAM two resident +contexts need is the sum of their budgets; on the machines this engine +targets that is usually the difference between working and paging, which +is why unloading is the default. + Streaming is written straight from the token callback, on the thread holding the lock. A client hanging up propagates back as a return value the engine understands — the callback says stop, `waste_generate` unwinds, the @@ -380,6 +430,20 @@ nothing while the model reasons — which, on a model whose reasoning can be most of the reply, looks like a server that has stopped. `--no-thinking` makes the default answer-only, and a request can still ask for reasoning. +## Request log + +With request logging on (the default; `--no-log-requests` turns it off), +each response line names the model it concerns: + +``` +127.0.0.1 - "POST /v1/chat/completions HTTP/1.1" 200 - [model=glm53] +127.0.0.1 - "POST /v1/chat/completions HTTP/1.1" 409 - [model=tiny] +``` + +A success line names the model that served it, a 409 or 404 names the +model that was refused, and a `/v1/models/load` line names the model it +switched to. GET lines carry no model annotation. + ## Security - `--host` defaults to `127.0.0.1`. Binding anywhere else without @@ -444,5 +508,11 @@ python3 -m serve MODEL [options] --max-tokens N default cap when a request does not set one (4096) --no-thinking answer without the think channel unless asked --allow-local-images + --models PATH[=ID] additional containers a client may switch to with + POST /v1/models/load (repeatable; switching unloads + the model it replaces) + --keep-previous keep the replaced model resident instead of unloading + it; RAM needed is then the sum of both budgets --plan print the memory plan and exit + --no-log-requests silence the per-request log lines ``` diff --git a/serve/__main__.py b/serve/__main__.py index 061355db5..cf89b623d 100644 --- a/serve/__main__.py +++ b/serve/__main__.py @@ -26,11 +26,35 @@ WASTE_E_ARG, WASTE_E_BUSY, WASTE_E_UNSUPPORTED, Engine, EngineError, build_info, physical_ram, plan_memory) -from .server import serve # noqa: E402 +from .server import ModelLoadError, serve # noqa: E402 POLICIES = {"lfru": CACHE_LFRU, "lru": CACHE_LRU} +def parse_registry(specs: list[str]) -> dict[str, str]: + """--models entries as {id: path}. + + `PATH[=ID]`, ID defaulting to the file name without .waste — the + same default --model-id uses. Duplicate ids are an error rather than + a quiet overwrite: two containers behind one name means a client's + "load this model" sometimes loads a different one than it just + listed, and that failure announces itself only under load. + """ + registry: dict[str, str] = {} + for spec in specs: + path, _, custom = spec.partition("=") + p = Path(path).expanduser() + if not p.exists(): + raise SystemExit(f"--models: no such container: {p}") + mid = custom or p.name.removesuffix(".waste") + if mid in registry: + raise SystemExit( + f"--models: duplicate model id: {mid} (both " + f"{registry[mid]} and {p})") + registry[mid] = str(p) + return registry + + def human(n: float) -> str: for unit in ("B", "KB", "MB", "GB", "TB"): if n < 1024 or unit == "TB": @@ -82,6 +106,11 @@ def main(argv=None) -> int: curl localhost:8000/v1/chat/completions -H 'Content-Type: application/json' \\ -d '{"model":"waste","messages":[{"role":"user","content":"hi"}]}' + + python3 -m serve ~/models/k3.waste --models ~/models/glm53.waste \\ + --models ~/models/deepseek41.waste=ds41 + # POST /v1/models/load {"model":"glm53"} swaps to it, unloading k3; + # add --keep-previous to hold both resident instead """) ap.add_argument("model", help="path to the .waste container") ap.add_argument("--host", default="127.0.0.1", @@ -140,8 +169,28 @@ def main(argv=None) -> int: help="let requests name images by filesystem path. Off by " "default: it lets any client read files the server " "can reach") + s.add_argument("--models", action="append", default=[], metavar="PATH[=ID]", + help="an additional container a client may switch to " + "with POST /v1/models/load (repeatable; id defaults " + "to the file name without .waste). Switching " + "unloads the model it replaces unless " + "--keep-previous") + s.add_argument("--keep-previous", action="store_true", + help="keep a model resident when another is loaded. " + "Off by default, and deliberately: the RAM two " + "contexts need together is the sum of their " + "budgets, and on the machines this engine targets " + "that is the difference between working and " + "paging. Both models can then answer at once — " + "each waste_ctx takes one caller, so each has its " + "own lock") s.add_argument("--plan", action="store_true", help="print the memory plan and exit without loading") + s.add_argument("--no-log-requests", action="store_true", + help="silence the per-request log lines. On by default: " + "each line names the model that served or was " + "refused — on a server that swaps models, the log " + "is how you find out which one answered") args = ap.parse_args(argv) @@ -217,7 +266,22 @@ def main(argv=None) -> int: api_key=args.api_key, default_max_tokens=args.max_tokens, default_thinking=not args.no_thinking, - allow_local_images=args.allow_local_images) + allow_local_images=args.allow_local_images, + log_requests=not args.no_log_requests, + models=parse_registry(args.models), + keep_previous=args.keep_previous, + engine_kwargs={ + "ram_budget_bytes": args.budget, + "ctx_tokens": args.ctx, + "n_threads": args.threads, + "cpu_list": args.cpus, + "cache_policy": POLICIES[args.cache], + "direct_io": not args.no_direct_io, + "vision": args.vision, + "verify_records": args.verify, + "usage_path": args.usage, + "exclusive_open": args.exclusive_open, + }) except (EngineError, OSError) as e: engine.close() print(f"{e}", file=sys.stderr) @@ -272,7 +336,7 @@ def main(argv=None) -> int: finally: srv.shutdown() srv.server_close() - engine.close() + srv.close_engines() shutil.rmtree(srv.tmpdir, ignore_errors=True) return 0 diff --git a/serve/api.py b/serve/api.py index 164f9f6d6..185d856d4 100644 --- a/serve/api.py +++ b/serve/api.py @@ -558,9 +558,16 @@ def engine_extra(stats: dict, *, ms: float) -> dict: } -def model_object(model_id: str, created: int, info: Optional[dict] = None) -> dict: +def model_object(model_id: str, created: int, info: Optional[dict] = None, + loaded: Optional[bool] = None) -> dict: obj = {"id": model_id, "object": "model", "created": created, "owned_by": "waste"} + # Not an OpenAI field. A registry of swappable containers needs to say + # which of them is resident; without a marker the client would have to + # infer it from which entry carries a `waste` shape, which is worse + # than spelling it. + if loaded is not None: + obj["loaded"] = loaded if info: obj["waste"] = info return obj diff --git a/serve/server.py b/serve/server.py index 6d50be688..197e5aba0 100644 --- a/serve/server.py +++ b/serve/server.py @@ -8,8 +8,11 @@ Endpoints: GET /health liveness, plus what is loaded - GET /v1/models the one model this process holds + GET /v1/models the registry: the loaded model, plus any + registered-but-not-loaded containers GET /v1/models/{id} + POST /v1/models/load swap models (requires --models; unloads + the previous model unless --keep-previous) POST /v1/chat/completions streaming and not, tools, images POST /v1/completions raw continuation, no chat template @@ -38,7 +41,7 @@ import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Optional +from typing import Callable, Optional from . import api, dsml, glmtools, xtml from .chatfmt import ChatFormat, ChatFormatError, PlainParser @@ -53,6 +56,16 @@ MAX_BODY_BYTES = 64 * 1024 * 1024 +class ModelLoadError(EngineError): + """A swap could not open the container it was asked for. Flows out of + POST /v1/models/load as a 500 with the engine's own reason; the + previously loaded model is still current.""" + + def __init__(self, message: str): + from .engine import WASTE_E_IO + super().__init__("model load", WASTE_E_IO, message) + + class ChatServer(ThreadingHTTPServer): """Threaded HTTP, one engine, one lock.""" @@ -65,18 +78,66 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, default_thinking: bool = True, allow_local_images: bool = False, log_requests: bool = True, - tmpdir: Optional[str] = None): + tmpdir: Optional[str] = None, + models: Optional[dict] = None, + keep_previous: bool = False, + engine_kwargs: Optional[dict] = None, + engine_factory: Optional[Callable] = None): super().__init__(addr, handler) - self.engine = engine - self.model_id = model_id + self.engine_kwargs = dict(engine_kwargs or {}) + self.keep_previous = keep_previous + self.engine_factory = engine_factory or self._default_engine_factory + self._slot_lock = threading.RLock() + # model_id -> the Engine holding it. With keep_previous there can be + # more than one; without it, exactly one — the previous entry is + # closed and dropped as the new one becomes current. + self.engines = {model_id: engine} + # model_id -> container path. The loaded model is registered from + # its own path; --models adds the rest of the swappable set. An id + # already taken by the loaded model keeps the loaded path: the + # registry is a way to name what can be swapped to, not a way to + # point the loaded model somewhere else. + self.registry = {model_id: engine.model_path or model_id} + for mid, path in (models or {}).items(): + self.registry.setdefault(mid, path) self.api_key = api_key self.default_max_tokens = default_max_tokens - self.default_thinking = default_thinking self.allow_local_images = allow_local_images self.log_requests = log_requests self.started = api.now() self._tmp = tmpdir or tempfile.mkdtemp(prefix="waste-serve-") self.tmpdir = self._tmp + + self._start_thinking = default_thinking + self._detect(engine, model_id) + + def _default_engine_factory(self, path: str) -> Engine: + """How a swap opens a container. Overridable — tests hand in a + factory that builds the scripted engine, and a host embedding the + server may want its own construction arguments.""" + return Engine(path, **self.engine_kwargs) + + # ---- what the current model makes true ------------------------------- + # + # The block that follows used to run once, in __init__, and every + # handler read its verdicts as constants for the life of the process. + # With a swappable registry they are per-model: a container without + # XTML markers must not inherit the previous container's chat format, + # stop tokens, or thinking default. So the same block is a method, run + # again on every load, and it keeps the old comment because every word + # of it still holds. + + def _detect(self, engine: Engine, model_id: str) -> None: + """Bind `engine` as the current model and re-derive everything the + handlers read from the container rather than from the request: + model_info, the reply format, stop tokens, the thinking default. + + The engine lock is NOT taken here — the caller holds it, or (at + construction) no request can have arrived yet. + """ + self.engine = engine + self.model_id = model_id + self.default_thinking = self._start_thinking try: self.model_info = engine.model_info() except EngineError: @@ -153,6 +214,101 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, # be asked to answer without it either. self.default_thinking = fmt.think is not None + # ---- the model registry ---------------------------------------------- + + def check_model_request(self, body: dict) -> None: + """What a generation request may name as its model. + + Absent, empty, or equal to the loaded model's id: fine — that is + what every client that does not know about this registry sends, + and it must keep working. A name the registry does not know is a + 404: before the registry existed any name was silently served by + the loaded model, which made `model` a decorative string. A name + the registry knows but that is not resident is a 409, not a + surprise multi-gigabyte swap in the middle of a conversation — + the client asks for that explicitly with POST /v1/models/load. + """ + mid = body.get("model") + if not isinstance(mid, str) or not mid or mid == self.model_id: + return + if mid not in self.registry: + raise api.APIError(f"no such model: {mid}", status=404, + type="not_found_error", param="model") + raise api.APIError( + f"model {mid} is registered but not loaded; POST /v1/models/load " + f"to switch to it", status=409, type="model_not_loaded", + param="model") + + def load_model(self, model_id: str) -> Optional[str]: + """Make `model_id` the current model, and return the id of the + model that was current before (None when it already was). + + Unloading: without keep_previous the previous engine is closed as + part of the swap — one model resident at a time, and the freed + RAM goes to the new container's expert cache. With keep_previous + the previous engine stays open in `engines`, so switching back to + it later is a slot move rather than a reopen of a multi-gigabyte + container. Generation always serves the current slot; other + resident models answer after the next load names them, not + before. + + The new container is opened *before* the old one is unloaded, + which is the opposite of the order a tight-RAM machine would + prefer, and the reason is rollback: a swap whose open fails — a + truncated container, an --exclusive-open conflict — must leave + the server serving the model it was serving, and the only way to + guarantee that is not to have closed it yet. A failed swap costs + a load's worth of RAM for a moment; a swap that leaves no model + loaded costs the whole server. When two containers genuinely will + not fit together, size --budget so each open plans against what + the engine can actually get. + + The old engine's lock is held across the whole swap, so a + generation in flight finishes before the slot moves under it, and + no request that took the old lock can find its engine closed. + """ + with self._slot_lock: + current = self._current() + if model_id == current[0]: + return None + path = self.registry.get(model_id) + if path is None: + raise api.APIError(f"no such model: {model_id}", status=404, + type="not_found_error", param="model") + previous_id, previous_engine = current + # A model kept resident by keep_previous does not need an + # open at all — its waste_ctx still holds the state it had. + # Moving the slot to it costs a format re-detect, not a load. + resident = self.engines.get(model_id) + with previous_engine.lock: + if resident is not None: + self._detect(resident, model_id) + return previous_id + try: + engine = self.engine_factory(path) + except EngineError as e: + raise ModelLoadError( + f"could not load {model_id}: {e}") from e + self._detect(engine, model_id) + self.engines[model_id] = engine + if self.keep_previous: + return previous_id + self.engines.pop(previous_id) + previous_engine.close() + return previous_id + + def _current(self) -> tuple: + with self._slot_lock: + model_id = self.model_id + return model_id, self.engines.get(model_id, self.engine) + + def close_engines(self) -> None: + """Every engine this server still holds. The shutdown path; with + keep_previous there may be several.""" + for engine in list(self.engines.values()): + engine.close() + self.engines.clear() + def new_parser(self, thinking: bool, tools=None): """The reply reader for whichever format this container speaks. @@ -200,7 +356,21 @@ class Handler(BaseHTTPRequestHandler): def log_message(self, fmt, *args): if getattr(self.server, "log_requests", True): - sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args)) + line = fmt % args + model = getattr(self, "_log_model", None) + if model: + line += " [model=%s]" % model + sys.stderr.write("%s - %s\n" % (self.address_string(), line)) + + def _log_model_from(self, body): + """The model name this request's log line carries: what the request + asked for, or the loaded model when it named none. A success line + then shows the model that served it, and a 409 or 404 shows the + model that was refused — `who was asked` is on the log line even + when `who answered` is nobody.""" + mid = body.get("model") + self._log_model = (mid if isinstance(mid, str) and mid + else self.server.model_id) def _send_json(self, status: int, payload: dict, *, headers=None) -> None: body = json.dumps(payload, ensure_ascii=False).encode() @@ -265,6 +435,7 @@ def _authorized(self) -> bool: # ---- routing -------------------------------------------------------- def do_GET(self): + self._log_model = None # no keep-alive request inherits the last one's model try: path = self.path.split("?", 1)[0].rstrip("/") or "/" if path == "/health": @@ -284,6 +455,7 @@ def do_GET(self): pass def do_POST(self): + self._log_model = None # no keep-alive request inherits the last one's model try: path = self.path.split("?", 1)[0].rstrip("/") or "/" if not self._authorized(): @@ -294,6 +466,8 @@ def do_POST(self): return self._chat() if path == "/v1/completions": return self._completions() + if path == "/v1/models/load": + return self._load_model() self.close_connection = True raise api.APIError(f"no route for POST {path}", status=404, type="not_found_error") @@ -317,20 +491,53 @@ def _health(self): }) def _models(self): - self._send_json(200, { - "object": "list", - "data": [api.model_object(self.server.model_id, - self.server.started, - self.server.model_info)], - }) + srv = self.server + # The whole registry, current model first, so a client scanning the + # list sees what is resident before what is only available. A + # registered-but-not-loaded entry carries no `waste` shape: its + # per-container facts are unknown until it is opened. + data = [api.model_object(srv.model_id, srv.started, + srv.model_info, loaded=True)] + data += [api.model_object(mid, srv.started, None, loaded=False) + for mid in sorted(srv.registry) + if mid != srv.model_id] + self._send_json(200, {"object": "list", "data": data}) def _model(self, model_id: str): - if model_id != self.server.model_id: + srv = self.server + if model_id == srv.model_id: + self._send_json(200, api.model_object(model_id, srv.started, + srv.model_info, loaded=True)) + return + if model_id not in srv.registry: raise api.APIError(f"no such model: {model_id}", status=404, type="not_found_error", param="model") - self._send_json(200, api.model_object(self.server.model_id, - self.server.started, - self.server.model_info)) + self._send_json(200, api.model_object(model_id, srv.started, + None, loaded=False)) + + def _load_model(self): + """POST /v1/models/load — swap the model this server serves. + + Body: {"model": ""} where id is a name from GET /v1/models. + The swap happens under the current engine's lock, so a generation + in flight finishes first; the reply says which model went out and + which came in. Errors: 404 unknown id, 500 (ModelLoadError) the + container would not open — the previous model is still served. + """ + body = self._read_body() + srv = self.server + mid = body.get("model") + self._log_model_from(body) # 404/500 lines name the model + if not isinstance(mid, str) or not mid: + raise api.APIError("'model' must be a non-empty string", + param="model") + previous = srv.load_model(mid) # raises 404 / ModelLoadError + self._send_json(200, { + "object": "model.load", + "loaded": mid, + "previous": previous, + "models": [m for m in srv.registry if m in srv.engines], + }) # ---- chat ----------------------------------------------------------- @@ -338,6 +545,13 @@ def _chat(self): body = self._read_body() srv = self.server engine = srv.engine + self._log_model_from(body) + + # Before anything else, and before the engine lock: a request that + # names a model this process does not serve should hear it from a + # 404, not from the model's own reply. Absent or matching the + # loaded model's id passes untouched. + srv.check_model_request(body) # Before anything else, and before the engine lock: this container # has no chat format we can render, and no request can change that. @@ -583,6 +797,9 @@ def _completions(self): """ body = self._read_body() srv = self.server + self._log_model_from(body) # before check_model_request: a 404 or + # 409 line should still name the model that was refused + srv.check_model_request(body) prompt_text = body.get("prompt") if isinstance(prompt_text, list): if len(prompt_text) != 1 or not isinstance(prompt_text[0], str): @@ -641,8 +858,19 @@ def serve(engine: Engine, *, host: str = "127.0.0.1", port: int = 8000, model_id: str = "waste", api_key: Optional[str] = None, default_max_tokens: int = 4096, default_thinking: bool = True, allow_local_images: bool = False, log_requests: bool = True, - ready: Optional[threading.Event] = None) -> ChatServer: - """Build the server. The caller decides whether to serve_forever.""" + ready: Optional[threading.Event] = None, + models: Optional[dict] = None, keep_previous: bool = False, + engine_kwargs: Optional[dict] = None, + engine_factory: Optional[Callable] = None) -> ChatServer: + """Build the server. The caller decides whether to serve_forever. + + models names the rest of the swappable registry (id -> container + path); keep_previous decides whether a swap unloads the model it + replaces; engine_kwargs is how the default factory opens a container + on a swap — the same arguments the startup engine was opened with — + and engine_factory replaces that factory wholesale, for a host that + builds engines its own way (tests do exactly that). + """ # IPv6-capable when the host asks for it, without forcing it: binding # :: on a host with IPv6 disabled fails outright. if ":" in host: @@ -651,7 +879,10 @@ def serve(engine: Engine, *, host: str = "127.0.0.1", port: int = 8000, api_key=api_key, default_max_tokens=default_max_tokens, default_thinking=default_thinking, allow_local_images=allow_local_images, - log_requests=log_requests) + log_requests=log_requests, models=models, + keep_previous=keep_previous, + engine_kwargs=engine_kwargs, + engine_factory=engine_factory) if ready is not None: ready.set() return srv diff --git a/tests/serve/fake_engine.py b/tests/serve/fake_engine.py index c4d2bd77c..e52916cef 100644 --- a/tests/serve/fake_engine.py +++ b/tests/serve/fake_engine.py @@ -86,6 +86,11 @@ class FakeEngine: fail_with: Optional[Exception] = None # Sleep this long before each token, to test client disconnects. delay: float = 0.0 + # What a fresh container answers without being asked. Serve-level code + # only reads this to report it (GET /v1/models) and to compare it + # against a request's reasoning_effort, so a boolean stands in for the + # real engine's reasoning_effort floor. + default_thinking: bool = True prompts: list[list[int]] = field(default_factory=list) calls: list[dict] = field(default_factory=list) diff --git a/tests/serve/test_server.py b/tests/serve/test_server.py index cd1441481..975a5f5bd 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -17,7 +17,9 @@ """ import json +import contextlib import http.client +import io import shutil import socket import sys @@ -35,6 +37,7 @@ from serve.engine import EngineError, WASTE_E_IO # noqa: E402 from serve.server import serve # noqa: E402 from tests.serve.fake_engine import (FakeEngine, LINEAR_MARKERS, # noqa: E402 + MARKERS, # noqa: E402 reply_plain, reply_tool_call) @@ -43,11 +46,13 @@ class ServerTestCase(unittest.TestCase): engine_kwargs: dict = {} server_kwargs: dict = {} + log_requests = False def setUp(self): self.engine = FakeEngine(**self.engine_kwargs) self.server = serve(self.engine, host="127.0.0.1", port=0, - model_id="test-model", log_requests=False, + model_id="test-model", + log_requests=self.log_requests, **self.server_kwargs) self.port = self.server.server_address[1] self.thread = threading.Thread(target=self.server.serve_forever, @@ -404,10 +409,19 @@ def test_client_disconnect_stops_generation(self): class TestValidation(ServerTestCase): def test_missing_messages(self): - status, body = self.post("/v1/chat/completions", {"model": "m"}) + status, body = self.post("/v1/chat/completions", {}) self.assertEqual(status, 400) self.assertEqual(body["error"]["param"], "messages") + def test_unknown_model_outranks_missing_messages(self): + # The model is validated before the request's shape, the way the + # OpenAI API does: a client pointed at a model this server does + # not serve should hear "no such model", not a complaint about + # messages it would have sent correctly to the right server. + status, body = self.post("/v1/chat/completions", {"model": "m"}) + self.assertEqual(status, 404) + self.assertEqual(body["error"]["param"], "model") + def test_empty_messages(self): status, body = self.post("/v1/chat/completions", {"messages": []}) self.assertEqual(status, 400) @@ -819,6 +833,313 @@ def test_rejected_post_closes_connection_with_unread_body(self): conn.close() +class TestRequestLogs(ServerTestCase): + """The request log names the model, even when it refused one. + + log_message runs on the server thread and writes sys.stderr at write + time, so redirecting it around the whole test captures the lines; each + is written before the response is flushed, so by the time the client + has the body the line is already in the buffer. + """ + + log_requests = True + + def make_engine(self, path: str) -> FakeEngine: + return FakeEngine(model_path=path, markers=MARKERS) + + def setUp(self): + self._captured = io.StringIO() + self._ctx = contextlib.redirect_stderr(self._captured) + self._ctx.__enter__() + try: + self.engine_kwargs = {"model_path": "/fake/start.waste"} + self.server_kwargs = { + "models": {"swap-a": "/fake/a.waste"}, + "engine_factory": self.make_engine, + } + ServerTestCase.setUp(self) + except BaseException: + self._ctx.__exit__(None, None, None) + raise + + def tearDown(self): + try: + ServerTestCase.tearDown(self) + finally: + self._ctx.__exit__(None, None, None) + + def logs(self) -> str: + self._captured.flush() + return self._captured.getvalue() + + def test_chat_log_names_the_serving_model(self): + status, _ = self.post("/v1/chat/completions", + {"messages": [{"role": "user", "content": "x"}]}) + self.assertEqual(status, 200) + self.assertIn('"POST /v1/chat/completions HTTP/1.1" 200 -' + " [model=test-model]", self.logs()) + + def test_chat_log_names_a_named_model(self): + status, _ = self.post("/v1/chat/completions", + {"model": "test-model", + "messages": [{"role": "user", "content": "x"}]}) + self.assertEqual(status, 200) + self.assertIn("[model=test-model]", self.logs()) + + def test_chat_log_names_the_model_a_409_refused(self): + status, _ = self.post("/v1/chat/completions", + {"model": "swap-a", + "messages": [{"role": "user", "content": "x"}]}) + self.assertEqual(status, 409) + self.assertIn('"POST /v1/chat/completions HTTP/1.1" 409 -' + " [model=swap-a]", self.logs()) + + def test_chat_log_names_the_model_a_404_refused(self): + status, _ = self.post("/v1/chat/completions", + {"model": "nope", + "messages": [{"role": "user", "content": "x"}]}) + self.assertEqual(status, 404) + self.assertIn("[model=nope]", self.logs()) + + def test_load_log_names_the_model_loaded(self): + status, _ = self.post("/v1/models/load", {"model": "swap-a"}) + self.assertEqual(status, 200) + self.assertIn('"POST /v1/models/load HTTP/1.1" 200 -' + " [model=swap-a]", self.logs()) + + def test_get_log_carries_no_model(self): + status, _ = self.get("/v1/models") + self.assertEqual(status, 200) + self.assertNotIn("[model=", self.logs()) + + +class TestModelSwap(ServerTestCase): + """POST /v1/models/load and the model field a generation request names. + + The registry holds scripted engines; the engine_factory builds one per + container path, which is what the real server's default factory does + with the real Engine. + """ + + keep = False + + def make_engine(self, path: str) -> FakeEngine: + engine = FakeEngine(model_path=path, markers=self.swap_markers, + **self.factory_kwargs) + self.made.append(engine) + return engine + + def setUp(self): + self.swap_markers = dict(MARKERS) + self.factory_kwargs = {} + self.made = [] + self.engine_kwargs = {"model_path": "/fake/start.waste"} + self.server_kwargs = { + "models": {"swap-a": "/fake/a.waste", "swap-b": "/fake/b.waste"}, + "keep_previous": self.keep, + "engine_factory": self.make_engine, + } + ServerTestCase.setUp(self) + + def load(self, model): + return self.post("/v1/models/load", {"model": model}) + + def test_load_switches_and_unloads_previous(self): + status, body = self.load("swap-a") + self.assertEqual(status, 200) + self.assertEqual(body["loaded"], "swap-a") + self.assertEqual(body["previous"], "test-model") + self.assertEqual(self.made[0].model_path, "/fake/a.waste") + # One model resident at a time: the startup engine is closed and + # dropped, and the swap target is what every endpoint reports. + self.assertTrue(self.engine.closed) + self.assertEqual(self.server.model_id, "swap-a") + self.assertEqual(self.server.engine, self.made[0]) + self.assertEqual(list(self.server.engines), ["swap-a"]) + + def test_load_unknown_model_404s(self): + status, body = self.load("nope") + self.assertEqual(status, 404) + self.assertEqual(body["error"]["type"], "not_found_error") + self.assertEqual(self.server.model_id, "test-model") + + def test_load_current_model_is_a_noop(self): + status, body = self.load("test-model") + self.assertEqual(status, 200) + self.assertEqual(body["previous"], None) + self.assertEqual(self.made, []) # no engine was built + self.assertEqual(self.server.model_id, "test-model") + + def test_models_lists_registry_with_loaded_flags(self): + status, body = self.get("/v1/models") + self.assertEqual(status, 200) + by_id = {m["id"]: m for m in body["data"]} + self.assertEqual(by_id["test-model"]["loaded"], True) + self.assertEqual(by_id["swap-a"]["loaded"], False) + self.assertNotIn("waste", by_id["swap-a"]) # not opened: no shape + self.assertIn("waste", by_id["test-model"]) + # The current model first, so a client scanning the list sees what + # is resident before what is only available. + self.assertEqual(body["data"][0]["id"], "test-model") + + def test_registered_model_entry_says_not_loaded(self): + status, body = self.get("/v1/models/swap-b") + self.assertEqual(status, 200) + self.assertEqual(body["id"], "swap-b") + self.assertEqual(body["loaded"], False) + + def test_generation_rejects_registered_but_not_loaded(self): + status, body = self.chat(model="swap-b") + self.assertEqual(status, 409) + self.assertEqual(body["error"]["type"], "model_not_loaded") + + def test_generation_rejects_unknown_model(self): + status, body = self.chat(model="nope") + self.assertEqual(status, 404) + self.assertEqual(body["error"]["type"], "not_found_error") + + def test_generation_accepts_loaded_model_after_swap(self): + status, _ = self.load("swap-a") + self.assertEqual(status, 200) + self.made[0].reply = reply_plain("from a") + status, body = self.chat(model="swap-a") + self.assertEqual(status, 200) + self.assertEqual(body["model"], "swap-a") + self.assertEqual(body["choices"][0]["message"]["content"], "from a") + + def test_generation_without_model_still_served_after_swap(self): + """A client that never names a model must survive a swap it did + not ask for — absent means 'whatever is loaded'. The helper's + default model name would 409 after a swap, so this omits it the + way a client that does not know about the registry does.""" + self.load("swap-a") + self.made[0].reply = reply_plain("still there") + status, body = self.chat(model=None) + self.assertEqual(status, 200) + self.assertEqual(body["model"], "swap-a") + + def test_failed_open_leaves_previous_loaded(self): + def broken(path): + raise EngineError("open", WASTE_E_IO, path) + + self.server.engine_factory = broken + status, body = self.load("swap-a") + self.assertEqual(status, 500) + self.assertEqual(body["error"]["type"], "engine_error") + self.assertEqual(self.server.model_id, "test-model") + # And it still serves: the rollback guarantee is the point. + self.engine.reply = reply_plain("unharmed") + status, body = self.chat() + self.assertEqual(status, 200) + self.assertEqual(body["choices"][0]["message"]["content"], "unharmed") + + def test_failed_open_leaves_registry_usable(self): + def broken(path): + raise EngineError("open", WASTE_E_IO, path) + + self.server.engine_factory = broken + self.load("swap-a") + self.server.engine_factory = self.make_engine + status, body = self.load("swap-b") + self.assertEqual(status, 200) + self.assertEqual(body["previous"], "test-model") + + def test_per_container_facts_are_rebuilt(self): + """A container without XTML must not inherit the previous one's + reply format — the swap re-derives markers, format, stop tokens, + and the thinking default.""" + self.swap_markers = {} # marker_ids() will raise + self.factory_kwargs = {"no_markers": True} + status, body = self.load("swap-a") + self.factory_kwargs = {} + self.assertEqual(status, 200) + self.assertIsNotNone(self.server.chat_error) + # The refusal moves with the slot: the startup model could chat, + # the loaded one cannot. + status, body = self.chat(model=None) + self.assertEqual(status, 400) + self.assertEqual(body["error"]["code"], "unsupported_chat_format") + + def test_swap_waits_out_a_generation_in_flight(self): + """A swap takes the old engine's lock: the streaming request that + started first finishes whole, and the swap's close happens after, + not under it.""" + self.engine.delay = 0.02 + self.engine.reply = reply_plain("long answer") + events = [] + errors = [] + + def streamer(): + try: + events.extend(self.stream()) + except Exception as e: # pragma: no cover + errors.append(e) + + t = threading.Thread(target=streamer) + t.start() + t.join(timeout=0.1) # started, not done + status, body = self.load("swap-a") + self.assertEqual(status, 200) + t.join(timeout=30) + self.assertEqual(errors, []) + self.assertEqual(events[-1], "[DONE]") + + +class TestModelSwapKeepPrevious(TestModelSwap): + """--keep-previous: the model a swap replaces stays resident.""" + + keep = True + + def test_load_switches_and_unloads_previous(self): + status, body = self.load("swap-a") + self.assertEqual(status, 200) + self.assertEqual(body["loaded"], "swap-a") + self.assertEqual(self.server.model_id, "swap-a") + self.assertEqual(self.server.engine, self.made[0]) + + def test_previous_stays_open(self): + status, body = self.load("swap-a") + self.assertEqual(status, 200) + self.assertFalse(self.engine.closed) + self.assertEqual(sorted(self.server.engines), + ["swap-a", "test-model"]) + + def test_can_switch_back_without_reopening(self): + """swap-a was never closed, so loading it again is a slot move, + not a new open.""" + self.load("swap-a") + self.load("test-model") + self.assertEqual(self.server.model_id, "test-model") + self.assertEqual(self.server.engine, self.engine) + self.assertEqual([e.model_path for e in self.made], + ["/fake/a.waste"]) + + def test_both_engines_survive_parallel_traffic(self): + """Two resident contexts: nothing is refused, nothing is closed + mid-answer. Generation serves the current slot, so every request + names it — the point here is that the resident-but-idle engine + does not interfere and is not touched.""" + self.load("swap-a") + self.made[0].reply = reply_plain("from a") + results = [] + lock = threading.Lock() + + def one(): + status, body = self.chat(model="swap-a") + with lock: + results.append(status) + + threads = [threading.Thread(target=one) for _ in range(3)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + self.assertEqual(len(results), 3) + for status in results: + self.assertEqual(status, 200) + self.assertFalse(self.engine.closed) + + class TestConcurrency(ServerTestCase): def test_parallel_requests_all_answered(self): """Requests queue on the engine lock; none is dropped or mixed up.""" From c45aba38ae056af25296920b161a72beaa9b0eba Mon Sep 17 00:00:00 2001 From: Hendrik Date: Mon, 21 Sep 2026 15:47:14 +0200 Subject: [PATCH 2/7] serve: per-model state as one immutable slot; a queued request races no swap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A request used to read srv.engine before taking any lock, and read the other per-model facts (chat_format, model_info, stop_tokens, default_thinking, markers, chat_error) one attribute at a time. Two interleavings broke: a request granted the old engine's lock only after a swap had closed it called state_reset()/generate() on a dead ctx and answered 500; and a request could see the new engine with the old chat format — building an XTML prompt for a container that speaks chat.json — because _detect() re-bound the attributes one at a time under a lock the reader had not taken. The per-model facts are now one immutable ModelSlot (engine, model_id, model_info, markers, chat_format, chat_error, stop_tokens, default_thinking, plus the parser factory), built whole by ModelSlot.detect() and published in a single assignment. Each request takes current_slot() once, locks that snapshot's engine, and re-checks with check_engine() that the slot is still current before touching the engine — answering 409 model_switched with the new current model id otherwise. The same discipline covers /v1/completions, and the streaming/blocking tails report stats and model from the slot rather than re-reading the moving server mid-generation. --- serve/server.py | 433 ++++++++++++++++++++++++------------- tests/serve/test_server.py | 147 +++++++++++++ 2 files changed, 430 insertions(+), 150 deletions(-) diff --git a/serve/server.py b/serve/server.py index 197e5aba0..0da1dabc6 100644 --- a/serve/server.py +++ b/serve/server.py @@ -41,7 +41,7 @@ import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Callable, Optional +from typing import Callable, NamedTuple, Optional from . import api, dsml, glmtools, xtml from .chatfmt import ChatFormat, ChatFormatError, PlainParser @@ -66,56 +66,31 @@ def __init__(self, message: str): super().__init__("model load", WASTE_E_IO, message) -class ChatServer(ThreadingHTTPServer): - """Threaded HTTP, one engine, one lock.""" - - daemon_threads = True - allow_reuse_address = True +class ModelSlot(NamedTuple): + """Everything one container makes true, published as one object. - def __init__(self, addr, handler, *, engine: Engine, model_id: str, - api_key: Optional[str] = None, - default_max_tokens: int = 4096, - default_thinking: bool = True, - allow_local_images: bool = False, - log_requests: bool = True, - tmpdir: Optional[str] = None, - models: Optional[dict] = None, - keep_previous: bool = False, - engine_kwargs: Optional[dict] = None, - engine_factory: Optional[Callable] = None): - super().__init__(addr, handler) - self.engine_kwargs = dict(engine_kwargs or {}) - self.keep_previous = keep_previous - self.engine_factory = engine_factory or self._default_engine_factory - self._slot_lock = threading.RLock() - # model_id -> the Engine holding it. With keep_previous there can be - # more than one; without it, exactly one — the previous entry is - # closed and dropped as the new one becomes current. - self.engines = {model_id: engine} - # model_id -> container path. The loaded model is registered from - # its own path; --models adds the rest of the swappable set. An id - # already taken by the loaded model keeps the loaded path: the - # registry is a way to name what can be swapped to, not a way to - # point the loaded model somewhere else. - self.registry = {model_id: engine.model_path or model_id} - for mid, path in (models or {}).items(): - self.registry.setdefault(mid, path) - self.api_key = api_key - self.default_max_tokens = default_max_tokens - self.allow_local_images = allow_local_images - self.log_requests = log_requests - self.started = api.now() - self._tmp = tmpdir or tempfile.mkdtemp(prefix="waste-serve-") - self.tmpdir = self._tmp + The per-model facts used to be attributes the server re-bound one at + a time on every load, and handlers re-read one at a time. That let a + request racing a swap see the new engine with the old chat format — + or the old engine with the new one — and build a prompt for a format + the container it then generated on does not speak. So the block is a + value instead: `detect` builds a fully-resolved slot and the swap + publishes it in a single assignment. A request takes the slot once, + locks *its* engine, and re-checks the slot is still current before + generating; every fact it uses thereafter comes from the slot, not + from the moving server. + """ - self._start_thinking = default_thinking - self._detect(engine, model_id) - - def _default_engine_factory(self, path: str) -> Engine: - """How a swap opens a container. Overridable — tests hand in a - factory that builds the scripted engine, and a host embedding the - server may want its own construction arguments.""" - return Engine(path, **self.engine_kwargs) + model_id: str + engine: Engine + model_info: dict + markers: dict + # One of: the xtml module, the dsml module, a ChatFormat, or None — + # None with a chat_error, which refuses chat completions. + chat_format: object + chat_error: Optional[str] + stop_tokens: list + default_thinking: bool # ---- what the current model makes true ------------------------------- # @@ -123,25 +98,25 @@ def _default_engine_factory(self, path: str) -> Engine: # handler read its verdicts as constants for the life of the process. # With a swappable registry they are per-model: a container without # XTML markers must not inherit the previous container's chat format, - # stop tokens, or thinking default. So the same block is a method, run - # again on every load, and it keeps the old comment because every word - # of it still holds. + # stop tokens, or thinking default. So the same block is a factory, + # run again on every load, and it keeps the old comment because every + # word of it still holds. - def _detect(self, engine: Engine, model_id: str) -> None: - """Bind `engine` as the current model and re-derive everything the - handlers read from the container rather than from the request: - model_info, the reply format, stop tokens, the thinking default. + @classmethod + def detect(cls, engine: Engine, model_id: str, + start_thinking: bool) -> "ModelSlot": + """Derive everything the handlers read from the container rather + than from the request: model_info, the reply format, stop tokens, + the thinking default. The engine lock is NOT taken here — the caller holds it, or (at construction) no request can have arrived yet. """ - self.engine = engine - self.model_id = model_id - self.default_thinking = self._start_thinking + default_thinking = start_thinking try: - self.model_info = engine.model_info() + model_info = engine.model_info() except EngineError: - self.model_info = {} + model_info = {} # Markers by token id: the parser decides structure from ids, not # from what the text happens to spell. See regions.py. # @@ -162,16 +137,15 @@ def _detect(self, engine: Engine, model_id: str) -> None: # model would read its own turn structure as prose and answer # anyway. See #34. try: - self.markers = engine.marker_ids() - self.chat_format = xtml - self.chat_error = None - self.stop_tokens = [tid for tid, text in self.markers.items() - if text == "<|end_of_msg|>"] + markers = engine.marker_ids() + chat_format, chat_error = xtml, None + stop_tokens = [tid for tid, text in markers.items() + if text == "<|end_of_msg|>"] except EngineError as e: - self.markers = {} - self.chat_format = None - self.chat_error = str(e) - self.stop_tokens = [] + markers = {} + chat_format = None + chat_error = str(e) + stop_tokens = [] # DeepSeek-V4.1's DSML, before the declarative fallback and for # the same reason XTML comes before both: it is a whole protocol # — turns, thinking, tools, images — where chat.json is a plain @@ -179,21 +153,21 @@ def _detect(self, engine: Engine, model_id: str) -> None: # container that is not this release falls through rather than # half-resolving. try: - self.markers = dsml.detect(engine) + markers = dsml.detect(engine) except EngineError as e_ds: - self.markers = {} + markers = {} e = f"{e}; and {e_ds}" else: - self.chat_format = dsml - self.chat_error = None - self.stop_tokens = [tid for tid, text in self.markers.items() - if text == dsml.EOS] - # The generation prompt always opens a channel — or - # — so a DSML container cannot be asked to answer - # without one being chosen. Default it on, as the release - # does. - self.default_thinking = True - return + chat_format, chat_error = dsml, None + stop_tokens = [tid for tid, text in markers.items() + if text == dsml.EOS] + # The generation prompt always opens a channel — so a DSML + # container cannot be asked to answer without one being + # chosen. Default it on, as the release does. + default_thinking = True + return cls(model_id, engine, model_info, markers, + chat_format, chat_error, stop_tokens, + default_thinking) try: fmt = ChatFormat.load(engine) except ChatFormatError as e2: @@ -201,18 +175,171 @@ def _detect(self, engine: Engine, model_id: str) -> None: # reads as "wrong model" when the chat.json is simply # missing, and the chat.json reason alone hides that the # richer formats were tried first. - self.chat_error = f"{e}; and {e2}" + chat_error = f"{e}; and {e2}" else: - self.markers = fmt.markers - self.chat_format = fmt - self.chat_error = None - self.stop_tokens = [fmt.stop_id] + markers = fmt.markers + chat_format, chat_error = fmt, None + stop_tokens = [fmt.stop_id] # On by default only when the format names a channel. A # container without one refuses a request that asks for it # rather than answering without it; a container whose # generation prompt always opens one — GLM's does — cannot # be asked to answer without it either. - self.default_thinking = fmt.think is not None + default_thinking = fmt.think is not None + return cls(model_id, engine, model_info, markers, + chat_format, chat_error, stop_tokens, default_thinking) + + def new_parser(self, thinking: bool, tools=None): + """The reply reader for whichever format this container speaks. + + `thinking` says which channel the generation prompt left open. + XTML and DSML both have channels to leave open; a chat.json format + has one only when it names a think marker. + """ + if self.chat_format is xtml: + return RegionParser(in_think=thinking, in_response=not thinking, + markers=self.markers) + if self.chat_format is dsml: + return dsml.DSMLParser(thinking=thinking, markers=self.markers) + fmt = self.chat_format + # Which tool protocol the reply reader should own: a GLM container + # speaks its own `` grammar, anything else that reaches + # a PlainParser speaks Kimi K2's five control tokens. + tool_parser = None + if getattr(fmt, "tool_protocol", "") == "glm": + tool_parser = glmtools.ToolParser(tools=tools) + return PlainParser(markers=self.markers, + think_close_id=getattr(fmt, "think_close_id", -1), + in_think=thinking and getattr(fmt, "think", None) + is not None, + tool_parser=tool_parser) + + +class ChatServer(ThreadingHTTPServer): + """Threaded HTTP, one engine, one lock.""" + + daemon_threads = True + allow_reuse_address = True + + def __init__(self, addr, handler, *, engine: Engine, model_id: str, + api_key: Optional[str] = None, + default_max_tokens: int = 4096, + default_thinking: bool = True, + allow_local_images: bool = False, + log_requests: bool = True, + tmpdir: Optional[str] = None, + models: Optional[dict] = None, + keep_previous: bool = False, + engine_kwargs: Optional[dict] = None, + engine_factory: Optional[Callable] = None): + super().__init__(addr, handler) + self.engine_kwargs = dict(engine_kwargs or {}) + self.keep_previous = keep_previous + self.engine_factory = engine_factory or self._default_engine_factory + self._slot_lock = threading.RLock() + # model_id -> the Engine holding it. With keep_previous there can be + # more than one; without it, exactly one — the previous entry is + # closed and dropped as the new one becomes current. + self.engines = {model_id: engine} + # model_id -> container path. The loaded model is registered from + # its own path; --models adds the rest of the swappable set. An id + # already taken by the loaded model keeps the loaded path: the + # registry is a way to name what can be swapped to, not a way to + # point the loaded model somewhere else. + self.registry = {model_id: engine.model_path or model_id} + for mid, path in (models or {}).items(): + self.registry.setdefault(mid, path) + self.api_key = api_key + self.default_max_tokens = default_max_tokens + self.allow_local_images = allow_local_images + self.log_requests = log_requests + self.started = api.now() + self._tmp = tmpdir or tempfile.mkdtemp(prefix="waste-serve-") + self.tmpdir = self._tmp + + self._start_thinking = default_thinking + self._detect(engine, model_id) + + # The per-model facts live on the current ModelSlot and move only by + # whole-slot assignment; these accessors are how the handlers and the + # tests read them without taking a slot themselves. A request path + # must not use them per-fact — it takes `current_slot()` once. + + @property + def engine(self) -> Engine: + return self._slot.engine + + @property + def model_id(self) -> str: + return self._slot.model_id + + @property + def model_info(self) -> dict: + return self._slot.model_info + + @property + def markers(self) -> dict: + return self._slot.markers + + @property + def chat_format(self): + return self._slot.chat_format + + @property + def chat_error(self): + return self._slot.chat_error + + @property + def stop_tokens(self) -> list: + return self._slot.stop_tokens + + @property + def default_thinking(self) -> bool: + return self._slot.default_thinking + + def current_slot(self) -> ModelSlot: + """The current model's facts, as one consistent value.""" + with self._slot_lock: + return self._slot + + def check_engine(self, slot: ModelSlot) -> ModelSlot: + """Call with slot.engine.lock held, after acquiring it: is this + slot still the one the server serves? + + Without keep_previous the swap target may have been *closed* by + the time this request's lock was granted — an engine closes under + its own lock, so a request queued on it either ran first (fine) or + gets the lock after close. Answering on the closed engine is the + 500 the check exists to prevent; answering on the other engine is + the torn-state bug. Either way the request is not served: 409 with + the model that is current now, so the client re-reads + GET /v1/models and retries explicitly rather than being silently + migrated to a different model mid-conversation. + """ + current = self.current_slot() + if slot is not current: + raise api.APIError( + f"model switched to {current.model_id} while this request " + f"was queued; re-read GET /v1/models and retry", + status=409, type="model_switched", param="model") + return current + + def _default_engine_factory(self, path: str) -> Engine: + """How a swap opens a container. Overridable — tests hand in a + factory that builds the scripted engine, and a host embedding the + server may want its own construction arguments.""" + return Engine(path, **self.engine_kwargs) + + def _detect(self, engine: Engine, model_id: str) -> None: + """Bind `engine` as the current model and publish everything the + handlers read from the container — as one slot, in one assignment. + + The engine lock is NOT taken here — the caller holds it, or (at + construction) no request can have arrived yet. + """ + slot = ModelSlot.detect(engine, model_id, self._start_thinking) + with self._slot_lock: + self._slot = slot # ---- the model registry ---------------------------------------------- @@ -264,18 +391,23 @@ def load_model(self, model_id: str) -> Optional[str]: the engine can actually get. The old engine's lock is held across the whole swap, so a - generation in flight finishes before the slot moves under it, and - no request that took the old lock can find its engine closed. + generation in flight finishes before the slot moves under it. + A request that queued on the old engine and is granted the lock + only after the swap finds the slot moved and answers 409 — see + check_engine — rather than generating on a closed or wrong + container. The per-model facts move with the slot: one assignment + publishes engine, format, markers, stop tokens and thinking + default together, so no reader can see a half-updated set. """ with self._slot_lock: current = self._current() - if model_id == current[0]: + if model_id == current[0].model_id: return None path = self.registry.get(model_id) if path is None: raise api.APIError(f"no such model: {model_id}", status=404, type="not_found_error", param="model") - previous_id, previous_engine = current + previous_slot, previous_engine = current # A model kept resident by keep_previous does not need an # open at all — its waste_ctx still holds the state it had. # Moving the slot to it costs a format re-detect, not a load. @@ -283,7 +415,7 @@ def load_model(self, model_id: str) -> Optional[str]: with previous_engine.lock: if resident is not None: self._detect(resident, model_id) - return previous_id + return previous_slot.model_id try: engine = self.engine_factory(path) except EngineError as e: @@ -292,15 +424,15 @@ def load_model(self, model_id: str) -> Optional[str]: self._detect(engine, model_id) self.engines[model_id] = engine if self.keep_previous: - return previous_id - self.engines.pop(previous_id) + return previous_slot.model_id + self.engines.pop(previous_slot.model_id) previous_engine.close() - return previous_id + return previous_slot.model_id def _current(self) -> tuple: - with self._slot_lock: - model_id = self.model_id - return model_id, self.engines.get(model_id, self.engine) + """(the current ModelSlot, its engine). Call with _slot_lock held.""" + slot = self._slot + return slot, slot.engine def close_engines(self) -> None: """Every engine this server still holds. The shutdown path; with @@ -310,29 +442,10 @@ def close_engines(self) -> None: self.engines.clear() def new_parser(self, thinking: bool, tools=None): - """The reply reader for whichever format this container speaks. - - `thinking` says which channel the generation prompt left open. - XTML and DSML both have channels to leave open; a chat.json format - has one only when it names a think marker. - """ - if self.chat_format is xtml: - return RegionParser(in_think=thinking, in_response=not thinking, - markers=self.markers) - if self.chat_format is dsml: - return dsml.DSMLParser(thinking=thinking, markers=self.markers) - fmt = self.chat_format - # Which tool protocol the reply reader should own: a GLM container - # speaks its own `` grammar, anything else that reaches - # a PlainParser speaks Kimi K2's five control tokens. - tool_parser = None - if getattr(fmt, "tool_protocol", "") == "glm": - tool_parser = glmtools.ToolParser(tools=tools) - return PlainParser(markers=self.markers, - think_close_id=getattr(fmt, "think_close_id", -1), - in_think=thinking and getattr(fmt, "think", None) - is not None, - tool_parser=tool_parser) + """The reply reader for the current container. Reads the slot under + the lock, so a request that calls it (as the chat path does, inside + its locked section) gets the parser of the model it is serving.""" + return self.current_slot().new_parser(thinking, tools=tools) def handle_error(self, request, client_address): """A client hanging up is not an error worth a traceback. @@ -544,7 +657,13 @@ def _load_model(self): def _chat(self): body = self._read_body() srv = self.server - engine = srv.engine + # One snapshot, taken before the lock and used for everything this + # request decides: the engine it will lock, the format it renders + # the prompt in, the ctx it reads, the tokens it stops on. Reading + # them one attribute at a time off the server is what let a request + # racing a swap see the new engine with the old chat format. + slot = srv.current_slot() + engine = slot.engine self._log_model_from(body) # Before anything else, and before the engine lock: a request that @@ -558,10 +677,10 @@ def _chat(self): # 400 rather than 501 because for an OpenAI client the unsupported # thing is the model, which is a request parameter — and a 501 is # the one status those clients tend to retry. - if srv.chat_error: + if slot.chat_error: raise api.APIError( f"this model cannot be used for chat completions: " - f"{srv.chat_error}. serve/ renders Kimi K3's XTML prompt " + f"{slot.chat_error}. serve/ renders Kimi K3's XTML prompt " f"format and no other. POST /v1/completions for raw " f"continuation, or use `waste chat`, which reads the " f"container's own chat.json", @@ -584,39 +703,47 @@ def _chat(self): # between its own build and generate would hand the second # request the first one's pictures. with engine.lock: + # First thing under the lock: did the slot move between the + # snapshot and this grant? Without keep_previous the swap may + # also have *closed* this engine — close takes the same lock, + # so we are here either before it (our generation runs to + # completion, the swap waits) or after it, and the check + # refuses before state_reset touches the dead ctx. + srv.check_engine(slot) + engine.state_reset() prompt = api.build_prompt( engine, body, - default_thinking=srv.default_thinking, + default_thinking=slot.default_thinking, allow_local_images=srv.allow_local_images, - tmpdir=srv.tmpdir, fmt=srv.chat_format) + tmpdir=srv.tmpdir, fmt=slot.chat_format) opts = api.generation_options( body, default_max_tokens=srv.default_max_tokens, - ctx_max=srv.model_info.get("ctx_max", 0), + ctx_max=slot.model_info.get("ctx_max", 0), prompt_len=len(prompt.tokens)) stops = api.stop_strings(body) request_id = api.new_id("chatcmpl") created = api.now() - parser = srv.new_parser(prompt.thinking, tools=body.get("tools")) + parser = slot.new_parser(prompt.thinking, tools=body.get("tools")) if stream: - self._chat_stream(body, prompt, opts, stops, parser, + self._chat_stream(slot, body, prompt, opts, stops, parser, request_id, created) else: - self._chat_blocking(body, prompt, opts, stops, parser, + self._chat_blocking(slot, body, prompt, opts, stops, parser, request_id, created) - def _run(self, prompt, opts, stops, parser, on_delta): + def _run(self, slot, prompt, opts, stops, parser, on_delta): """Drive one generation. Returns (n_tokens, hit_limit, stopped). `on_delta(delta)` is called on the engine thread for each token, and may raise Cancelled to stop — which is how a disconnected streaming client stops the generation rather than paying for all of it. """ - engine = self.server.engine + engine = slot.engine stops = [s for s in stops if s] state = {"n": 0, "stopped": False, "content_sent": 0} @@ -664,7 +791,7 @@ def on_token(token_id, piece, info): temperature=opts["temperature"], top_p=opts["top_p"], top_k=opts["top_k"], seed=opts["seed"], max_tokens=opts["max_tokens"], - stop_tokens=self.server.stop_tokens or None) + stop_tokens=slot.stop_tokens or None) tail = parser.finish() if not state["stopped"]: if deliver(tail, final=True): @@ -672,23 +799,23 @@ def on_token(token_id, piece, info): hit_limit = completed and state["n"] >= opts["max_tokens"] return state["n"], hit_limit, state["stopped"] - def _chat_blocking(self, body, prompt, opts, stops, parser, + def _chat_blocking(self, slot, body, prompt, opts, stops, parser, request_id, created): t0 = time.time() - n, hit_limit, stopped = self._run(prompt, opts, stops, parser, + n, hit_limit, stopped = self._run(slot, prompt, opts, stops, parser, lambda d: None) reason = api.finish_reason(parser, hit_limit=hit_limit, stopped=stopped) usage = api.usage_block(len(prompt.tokens), n) payload = api.chat_completion( - parser, model=self.server.model_id, request_id=request_id, + parser, model=slot.model_id, request_id=request_id, created=created, reason=reason, usage=usage, - extra=api.engine_extra(self.server.engine.stats(), + extra=api.engine_extra(slot.engine.stats(), ms=(time.time() - t0) * 1000)) self._send_json(200, payload) - def _chat_stream(self, body, prompt, opts, stops, parser, + def _chat_stream(self, slot, body, prompt, opts, stops, parser, request_id, created): - model = self.server.model_id + model = slot.model_id self.send_response(200) self.send_header("Content-Type", "text/event-stream") self.send_header("Cache-Control", "no-cache") @@ -742,8 +869,8 @@ def on_delta(delta): raise Cancelled() try: - n, hit_limit, stopped = self._run(prompt, opts, stops, parser, - on_delta) + n, hit_limit, stopped = self._run(slot, prompt, opts, stops, + parser, on_delta) except Cancelled: # The client is gone. Nothing left to write to. return @@ -778,7 +905,7 @@ def on_delta(delta): "usage": usage}) write({"id": request_id, "object": "chat.completion.chunk", "created": created, "model": model, "choices": [], - "waste": api.engine_extra(self.server.engine.stats(), + "waste": api.engine_extra(slot.engine.stats(), ms=(time.time() - t0) * 1000)}) write("[DONE]") self.wfile.write(b"0\r\n\r\n") @@ -797,6 +924,11 @@ def _completions(self): """ body = self._read_body() srv = self.server + # Same discipline as _chat: one snapshot before the lock, and a + # re-check under it, so a request queued behind a swap answers 409 + # instead of generating on a closed (or wrong) engine. + slot = srv.current_slot() + engine = slot.engine self._log_model_from(body) # before check_model_request: a 404 or # 409 line should still name the model that was refused srv.check_model_request(body) @@ -809,15 +941,16 @@ def _completions(self): raise api.APIError("'prompt' must be a non-empty string", param="prompt") - with srv.engine.lock: - srv.engine.state_reset() # each request stands alone - tokens = srv.engine.tokenize(prompt_text) + with engine.lock: + srv.check_engine(slot) + engine.state_reset() # each request stands alone + tokens = engine.tokenize(prompt_text) if not tokens: raise api.APIError("'prompt' encoded to no tokens", param="prompt") opts = api.generation_options( body, default_max_tokens=srv.default_max_tokens, - ctx_max=srv.model_info.get("ctx_max", 0), + ctx_max=slot.model_info.get("ctx_max", 0), prompt_len=len(tokens)) stops = api.stop_strings(body) @@ -832,11 +965,11 @@ def on_token(token_id, piece, info): return False return True - completed = srv.engine.generate( + completed = engine.generate( tokens, on_token, temperature=opts["temperature"], top_p=opts["top_p"], top_k=opts["top_k"], seed=opts["seed"], max_tokens=opts["max_tokens"], - stop_tokens=srv.stop_tokens or None) + stop_tokens=slot.stop_tokens or None) text = "".join(pieces) for s in stops: @@ -847,7 +980,7 @@ def on_token(token_id, piece, info): "id": api.new_id("cmpl"), "object": "text_completion", "created": api.now(), - "model": srv.model_id, + "model": slot.model_id, "choices": [{"index": 0, "text": text, "logprobs": None, "finish_reason": "length" if hit_limit else "stop"}], "usage": api.usage_block(len(tokens), len(pieces)), diff --git a/tests/serve/test_server.py b/tests/serve/test_server.py index 975a5f5bd..ba2eccb68 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -18,6 +18,7 @@ import json import contextlib +import dataclasses import http.client import io import shutil @@ -29,6 +30,7 @@ import urllib.error import urllib.request from pathlib import Path +from typing import Callable, Optional REPO = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO)) @@ -1140,6 +1142,151 @@ def one(): self.assertFalse(self.engine.closed) +# A test-only engine that turns the request-vs-swap interleaving from a +# scheduler race into a certainty: the first acquire of its lock performs +# the swap before the lock is granted, so the request that snapshots the +# slot always finds the slot moved when it finally gets the engine. + +@dataclasses.dataclass +class RaceOnFirstLockEngine(FakeEngine): + """The interleaving that used to 500: the request reads srv.engine, + a swap takes the slot — and, without keep_previous, closes this + engine — and only then is the request granted the lock.""" + + on_first_lock: Optional[Callable] = None + + def __post_init__(self): + self._fired = False + @property + def lock(self): + outer = self + + class _Lock: + def acquire(self, *args, **kwargs): + if not outer._fired: + outer._fired = True + if outer.on_first_lock: + outer.on_first_lock() + return outer._lock.acquire(*args, **kwargs) + + def release(self): + outer._lock.release() + + def __enter__(self): + self.acquire() + return outer._lock + + def __exit__(self, *exc): + self.release() + return _Lock() + + +class TestRequestQueuedBehindSwap(ServerTestCase): + """A request that queued on the old engine's lock and is granted it + only after a swap must answer 409 naming the new model — not 500 on + a closed engine, and not a generation on the wrong container. The + engine is a RaceOnFirstLockEngine, so the swap always lands between + the request's snapshot and its lock grant.""" + + keep = False + log_requests = False + + def setUp(self): + self.swap_markers = dict(MARKERS) + self.made: list[FakeEngine] = [] + self.swapped_in: list[FakeEngine] = [] + self.engine = RaceOnFirstLockEngine(model_path="/fake/start.waste") + self.engine.on_first_lock = self.swap_behind_the_request + self.server = serve(self.engine, host="127.0.0.1", port=0, + model_id="test-model", + log_requests=False, + models={"swap-a": "/fake/a.waste", + "swap-b": "/fake/b.waste"}, + keep_previous=self.keep, + engine_factory=self.make_engine) + self.port = self.server.server_address[1] + self.thread = threading.Thread(target=self.server.serve_forever, + daemon=True) + self.thread.start() + + def make_engine(self, path: str) -> FakeEngine: + engine = FakeEngine(model_path=path, markers=self.swap_markers) + self.made.append(engine) + return engine + + def swap_behind_the_request(self): + """The swap, minus the close: it runs while the queued request + holds the old engine's lock, so the close is the test's job — + exactly as the real path defers it past the lock.""" + srv = self.server + engine = FakeEngine(model_path="/fake/a.waste", + markers=self.swap_markers) + self.swapped_in.append(engine) + srv._detect(engine, "swap-a") + srv.engines["swap-a"] = engine + if not self.keep: + srv.engines.pop("test-model") + + def test_queued_request_gets_409_not_a_closed_engine(self): + status, body = self.chat(model=None) + self.assertEqual(status, 409) + self.assertEqual(body["error"]["type"], "model_switched") + self.assertIn("swap-a", body["error"]["message"]) + # The old engine's state was never touched: no reset, no + # generation — the check fires before either. + self.assertEqual(self.engine.resets, 0) + self.assertEqual(self.engine.calls, []) + self.assertEqual(self.swapped_in[0].calls, []) + # And the server still serves, on the model it now holds. + self.engine.close() + self.engine.reply = reply_plain("after swap") + self.server.engine.reply = self.engine.reply + status, body = self.chat(model=None) + self.assertEqual(status, 200) + self.assertEqual(body["model"], "swap-a") + + def test_stale_slot_is_refused_even_when_engine_still_open(self): + """With keep_previous the old engine is not closed, but the + generation still must not happen on it: a stale slot is 409 + whichever way the engine lives.""" + # keep=True turns this into the keep_previous variant below. + + +class TestRequestQueuedBehindSwapKeepPrevious(TestRequestQueuedBehindSwap): + """Same race with --keep-previous: the old engine stays open, and a + request queued on it is still refused rather than served by the + container that is no longer current.""" + + keep = True + + def test_queued_request_gets_409_not_a_closed_engine(self): + status, body = self.chat(model=None) + self.assertEqual(status, 409) + self.assertEqual(body["error"]["type"], "model_switched") + self.assertIn("swap-a", body["error"]["message"]) + self.assertEqual(self.engine.resets, 0) + self.assertEqual(self.engine.calls, []) + # The engine was never closed — the refusal is about the slot, + # not about a dead ctx. + self.assertFalse(self.engine.closed) + self.server.engine.reply = self.engine.reply + status, body = self.chat(model=None) + self.assertEqual(status, 200) + self.assertEqual(body["model"], "swap-a") + + def test_stale_slot_is_refused_even_when_engine_still_open(self): + """check_engine refuses a stale slot directly, open engine and + all — the property the keep_previous torn-state bug violated.""" + from serve import api as api_mod + stale = self.server.current_slot() + self.server.load_model("swap-a") + with self.assertRaises(api_mod.APIError) as cm: + self.server.check_engine(stale) + self.assertEqual(cm.exception.status, 409) + self.assertEqual(cm.exception.type, "model_switched") + self.assertFalse(stale.engine.closed) + + class TestConcurrency(ServerTestCase): def test_parallel_requests_all_answered(self): """Requests queue on the engine lock; none is dropped or mixed up.""" From 980c8f8cabcd6b06ca9e353f6dc0e9745e219835 Mon Sep 17 00:00:00 2001 From: Hendrik Date: Mon, 21 Sep 2026 16:17:32 +0200 Subject: [PATCH 3/7] serve: a swap has to fit, and --models has to say so with --budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A swap opens the new container before closing the old one — the order that makes a failed open leave the server serving what it was serving — so two contexts are resident at once. docs/SERVE.md asked the operator to size --budget so that moment fits, and nothing checked it. With the default --budget 0 the moment was worse than "the sum of the two": each context sizes itself to up to 3/4 of waste_usable_ram (waste.h, waste_cfg.ram_budget_bytes), so a swap ran at ~1.5x what the process may use — a paging run, not a slow one. With --keep-previous the total is not two but every model ever loaded, which no startup check can price at all. - --models now requires an explicit --budget, and the server refuses to start unless 2 x budget fits, naming the largest budget that does; the same lines print under the startup banner and under --plan, which is where a budget gets chosen. - The resident set is counted at every load — each engine's declared budget, or the floor plus expert cache waste_memory_used reports for one that chose its own — and a load that would not fit is refused with 507 (insufficient_memory) *before* the container is opened: nothing is closed, nothing is half-loaded, the previous model keeps serving, and the message says what it needed next to what was already held. 507 rather than 503 because asking again cannot help. - A slot move to a model that is already resident is never refused: it allocates nothing. Evicting a resident model to make room is not done — changing what is resident behind a client's back is the failure --keep-previous exists to prevent. serve/api.py gains human_bytes, shared by the banner and the refusal, so the two cannot quote different figures for the same machine. A platform that will not report its RAM is warned about, not refused: a limit nobody can measure is not a limit. Tests: 9 new server tests (144, from 135) and tests/serve/test_main.py with 12 of its own — the arithmetic needs no container, no libwaste and no particular machine. Checked end to end against real containers too: the banner, /v1/models, a --keep-previous swap, and a 507 from the third resident at a declared budget of 24G on a 64G machine. --- CHANGELOG.md | 34 +++++++ docs/SERVE.md | 48 ++++++++- serve/__main__.py | 176 ++++++++++++++++++++++++++++++--- serve/api.py | 16 +++ serve/server.py | 163 ++++++++++++++++++++++++++++-- tests/serve/test_main.py | 144 +++++++++++++++++++++++++++ tests/serve/test_server.py | 197 +++++++++++++++++++++++++++++++++++++ 7 files changed, 754 insertions(+), 24 deletions(-) create mode 100644 tests/serve/test_main.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 84cf33a89..ad96b53c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,6 +27,40 @@ changed. Each entry names the section to read for the numbers behind it. ### Added +- **`--models` now has to prove its swap fits, and a load that would not + is a 507.** A swap opens the new container before closing the old one — + the order that makes a failed open leave the server serving what it was + serving — so two contexts are resident at once, and `docs/SERVE.md` + asked the operator to size `--budget` so that moment fits while nothing + checked. The default made it worse than "the sum of the two": with + `--budget 0` each context sizes itself to as much as 3/4 of + `waste_usable_ram()`, so a swap ran at ~1.5x what the process may use — + a paging run, not a slow one. + + `--models` now requires an explicit `--budget`, and the server refuses + to start unless `2 x budget` fits, naming the largest budget that does; + the same lines are printed under the startup banner and by `--plan`, + which is where a budget gets chosen. At runtime the resident set is + counted at every load — each engine's budget, or the floor and expert + cache `waste_memory_used` reports for one that chose its own — and a + load that would not fit is refused with **507** + (`insufficient_memory`) *before* the container is opened: nothing is + closed, nothing is half-loaded, the previous model keeps serving, and + the message says what it needed next to what was already held. 507 + rather than 503 because asking again cannot help. + + That is also what bounds `--keep-previous`, whose resident set grows + with every model ever switched to and which no startup check can price + in advance: the cap is derived from the budgets rather than from a + separate `--max-resident` count that could disagree with them, and a + slot move to a model that is already resident is never refused because + it allocates nothing. Evicting a resident model to make room was the + alternative and is not done — changing what is resident behind a + client's back is the failure `--keep-previous` exists to prevent. + + `serve/` goes to 144 server tests (from 135) plus a new + `tests/serve/test_main.py` of 12, for the arithmetic on its own. + - **A strict CI job for the K2 tool protocol**, the one GLM has had and which `ci.yml` used to have to exempt K2 from in so many words: "the same ground-truth rule the K2 template check in tests/run.sh applies, except diff --git a/docs/SERVE.md b/docs/SERVE.md index dcf850eaf..4134e16f1 100644 --- a/docs/SERVE.md +++ b/docs/SERVE.md @@ -350,7 +350,26 @@ What a swap does: that fails — a truncated container, an `--exclusive-open` conflict — leaves the server serving what it was serving and answers 500 with the engine's own reason. The cost of that guarantee is a moment where both - containers are resident; size `--budget` so that moment fits. + containers are resident, and that moment is checked rather than + assumed: `--models` **requires `--budget`**, and `2 x budget` has to + fit in `waste_usable_ram()` or the server refuses to start. The reason + is the default. With `--budget 0` each context sizes itself to as much + as 3/4 of usable RAM, so two of them at once is ~1.5x what the process + may use — paging, not slowness. At startup the numbers are printed: + + registry glm53, ds41 + glm53 floor 12.3 GB, recommended 30.1 GB + ds41 floor 18.7 GB, recommended 44.2 GB + two at once: 48.0 GB against 64.0 GB usable — fits + + `--plan` prints the same lines without starting anything, which is how + to choose the budget in the first place. +- A load that would not fit is refused **before the container is + opened**, with **507** and `type: insufficient_memory`: nothing is + closed, nothing is half-loaded, the previous model keeps serving, and + the message says what it needed next to what was already held. 507 + rather than 503 because asking again cannot help — an operator has to + lower `--budget`, drop `--keep-previous`, or restart. - Generation always serves the current model. A request naming a registered-but-not-loaded model is a 409, telling the client to `POST /v1/models/load` first, rather than an unnoticed multi-gigabyte @@ -363,11 +382,24 @@ What a swap does: `--keep-previous` keeps the replaced model resident instead of unloading it. Switching back to it is then a slot move rather than a reopen — its -`waste_ctx` and the state it holds are still there. The RAM two resident +`waste_ctx` and the state it holds are still there — and a slot move is +never refused, because it allocates nothing. The RAM two resident contexts need is the sum of their budgets; on the machines this engine targets that is usually the difference between working and paging, which is why unloading is the default. +It also makes the total unbounded by anything a startup check can know, +since every model switched to stays resident: three models at a budget of +20 GB under a 64 GB machine is not a pair anyone can validate in advance. +So the resident set is counted at each load — every engine's budget, or, +when it chose its own, the floor and expert cache `waste_memory_used` +reports for it — and the load that would put the sum over +`waste_usable_ram()` is refused with the same 507. There is no separate +`--max-resident` number to keep in step with the budgets: the cap is +derived from them. A model that does not fit is not evicted to make room +either — changing what is resident behind a client's back is the failure +`--keep-previous` exists to prevent. + Streaming is written straight from the token callback, on the thread holding the lock. A client hanging up propagates back as a return value the engine understands — the callback says stop, `waste_generate` unwinds, the @@ -493,7 +525,9 @@ K3_DIR=/Volumes/WasteDisk/k3 python3 tools/gen_xtml_goldens.py python3 -m serve MODEL [options] --host, --port, --model-id, --api-key - --budget SIZE hard RAM ceiling, e.g. 48G (0 = the engine chooses) + --budget SIZE hard RAM ceiling, e.g. 48G (0 = the engine chooses, + which is up to 3/4 of usable RAM per context — so + --models needs this set explicitly) --ctx N context tokens --threads N compute threads (0 = one per core) --cpus LIST restrict them to a cpu list, e.g. 0-5 or 0-2,6-8; @@ -510,9 +544,13 @@ python3 -m serve MODEL [options] --allow-local-images --models PATH[=ID] additional containers a client may switch to with POST /v1/models/load (repeatable; switching unloads - the model it replaces) + the model it replaces). Requires --budget, because a + swap holds two contexts at once: 2 x budget has to + fit in RAM or the server refuses to start --keep-previous keep the replaced model resident instead of unloading - it; RAM needed is then the sum of both budgets + it; every model switched to stays resident, so the + load that would put the set over the machine's RAM is + refused with 507 --plan print the memory plan and exit --no-log-requests silence the per-request log lines ``` diff --git a/serve/__main__.py b/serve/__main__.py index cf89b623d..f97211549 100644 --- a/serve/__main__.py +++ b/serve/__main__.py @@ -25,7 +25,7 @@ from .engine import (CACHE_LFRU, CACHE_LRU, # noqa: E402 WASTE_E_ARG, WASTE_E_BUSY, WASTE_E_UNSUPPORTED, Engine, EngineError, build_info, physical_ram, - plan_memory) + plan_memory, usable_ram) from .server import ModelLoadError, serve # noqa: E402 POLICIES = {"lfru": CACHE_LFRU, "lru": CACHE_LRU} @@ -55,12 +55,107 @@ def parse_registry(specs: list[str]) -> dict[str, str]: return registry -def human(n: float) -> str: - for unit in ("B", "KB", "MB", "GB", "TB"): - if n < 1024 or unit == "TB": - return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}" - n /= 1024 - return f"{n:.1f} TB" +class RegistryBudgetError(Exception): + """--models and --budget are not a pair that fits in this machine.""" + + +def check_registry_budget(models: dict[str, str], *, budget: int, + usable: int) -> None: + """Refuse a registry whose swap window does not fit. + + A swap holds two contexts at once: the new container is opened before + the old one is closed, deliberately, so that a failed open leaves the + server serving what it was serving. docs/SERVE.md asked the operator to + size --budget so that moment fits and nothing checked it — and with the + default budget of 0 the moment is far worse than "the sum of the two": + 0 means the engine sizes each context itself, up to 3/4 of + waste_usable_ram (waste.h, waste_cfg.ram_budget_bytes), so two of them + is ~1.5x what the process may use. That is a paging run, not a slow one. + + So, two failures, and they are different things to say: + + - No budget at all: the pair cannot be computed, and the number the + engine would pick is not a number this process should be allowed to + pick twice. Refused rather than guessed at. + - An explicit budget with no room for its pair: 2 x budget over usable. + The largest budget that fits is usable // 2, and naming that figure + is the difference between an error and a puzzle. + + `usable` is passed in rather than measured here: the caller prints it, + and a test needs no machine of its own. 0 means the platform would not + say (see main) and is neither a pass nor a failure — refusing to start + because a machine will not report its RAM is worse than the thing this + exists to prevent. The runtime check in serve/server.py (check_room) + still refuses a load that would exceed what is left. + """ + if not models or not usable: + return + machine = api.human_bytes(usable) + half = api.human_bytes(usable // 2) + if not budget: + per_ctx = api.human_bytes(usable - usable // 4) + raise RegistryBudgetError( + f"--models needs an explicit --budget: with 0 the engine sizes " + f"each context itself, up to {per_ctx} of the {machine} this " + f"process may use, and a swap holds two of them at once — the " + f"new container is opened before the old one is closed, so that " + f"a failed open leaves the server serving what it was serving. " + f"Give --budget {half} or less, or drop --models and serve one " + f"container.") + if 2 * budget > usable: + raise RegistryBudgetError( + f"--budget {api.human_bytes(budget)} does not fit twice: a swap " + f"holds the container being loaded and the one it replaces at " + f"the same time, which is {api.human_bytes(2 * budget)} against " + f"the {machine} this process may use. Use {half} or less, or " + f"drop --models.") + + +def describe_registry(models: dict[str, str], *, budget: int, usable: int, + ctx: int = 0, plan=plan_memory) -> list[str]: + """What --models will cost, as lines: one per container, then the + arithmetic a swap performs. + + Returned rather than printed so the same lines can stand under the + startup banner, under --plan, and under a refusal — the last being + where they are worth most, because that is when the operator is + choosing a --budget. + """ + lines = [] + widest = max((len(mid) for mid in models), default=0) + for mid, path in models.items(): + try: + p = plan(path, ctx) + except EngineError as e: + lines.append(f"{mid:<{widest}} unreadable: {e}") + continue + lines.append(f"{mid:<{widest}} floor {api.human_bytes(p.floor_bytes)}" + f", recommended {api.human_bytes(p.recommended_bytes)}") + if not usable: + lines.append("this platform reports no usable-RAM figure; the pair " + "is not checked") + elif not budget: + lines.append(f"no --budget: each context sizes itself to up to " + f"{api.human_bytes(usable - usable // 4)} of the " + f"{api.human_bytes(usable)} this process may use") + else: + pair = 2 * budget + if pair <= usable: + lines.append(f"two at once: {api.human_bytes(pair)} against " + f"{api.human_bytes(usable)} usable — fits") + else: + lines.append(f"two at once: {api.human_bytes(pair)} against " + f"{api.human_bytes(usable)} usable — does not fit; " + f"the largest --budget is " + f"{api.human_bytes(usable // 2)}") + return lines + + +# api.human_bytes, under the name the banner lines below were written +# with. One formatter for the plans, the registry lines and the 507 +# refusal a swap can answer with: an operator comparing them should not be +# doing two conversions. +human = api.human_bytes def parse_size(text: str) -> int: @@ -108,9 +203,14 @@ def main(argv=None) -> int: -d '{"model":"waste","messages":[{"role":"user","content":"hi"}]}' python3 -m serve ~/models/k3.waste --models ~/models/glm53.waste \\ - --models ~/models/deepseek41.waste=ds41 + --models ~/models/deepseek41.waste=ds41 --budget 24G + # --budget is required with --models: a swap opens the new + # container before closing the old one, so 2 x budget has to fit + # in RAM or the server refuses to start; # POST /v1/models/load {"model":"glm53"} swaps to it, unloading k3; - # add --keep-previous to hold both resident instead + # add --keep-previous to hold both resident instead — every model + # switched to stays resident, and a load that would put the set + # over this machine's RAM answers 507 """) ap.add_argument("model", help="path to the .waste container") ap.add_argument("--host", default="127.0.0.1", @@ -124,7 +224,9 @@ def main(argv=None) -> int: g = ap.add_argument_group("engine") g.add_argument("--budget", type=parse_size, default=0, metavar="SIZE", - help="hard RAM ceiling, e.g. 48G. 0 lets the engine choose") + help="hard RAM ceiling, e.g. 48G. 0 lets the engine choose " + "— up to 3/4 of the RAM this process may use, so a " + "swap (below) needs it set explicitly") g.add_argument("--ctx", type=bounded_int(0, (1 << 32) - 1), default=0, metavar="N", help="context tokens (0 = container default)") @@ -174,7 +276,10 @@ def main(argv=None) -> int: "with POST /v1/models/load (repeatable; id defaults " "to the file name without .waste). Switching " "unloads the model it replaces unless " - "--keep-previous") + "--keep-previous. Requires --budget: a swap holds " + "the new container and the old one at once, so " + "2 x budget must fit in RAM, and refusing to start " + "otherwise is the point — see docs/SERVE.md") s.add_argument("--keep-previous", action="store_true", help="keep a model resident when another is loaded. " "Off by default, and deliberately: the RAM two " @@ -183,7 +288,9 @@ def main(argv=None) -> int: "that is the difference between working and " "paging. Both models can then answer at once — " "each waste_ctx takes one caller, so each has its " - "own lock") + "own lock. Every model switched to stays resident, " + "so the load that would put the set over the " + "machine's RAM is refused with 507") s.add_argument("--plan", action="store_true", help="print the memory plan and exit without loading") s.add_argument("--no-log-requests", action="store_true", @@ -200,6 +307,30 @@ def main(argv=None) -> int: return 2 model_id = args.model_id or model.name.removesuffix(".waste") + registry = parse_registry(args.models) + # What this process may use, measured once: the startup check below, + # the banner, and every swap this server will perform all count + # against one number rather than three readings that can disagree. + # 0 means the platform would not say — see check_registry_budget. + try: + usable = usable_ram() + except EngineError: + usable = 0 + + # Priced before anything is opened, and skipped under --plan, which is + # the command that exists to tell an operator the numbers *before* + # they pick a --budget. + registry_lines = describe_registry(registry, budget=args.budget, + usable=usable, ctx=args.ctx) + if registry and not args.plan: + try: + check_registry_budget(registry, budget=args.budget, usable=usable) + except RegistryBudgetError as e: + print(f"{e}\n", file=sys.stderr) + for line in registry_lines: + print(f" {line}", file=sys.stderr) + return 2 + try: if args.plan: plan = plan_memory(str(model), args.ctx) @@ -216,6 +347,10 @@ def main(argv=None) -> int: f"(only with --vision)") if ram: print(f"\n this machine has {human(ram)}") + if registry: + print("\n registry") + for line in registry_lines: + print(f" {line}") return 0 engine = Engine( @@ -268,8 +403,9 @@ def main(argv=None) -> int: default_thinking=not args.no_thinking, allow_local_images=args.allow_local_images, log_requests=not args.no_log_requests, - models=parse_registry(args.models), + models=registry, keep_previous=args.keep_previous, + usable_ram=usable, engine_kwargs={ "ram_budget_bytes": args.budget, "ctx_tokens": args.ctx, @@ -321,6 +457,20 @@ def main(argv=None) -> int: print(f"chat from {model}/chat.json — plain conversation, " f"{think},\n {images}, {tools}") + # What a client may switch to, and what that costs — on the same + # lines as the banner rather than in a manual, because the swap + # window is the one number here an operator can still get wrong. + if registry: + print(f"{'registry':<9} {', '.join(sorted(registry))}") + for line in registry_lines: + print(f"{'':<9} {line}") + if registry and args.keep_previous: + over = (f"a load that would put them over {human(usable)} answers 507" + if usable else "the resident set is not checked on this " + "platform") + print(f"{'':<9} keep-previous: every model switched to stays " + f"resident;\n{'':<9} {over}") + shown = args.host if ":" not in args.host else f"[{args.host}]" print(f"\nlistening on http://{shown}:{args.port} " f"(POST {'/v1/completions' if srv.chat_error else '/v1/chat/completions'})") diff --git a/serve/api.py b/serve/api.py index 185d856d4..ad38abee6 100644 --- a/serve/api.py +++ b/serve/api.py @@ -55,6 +55,22 @@ def new_id(prefix: str) -> str: return f"{prefix}-{uuid.uuid4().hex[:24]}" +def human_bytes(n: float) -> str: + """A byte count in the units a person reads. + + Shared on purpose: `serve/__main__.py` prints memory plans with it and + `serve/server.py` refuses an over-budget swap with it. An operator + comparing "this budget fits" against "it does not" should not be doing + two conversions, and a refusal that quotes a different figure than the + banner it is arguing with is worse than no message. + """ + for unit in ("B", "KB", "MB", "GB", "TB"): + if n < 1024 or unit == "TB": + return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}" + n /= 1024 + return f"{n:.1f} TB" + + # ---- request validation -------------------------------------------------- diff --git a/serve/server.py b/serve/server.py index 0da1dabc6..baca845ab 100644 --- a/serve/server.py +++ b/serve/server.py @@ -45,7 +45,8 @@ from . import api, dsml, glmtools, xtml from .chatfmt import ChatFormat, ChatFormatError, PlainParser -from .engine import Cancelled, Engine, EngineError +from .engine import (Cancelled, Engine, EngineError, plan_memory, + usable_ram) from .regions import RegionParser SERVER_NAME = "waste" @@ -231,9 +232,28 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, models: Optional[dict] = None, keep_previous: bool = False, engine_kwargs: Optional[dict] = None, - engine_factory: Optional[Callable] = None): + engine_factory: Optional[Callable] = None, + usable_ram: Optional[int] = None, + memory_plan: Optional[Callable] = None): super().__init__(addr, handler) self.engine_kwargs = dict(engine_kwargs or {}) + # The ceiling every swap-open is given, when there is one. It is + # also the number the resident-set check counts with: an explicit + # budget is used exactly as given (waste.h), so a resident model's + # footprint is known without asking the container. 0 means the + # engine sizes itself, and then the check asks instead. + self.ram_budget_bytes = int(self.engine_kwargs.get("ram_budget_bytes") + or 0) + # What the machine may use, in bytes: passed in by `serve/__main__` + # (which measures it once for its own startup checks), or measured + # on first use. None = not known yet; 0 = this platform would not + # say, and nothing is refused on a number nobody has. + self._usable_ram = usable_ram + # How a container's own plan is read, for the check below and for + # a host whose containers are not files on this machine. The real + # one reads the manifest only — no weights — so asking it before + # an open costs nothing next to the open. + self.memory_plan = memory_plan or plan_memory self.keep_previous = keep_previous self.engine_factory = engine_factory or self._default_engine_factory self._slot_lock = threading.RLock() @@ -366,6 +386,123 @@ def check_model_request(self, body: dict) -> None: f"to switch to it", status=409, type="model_not_loaded", param="model") + # ---- what the machine has room for ----------------------------------- + # + # A swap holds two contexts at once: the new container is opened before + # the old one is closed, deliberately, so that a failed open leaves the + # server serving what it was serving. docs/SERVE.md used to ask the + # operator to size --budget so that moment fits, and nothing checked; + # with the default budget of 0 each context sizes itself to up to 3/4 + # of waste_usable_ram (waste.h), so the moment was ~1.5x what the + # process may use. With --keep-previous the total is not two but every + # model ever loaded, which no amount of sizing at startup can bound. + # + # So the resident set is counted, against the budgets it was opened + # with, and a load that would not fit is refused *before* the open. + + def usable_ram_bytes(self) -> int: + """What this process may use, in bytes. 0 = this platform will not + say, in which case nothing here refuses anything: a limit nobody + can measure is not a limit, and refusing to serve a model because + the machine would not report its RAM is worse than the thing this + check exists to prevent.""" + if self._usable_ram is None: + try: + self._usable_ram = int(usable_ram()) + except EngineError: + self._usable_ram = 0 + return self._usable_ram + + def engine_budget(self, engine: Engine) -> int: + """What one resident model holds. + + An explicit budget is exactly what the engine was opened with — + waste.h calls it a hard ceiling on all engine allocations, used as + given — and every open goes through `engine_kwargs`, whose + contract is that it holds the same arguments the startup engine + was opened with. When the engine chose its own, the ceiling is + whatever its ladder resolved to, and waste_memory_used reports the + result: the plan's floor plus the expert cache it actually + allocated (n_slots x record bytes). Neither is an estimate; both + come from the context that holds the memory. + """ + if self.ram_budget_bytes: + return self.ram_budget_bytes + try: + used = engine.memory_used() + except EngineError: + return 0 + return int(used.get("floor_bytes", 0) + + used.get("min_expert_cache", 0)) + + def planned_budget(self, path: str) -> int: + """What a container would ask for when no --budget is configured. + + A prediction, and named one. With 0 the engine walks its own ladder + — floor plus whole expert working sets, under 3/4 of usable RAM — + and nothing here can run that ladder without opening the container. + It uses the ladder's own definition of "worth having", + recommended_bytes, which is the number waste_plan_memory exists to + give. The CLI never relies on this: --models requires an explicit + --budget, so every figure on that path is exact. + """ + try: + ctx = int(self.engine_kwargs.get("ctx_tokens") or 0) + plan = self.memory_plan(path, ctx) + except (EngineError, OSError): + return 0 + return int(getattr(plan, "recommended_bytes", 0)) + + def check_room(self, model_id: str, path: str) -> None: + """Refuse a load that would not fit next to what stays resident. + + Call with _slot_lock held, *before* the container is opened. Only + the case that needs a new context is checked: a load that moves the + slot to an engine already in `engines` — what --keep-previous makes + free — adds no memory and is never refused. + + What stays resident is the whole point of the two modes: without + keep_previous the previous model alone is resident alongside the + new one until the swap completes, so the number is 2 x budget; with + it, every model in `engines` stays and the total is their sum. That + sum is the cap --keep-previous needs — it derives one from the + budgets rather than from a hand-set count, so it cannot disagree + with what the engines actually hold. + + Evicting a resident model to make room was the alternative and is + rejected: it changes what is resident behind a client's back, which + is the exact failure --keep-previous exists to prevent. + + 507 rather than 503: the request is well-formed and the server will + never satisfy it by being asked again — an operator has to change + --budget, drop --keep-previous, or restart. A 503 invites a retry + loop that cannot end. + + Nothing is opened, closed or moved when this refuses, so the model + that was serving keeps serving and answers with the reason. + """ + usable = self.usable_ram_bytes() + if not usable: + return + if self.keep_previous: + stays = sorted(self.engines) + else: + stays = [self._current()[0].model_id] + held = sum(self.engine_budget(self.engines[m]) for m in stays) + need = self.ram_budget_bytes or self.planned_budget(path) + if held + need <= usable: + return + raise api.APIError( + f"loading {model_id} would need " + f"{api.human_bytes(need)} next to the " + f"{api.human_bytes(held)} held by " + f"{', '.join(stays)}: {api.human_bytes(held + need)} against " + f"the {api.human_bytes(usable)} this machine may use. The " + f"current model is still serving. Lower --budget, drop " + f"--keep-previous, or start a second server for a resident set " + f"this size cannot hold", + status=507, type="insufficient_memory", param="model") + def load_model(self, model_id: str) -> Optional[str]: """Make `model_id` the current model, and return the id of the model that was current before (None when it already was). @@ -387,8 +524,8 @@ def load_model(self, model_id: str) -> Optional[str]: guarantee that is not to have closed it yet. A failed swap costs a load's worth of RAM for a moment; a swap that leaves no model loaded costs the whole server. When two containers genuinely will - not fit together, size --budget so each open plans against what - the engine can actually get. + not fit together the load is refused before the open — see + check_room, and 507 — rather than being attempted into paging. The old engine's lock is held across the whole swap, so a generation in flight finishes before the slot moves under it. @@ -416,6 +553,11 @@ def load_model(self, model_id: str) -> Optional[str]: if resident is not None: self._detect(resident, model_id) return previous_slot.model_id + # Room for it, before a byte of it is allocated — and + # before anything is closed, so a refusal here is + # indistinguishable from never having been asked: same + # current model, same resident set, same open containers. + self.check_room(model_id, path) try: engine = self.engine_factory(path) except EngineError as e: @@ -994,7 +1136,9 @@ def serve(engine: Engine, *, host: str = "127.0.0.1", port: int = 8000, ready: Optional[threading.Event] = None, models: Optional[dict] = None, keep_previous: bool = False, engine_kwargs: Optional[dict] = None, - engine_factory: Optional[Callable] = None) -> ChatServer: + engine_factory: Optional[Callable] = None, + usable_ram: Optional[int] = None, + memory_plan: Optional[Callable] = None) -> ChatServer: """Build the server. The caller decides whether to serve_forever. models names the rest of the swappable registry (id -> container @@ -1003,6 +1147,12 @@ def serve(engine: Engine, *, host: str = "127.0.0.1", port: int = 8000, on a swap — the same arguments the startup engine was opened with — and engine_factory replaces that factory wholesale, for a host that builds engines its own way (tests do exactly that). + + usable_ram is what the resident-set check counts against — pass the + figure already measured so a load never probes for it — and + memory_plan is how a container with no explicit budget is priced. + Both default to the real thing, and neither is needed by a caller + that never swaps models. """ # IPv6-capable when the host asks for it, without forcing it: binding # :: on a host with IPv6 disabled fails outright. @@ -1015,7 +1165,8 @@ def serve(engine: Engine, *, host: str = "127.0.0.1", port: int = 8000, log_requests=log_requests, models=models, keep_previous=keep_previous, engine_kwargs=engine_kwargs, - engine_factory=engine_factory) + engine_factory=engine_factory, + usable_ram=usable_ram, memory_plan=memory_plan) if ready is not None: ready.set() return srv diff --git a/tests/serve/test_main.py b/tests/serve/test_main.py new file mode 100644 index 000000000..cde2c52b3 --- /dev/null +++ b/tests/serve/test_main.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +""" +test_main.py — what `python3 -m serve` decides before it opens anything. + +The decision under test is the one that used to be advice in a manual: a +registry of swappable containers is only allowed to start if its swap +window fits. A swap holds two contexts at once — the new container is +opened before the old one is closed, which is the rollback guarantee — and +with the default `--budget 0` each context sizes itself to up to 3/4 of +waste_usable_ram, so the window was ~1.5x what the process may use. + +It is arithmetic over numbers, so it is tested as arithmetic: this file +needs no container, no libwaste and no particular machine. Everything that +does need those — loading, serving, the banner — is in test_engine.py and +test_server.py. + + python3 tests/serve/test_main.py +""" + +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from serve.__main__ import (RegistryBudgetError, check_registry_budget, + describe_registry, human) +from serve.engine import EngineError + +GB = 1 << 30 +MODELS = {"glm53": "/fake/glm53.waste", "ds41": "/fake/ds41.waste"} + + +class TestRegistryBudget(unittest.TestCase): + """--models is refused unless --budget shows that 2 x budget fits.""" + + def refuse(self, **kwargs) -> str: + with self.assertRaises(RegistryBudgetError) as cm: + check_registry_budget(MODELS, **kwargs) + return str(cm.exception) + + def test_a_registry_without_a_budget_is_refused(self): + """0 means "the engine chooses", and what it chooses is per-context + — which is exactly how the window got to 1.5x the machine.""" + message = self.refuse(budget=0, usable=64 * GB) + self.assertIn("--budget", message) + self.assertIn("48.0 GB", message) # what 0 resolves to per ctx + self.assertIn("64.0 GB", message) # what the process may use + self.assertIn("32.0 GB", message) # the largest that fits twice + + def test_a_pair_that_does_not_fit_is_refused(self): + message = self.refuse(budget=40 * GB, usable=64 * GB) + self.assertIn("40.0 GB", message) + self.assertIn("80.0 GB", message) # what the moment would cost + self.assertIn("32.0 GB", message) # and what to use instead + + def test_exactly_half_fits(self): + """2 x budget == usable is the boundary and it is a pass: the + engine's own cap already leaves a quarter of its budget to the OS, + so refusing this would refuse a machine with nothing else running.""" + check_registry_budget(MODELS, budget=32 * GB, usable=64 * GB) + + + def test_one_byte_over_the_boundary_is_refused(self): + self.assertIn("does not fit twice", + self.refuse(budget=32 * GB + 1, usable=64 * GB)) + + def test_a_single_model_is_never_checked(self): + """No registry, no swap: one context against the machine is the + engine's own business, whatever the budget says.""" + check_registry_budget({}, budget=0, usable=0) + check_registry_budget({}, budget=GB, usable=GB // 4) + + def test_an_unmeasurable_machine_is_not_refused(self): + """usable 0 = the platform would not say. Refusing to start over a + number nobody can read is worse than the failure this prevents; + serve/server.py's runtime check still refuses the load itself.""" + check_registry_budget(MODELS, budget=0, usable=0) + check_registry_budget(MODELS, budget=64 * GB, usable=0) + +class TestRegistryLines(unittest.TestCase): + """What the operator is shown when choosing a budget: the startup + banner, `--plan`, and the output of a refusal all print these.""" + + class Plan: + floor_bytes = 12 * GB + recommended_bytes = 30 * GB + + def plan(self, path: str, ctx: int): + return self.Plan() + + def test_each_container_is_priced(self): + lines = describe_registry(MODELS, budget=24 * GB, usable=64 * GB, + plan=self.plan) + self.assertEqual(len(lines), 3) + self.assertIn("glm53", lines[0]) + self.assertIn("floor 12.0 GB", lines[0]) + self.assertIn("recommended 30.0 GB", lines[0]) + self.assertIn("ds41", lines[1]) + + def test_the_arithmetic_verdicts(self): + fits = describe_registry(MODELS, budget=24 * GB, usable=64 * GB, + plan=self.plan) + self.assertIn("48.0 GB against 64.0 GB usable — fits", fits[2]) + over = describe_registry(MODELS, budget=48 * GB, usable=64 * GB, + plan=self.plan) + self.assertIn("does not fit", over[2]) + self.assertIn("32.0 GB", over[2]) # the largest that fits + + def test_no_budget_says_what_the_engine_would_take(self): + lines = describe_registry(MODELS, budget=0, usable=64 * GB, + plan=self.plan) + self.assertIn("no --budget", lines[2]) + self.assertIn("48.0 GB", lines[2]) # 3/4 of usable, per context + + def test_an_unreadable_container_does_not_hide_the_others(self): + def plan(path: str, ctx: int): + raise EngineError("plan_memory", -2, path) + + lines = describe_registry(MODELS, budget=8 * GB, usable=64 * GB, + plan=plan) + self.assertTrue(all("unreadable" in line for line in lines[:2])) + self.assertIn("fits", lines[2]) + + def test_an_unmeasurable_machine_says_so(self): + lines = describe_registry(MODELS, budget=8 * GB, usable=0, + plan=self.plan) + self.assertIn("not checked", lines[2]) + + +class TestHumanBytes(unittest.TestCase): + def test_units(self): + self.assertEqual(human(0), "0 B") + self.assertEqual(human(1024), "1.0 KB") + self.assertEqual(human(3 * GB), "3.0 GB") + self.assertEqual(human(2 * (1 << 40)), "2.0 TB") + + + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/serve/test_server.py b/tests/serve/test_server.py index ba2eccb68..787a498f6 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -1142,6 +1142,203 @@ def one(): self.assertFalse(self.engine.closed) +class TestSwapMemoryBudget(ServerTestCase): + """A swap has to fit in the machine. + + The new container is opened before the old one is closed — that order + is the rollback guarantee — so two contexts are resident at once, and + --keep-previous makes the set grow with every model ever loaded. The + server counts the resident budgets against what the process may use + and refuses a load that would exceed it *before* the open, so nothing + is closed and nothing is half-loaded: the previous model keeps + serving. + + The budgets are the test's own. `usable_ram` is passed in, so what is + asserted is the arithmetic rather than the machine the suite happens + to run on, and the engines are scripted, so a budget of 10 is 10 and + not whatever a real container would resolve to. + """ + + keep = False + usable = 25 + budget = 10 + + def make_engine(self, path: str) -> FakeEngine: + engine = FakeEngine(model_path=path, markers=dict(MARKERS)) + self.made.append(engine) + return engine + + def setUp(self): + self.made: list[FakeEngine] = [] + self.engine_kwargs = {"model_path": "/fake/start.waste"} + self.server_kwargs = { + "models": {"swap-a": "/fake/a.waste", "swap-b": "/fake/b.waste"}, + "keep_previous": self.keep, + "engine_factory": self.make_engine, + "engine_kwargs": {"ram_budget_bytes": self.budget}, + "usable_ram": self.usable, + } + ServerTestCase.setUp(self) + + def load(self, model): + return self.post("/v1/models/load", {"model": model}) + + +class TestSwapMemoryBudgetFits(TestSwapMemoryBudget): + """Two budgets of 10 under a 25-byte machine: the window fits, and the + swap behaves exactly as it did before the check existed.""" + + def test_a_pair_that_fits_loads(self): + status, body = self.load("swap-a") + self.assertEqual(status, 200) + self.assertEqual(body["loaded"], "swap-a") + self.assertTrue(self.engine.closed) # replaced, as always + self.assertEqual(list(self.server.engines), ["swap-a"]) + + def test_the_check_counts_a_budget_against_the_machine(self): + """The number is the ceiling the factory was told to open with, + not a measurement taken after the fact.""" + self.assertEqual(self.server.engine_budget(self.engine), self.budget) + self.assertEqual(self.server.usable_ram_bytes(), self.usable) + + +class TestSwapMemoryBudgetTooSmall(TestSwapMemoryBudget): + """2 x budget does not fit in usable: the swap window itself is over.""" + + usable = 15 + + def test_a_pair_that_does_not_fit_is_507(self): + status, body = self.load("swap-a") + self.assertEqual(status, 507) + self.assertEqual(body["error"]["type"], "insufficient_memory") + message = body["error"]["message"] + self.assertIn("swap-a", message) + self.assertIn("10 B", message) # what it would need + self.assertIn("20 B", message) # against 2 x budget + self.assertIn("15 B", message) # and what the machine has + + def test_nothing_moved_on_a_refusal(self): + self.load("swap-a") + # No engine was built, nothing was closed, and the model that was + # serving is still the one the server reports. + self.assertEqual(self.made, []) + self.assertFalse(self.engine.closed) + self.assertEqual(list(self.server.engines), ["test-model"]) + self.assertEqual(self.server.model_id, "test-model") + self.assertEqual(self.server.engine, self.engine) + # And it still answers — the refusal cost this server nothing. + status, body = self.chat() + self.assertEqual(status, 200) + self.assertEqual(body["model"], "test-model") + + +class TestSwapMemoryBudgetKeepPrevious(TestSwapMemoryBudget): + """--keep-previous: every model switched to stays resident, so the set + is unbounded by anything --budget alone can say. The ledger is what + bounds it, and it is derived from the same budgets.""" + + keep = True + + def test_the_resident_set_grows_until_it_would_not_fit(self): + status, _ = self.load("swap-a") + self.assertEqual(status, 200) # 2 x 10 <= 25 + status, body = self.load("swap-b") + self.assertEqual(status, 507) # 3 x 10 > 25 + self.assertEqual(body["error"]["type"], "insufficient_memory") + message = body["error"]["message"] + self.assertIn("swap-b", message) + self.assertIn("10 B", message) # wanted + self.assertIn("20 B", message) # held, by name: + self.assertIn("swap-a", message) + self.assertIn("test-model", message) + self.assertIn("30 B", message) # the moment's total + self.assertIn("25 B", message) # against what the machine has + # The set that was already resident is untouched, and one of its + # models is still the one serving. + self.assertEqual(sorted(self.server.engines), + ["swap-a", "test-model"]) + self.assertEqual(self.server.model_id, "swap-a") + self.assertFalse(any(e.closed for e in self.server.engines.values())) + + def test_switching_back_to_a_resident_model_is_never_refused(self): + """No open, no memory: a slot move costs a re-detect. This is the + property of --keep-previous that must not be taxed by the check.""" + self.load("swap-a") # now at the limit: 20/25 + status, body = self.load("test-model") + self.assertEqual(status, 200) + self.assertEqual(body["loaded"], "test-model") + self.assertEqual([e.model_path for e in self.made], + ["/fake/a.waste"]) + + +class TestSwapMemoryBudgetWithoutABudget(ServerTestCase): + """No --budget: the engine sizes each context itself, so there is no + ceiling to count with and the check prices the container with + waste_plan_memory instead — recommended_bytes, the ladder's own + definition of "worth having". The CLI does not rely on this, because + --models requires an explicit --budget; a host that calls serve() + directly is the case it covers.""" + + usable = 12 + recommended = 8 + + class Plan: + def __init__(self, recommended: int): + self.floor_bytes = 4 + self.recommended_bytes = recommended + + def plan(self, path: str, ctx: int): + return self.Plan(self.recommended) + + def make_engine(self, path: str) -> FakeEngine: + engine = FakeEngine(model_path=path, markers=dict(MARKERS)) + self.made.append(engine) + return engine + + def server_kwargs_for(self, usable: int) -> dict: + return { + "models": {"swap-a": "/fake/a.waste"}, + "engine_factory": self.make_engine, + "usable_ram": usable, + "memory_plan": self.plan, + } + + def setUp(self): + self.made: list[FakeEngine] = [] + self.engine_kwargs = {"model_path": "/fake/start.waste"} + self.server_kwargs = self.server_kwargs_for(self.usable) + ServerTestCase.setUp(self) + + def test_a_resident_without_a_budget_is_measured_from_the_context(self): + """FakeEngine reports floor 4 + cache 1, so a resident that chose + its own budget counts as 5 — the same figures waste_memory_used + returns for a real one.""" + self.assertEqual(self.server.engine_budget(self.engine), 5) + + def test_a_container_without_a_budget_is_priced_by_its_plan(self): + from serve import api as api_mod + with self.assertRaises(api_mod.APIError) as cm: + self.server.load_model("swap-a") + self.assertEqual(cm.exception.status, 507) + self.assertEqual(cm.exception.type, "insufficient_memory") + message = str(cm.exception) + self.assertIn("8 B", message) # recommended, from the plan + self.assertIn("5 B", message) # the resident, measured + self.assertIn("13 B", message) # against 12 B usable + self.assertEqual(self.made, []) + + def test_the_same_load_fits_a_slightly_larger_machine(self): + srv = serve(self.engine, host="127.0.0.1", port=0, + model_id="test-model", log_requests=False, + **self.server_kwargs_for(self.usable + 1)) + self.addCleanup(srv.server_close) + self.assertEqual(srv.load_model("swap-a"), "test-model") + self.assertEqual(len(self.made), 1) + + + + + # A test-only engine that turns the request-vs-swap interleaving from a # scheduler race into a certainty: the first acquire of its lock performs # the swap before the lock is granted, so the request that snapshots the From b34cb47dffd8263e9080aa4b73e2cfa5ea68e985 Mon Sep 17 00:00:00 2001 From: Hendrik Date: Tue, 22 Sep 2026 15:04:23 +0200 Subject: [PATCH 4/7] serve: strict model validation is what --models opts into MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The registry made check_model_request strict for every deployment: with no --models the registry holds only the loaded model, so a request naming anything else was a 404 — breaking every client that sends a fixed model name it cannot easily change, including serve --help's own example of "model": "waste". The server's id defaults to the container's file name (k3.waste -> k3), so even the correct name was one most clients did not know. Validation now applies only when the registry names a swappable set — len(registry) > 1, which is what --models gives — and a single-container server answers any name with the one model it has, as it did before the registry existed. Strict validation is unchanged where it is wanted: a 404 for a name the registry does not know, a 409 for one it knows but has not loaded, on both chat completions and raw completions, with the model-before-shape ordering the OpenAI API uses. docs/SERVE.md no longer describes the 404 as the default behaviour, and the CHANGELOG carries the break-and-restore under Unreleased. Tests: the no-registry case is covered where it lives now (a foreign name served, shape validation first), the registry case keeps its 404/409 and ordering tests, and /v1/completions gets the coverage chat had. Suite is 149 server tests, all green. --- CHANGELOG.md | 15 ++++++++++++++ docs/SERVE.md | 18 ++++++++++------- serve/server.py | 27 +++++++++++++++++++------ tests/serve/test_server.py | 41 +++++++++++++++++++++++++++++++------- 4 files changed, 81 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ad96b53c5..ee9f659ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,21 @@ changed. Each entry names the section to read for the numbers behind it. ### Fixed +- **A single-container server serves any model name again.** The model + registry made `check_model_request` strict for every deployment: with + no `--models`, the registry holds only the loaded model, so a request + naming anything else was a 404 — including the fixed name most OpenAI + clients are configured with and the `"model": "waste"` example in + `serve --help` itself, neither of which a client can easily change. + The server's own id defaults to the container's file name + (`k3.waste` → `k3`), so even the "correct" name was one most clients + did not know. What an existing deployment guaranteed was "any name is + answered by the one container"; that is restored: absent, the loaded + id, or any other string all serve. Strict validation — a 404 for a + name the registry does not know, a 409 for one it knows but has not + loaded — is what `--models` opts into, and is unchanged there. Tests + for both behaviours in `tests/serve/test_server.py`. + - **The K2 tool-grammar check defaulted to a path on an external volume**, which is a description of one machine rather than a default: everywhere else — CI, a fresh clone, this machine with the disk unplugged — it read diff --git a/docs/SERVE.md b/docs/SERVE.md index 4134e16f1..ef0af5144 100644 --- a/docs/SERVE.md +++ b/docs/SERVE.md @@ -321,13 +321,17 @@ the answer. ### Swapping models -By default the process serves the one container it was started with, and a -request naming any other model is a 404 — before the registry existed any -name was silently served by the loaded model, which made `model` a -decorative string. - -`--models PATH[=ID]` (repeatable) registers additional containers a client -may switch to: +By default the process serves the one container it was started with, and +a request may name it by any model string: absent, the container's own +id — the file name, `k3.waste` → `k3` — or anything else a client sends, +including the fixed name most OpenAI clients are configured with. There +is nothing to be strict about until there is a set to be strict against. + +`--models PATH[=ID]` (repeatable) registers additional containers a +client may switch to, and opts into strict validation: a request naming a +model the registry does not know is a 404, rather than being silently +served by the loaded model, which is what made `model` a decorative +string before the registry existed. python3 -m serve ~/models/k3.waste \ --models ~/models/glm53.waste --models ~/models/deepseek41.waste=ds41 diff --git a/serve/server.py b/serve/server.py index baca845ab..7d3014937 100644 --- a/serve/server.py +++ b/serve/server.py @@ -368,16 +368,31 @@ def check_model_request(self, body: dict) -> None: Absent, empty, or equal to the loaded model's id: fine — that is what every client that does not know about this registry sends, - and it must keep working. A name the registry does not know is a - 404: before the registry existed any name was silently served by - the loaded model, which made `model` a decorative string. A name - the registry knows but that is not resident is a 409, not a - surprise multi-gigabyte swap in the middle of a conversation — - the client asks for that explicitly with POST /v1/models/load. + and it must keep working. With a registry — `--models` was given — + a name it does not know is a 404: before the registry existed any + name was silently served by the loaded model, which made `model` + a decorative string, and an operator who names a swappable set + has said a wrong name should say so rather than answer as another + model. A name the registry knows but that is not resident is a + 409, not a surprise multi-gigabyte swap in the middle of a + conversation — the client asks for that explicitly with + POST /v1/models/load. + + Without `--models` the registry holds only the loaded model, and + nothing here rejects anything: a single-container server has no + swappable set to defend, its id defaults to the container's file + name, and clients — including the example in `serve --help` — + send a fixed model name they cannot easily change. Validating + every name would 404 deployments that answered them yesterday; + strict validation is what `--models` opts into. """ mid = body.get("model") if not isinstance(mid, str) or not mid or mid == self.model_id: return + # No --models, no registry to be strict about: the only entry is + # the loaded model itself, so anything named is served by it. + if len(self.registry) < 2: + return if mid not in self.registry: raise api.APIError(f"no such model: {mid}", status=404, type="not_found_error", param="model") diff --git a/tests/serve/test_server.py b/tests/serve/test_server.py index 787a498f6..c5bec5d8f 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -415,14 +415,24 @@ def test_missing_messages(self): self.assertEqual(status, 400) self.assertEqual(body["error"]["param"], "messages") - def test_unknown_model_outranks_missing_messages(self): - # The model is validated before the request's shape, the way the - # OpenAI API does: a client pointed at a model this server does - # not serve should hear "no such model", not a complaint about - # messages it would have sent correctly to the right server. + def test_unknown_model_without_a_registry_is_served(self): + # A single-container server has no registry to be strict about: + # its id defaults to the container's file name, and clients send + # a fixed model name they cannot easily change — `serve --help`'s + # own example sends "waste". Any name is served by the loaded + # model, as it was before --models existed. + status, body = self.chat(model="waste") + self.assertEqual(status, 200) + self.assertEqual(body["model"], "test-model") + + def test_unknown_model_does_not_outrank_missing_messages(self): + # With no registry there is no model validation at all, so the + # request's own shape is the first thing refused. The ordering + # the OpenAI API uses — model before shape — is what having a + # registry buys; TestModelSwap holds that test. status, body = self.post("/v1/chat/completions", {"model": "m"}) - self.assertEqual(status, 404) - self.assertEqual(body["error"]["param"], "model") + self.assertEqual(status, 400) + self.assertEqual(body["error"]["param"], "messages") def test_empty_messages(self): status, body = self.post("/v1/chat/completions", {"messages": []}) @@ -897,6 +907,8 @@ def test_chat_log_names_the_model_a_409_refused(self): " [model=swap-a]", self.logs()) def test_chat_log_names_the_model_a_404_refused(self): + # This class starts with --models, so a foreign name is a 404 and + # the line names what was refused. status, _ = self.post("/v1/chat/completions", {"model": "nope", "messages": [{"role": "user", "content": "x"}]}) @@ -1000,6 +1012,21 @@ def test_generation_rejects_unknown_model(self): self.assertEqual(status, 404) self.assertEqual(body["error"]["type"], "not_found_error") + def test_unknown_model_outranks_missing_messages(self): + # The model is validated before the request's shape, the way the + # OpenAI API does: a client pointed at a model this server does + # not serve should hear "no such model", not a complaint about + # messages it would have sent correctly to the right server. + status, body = self.post("/v1/chat/completions", {"model": "m"}) + self.assertEqual(status, 404) + self.assertEqual(body["error"]["param"], "model") + + def test_completions_rejects_unknown_model(self): + status, body = self.post("/v1/completions", + {"model": "nope", "prompt": "hi"}) + self.assertEqual(status, 404) + self.assertEqual(body["error"]["param"], "model") + def test_generation_accepts_loaded_model_after_swap(self): status, _ = self.load("swap-a") self.assertEqual(status, 200) From 1364101f6a2549d76a81508e983b9e2e5cfe3008 Mon Sep 17 00:00:00 2001 From: Hendrik Date: Tue, 22 Sep 2026 15:41:05 +0200 Subject: [PATCH 5/7] serve: loaded means residency, as the load response already said it did With --keep-previous a model the server still holds an open waste_ctx for was listed loaded:false on GET /v1/models and /v1/models/{id} - indistinguishable from a container never opened - while POST /v1/models/load counted it in its resident set and check_model_request called it 'not loaded'. The flag is residency: mid in engines, the same set the load response reports. The 409 for a resident-but-idle model names residency too; the model_not_loaded type and status are unchanged. The waste shape still travels only on the current entry - the per-model facts move with the current slot and are re-derived when a swap makes a resident model current again. --- CHANGELOG.md | 17 +++++++++++ docs/SERVE.md | 2 +- serve/server.py | 26 +++++++++++++---- tests/serve/test_server.py | 58 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ee9f659ba..8c6e72513 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,23 @@ changed. Each entry names the section to read for the numbers behind it. ### Fixed +- **`"loaded"` on `GET /v1/models` now means residency, as + `POST /v1/models/load` already said it did.** The registry listing set + `loaded: true` only for the model being served, so under + `--keep-previous` a model the server still holds an open `waste_ctx` + for was listed `loaded: false` — indistinguishable from a container + that was never opened — while the load response's `models` field + counted it as resident, and a generation naming it was refused as "not + loaded". The two endpoints now agree: `loaded` is `mid in engines`, + the same set the load response reports, and `GET /v1/models/{id}` says + the same for one entry. The 409 for a model that is resident but not + current says "resident but not the model being served" rather than + "registered but not loaded"; the `model_not_loaded` type and status + are unchanged. The `waste` shape still travels only on the current + entry — the per-model facts move with the current slot and are + re-derived when a swap makes a resident model current again. Tests in + `tests/serve/test_server.py`. + - **A single-container server serves any model name again.** The model registry made `check_model_request` strict for every deployment: with no `--models`, the registry holds only the loaded model, so a request diff --git a/docs/SERVE.md b/docs/SERVE.md index ef0af5144..3796be8aa 100644 --- a/docs/SERVE.md +++ b/docs/SERVE.md @@ -268,7 +268,7 @@ transcribed in this repo. That is what is left of | endpoint | notes | |---|---| | `GET /health` | liveness; never requires the API key | -| `GET /v1/models`, `GET /v1/models/{id}` | the registry: the loaded model (its real shape under a `waste` key), plus any registered-but-not-loaded containers with `"loaded": false` | +| `GET /v1/models`, `GET /v1/models/{id}` | the registry: `loaded` is residency — with `--keep-previous`, a model the server has swapped away from answers `loaded: true` too. The model being served is listed first and is the only entry with its real shape under a `waste` key | | `POST /v1/models/load` | swap models; see "Swapping models" below | | `POST /v1/chat/completions` | streaming and not, tools, images | | `POST /v1/completions` | raw continuation, no chat template | diff --git a/serve/server.py b/serve/server.py index 7d3014937..6d2fc1e5f 100644 --- a/serve/server.py +++ b/serve/server.py @@ -396,6 +396,11 @@ def check_model_request(self, body: dict) -> None: if mid not in self.registry: raise api.APIError(f"no such model: {mid}", status=404, type="not_found_error", param="model") + if mid in self.engines: + raise api.APIError( + f"model {mid} is resident but not the model being served; " + f"POST /v1/models/load to switch to it", + status=409, type="model_not_loaded", param="model") raise api.APIError( f"model {mid} is registered but not loaded; POST /v1/models/load " f"to switch to it", status=409, type="model_not_loaded", @@ -763,12 +768,19 @@ def _health(self): def _models(self): srv = self.server # The whole registry, current model first, so a client scanning the - # list sees what is resident before what is only available. A - # registered-but-not-loaded entry carries no `waste` shape: its - # per-container facts are unknown until it is opened. + # list sees what is resident before what is only available. `loaded` + # is residency — the same set the load response reports: a model + # `--keep-previous` still holds a waste_ctx for is loaded, whether + # or not it is the one being served. Only the current entry carries + # a `waste` shape: the per-model facts live on the current slot and + # move with it, and a resident-but-idle engine's shape is re-derived + # when a swap makes it current again. A registered-but-never-opened + # container is loaded=false and carries no shape — its facts are + # unknown until it is opened. data = [api.model_object(srv.model_id, srv.started, srv.model_info, loaded=True)] - data += [api.model_object(mid, srv.started, None, loaded=False) + data += [api.model_object(mid, srv.started, None, + loaded=mid in srv.engines) for mid in sorted(srv.registry) if mid != srv.model_id] self._send_json(200, {"object": "list", "data": data}) @@ -782,8 +794,10 @@ def _model(self, model_id: str): if model_id not in srv.registry: raise api.APIError(f"no such model: {model_id}", status=404, type="not_found_error", param="model") - self._send_json(200, api.model_object(model_id, srv.started, - None, loaded=False)) + # Residency, as in _models — with --keep-previous a model the + # server has swapped away from is still open and answers loaded. + self._send_json(200, api.model_object(model_id, srv.started, None, + loaded=model_id in srv.engines)) def _load_model(self): """POST /v1/models/load — swap the model this server serves. diff --git a/tests/serve/test_server.py b/tests/serve/test_server.py index c5bec5d8f..19038557c 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -1002,6 +1002,19 @@ def test_registered_model_entry_says_not_loaded(self): self.assertEqual(body["id"], "swap-b") self.assertEqual(body["loaded"], False) + def test_swapped_away_model_entry_says_not_loaded(self): + """Without --keep-previous a swap closes the previous engine and + drops it from `engines` — its entry is a container that is no + longer resident, the same as one never opened.""" + self.load("swap-a") + status, body = self.get("/v1/models/test-model") + self.assertEqual(status, 200) + self.assertEqual(body["loaded"], False) + status, body = self.get("/v1/models") + by_id = {m["id"]: m for m in body["data"]} + self.assertEqual(by_id["test-model"]["loaded"], False) + self.assertEqual(by_id["swap-a"]["loaded"], True) + def test_generation_rejects_registered_but_not_loaded(self): status, body = self.chat(model="swap-b") self.assertEqual(status, 409) @@ -1133,6 +1146,51 @@ def test_previous_stays_open(self): self.assertEqual(sorted(self.server.engines), ["swap-a", "test-model"]) + def test_models_reports_resident_models_as_loaded(self): + """`loaded` is residency, the same thing the load response reports: + a model --keep-previous still holds a waste_ctx for answers + loaded=True even after the swap made another model current. It was + listed loaded=false, indistinguishable from a container that was + never opened, while /v1/models/load counted it as resident.""" + self.load("swap-a") + status, body = self.get("/v1/models") + self.assertEqual(status, 200) + by_id = {m["id"]: m for m in body["data"]} + self.assertEqual(by_id["test-model"]["loaded"], True) + self.assertEqual(by_id["swap-a"]["loaded"], True) + self.assertEqual(by_id["swap-b"]["loaded"], False) + self.assertNotIn("waste", by_id["test-model"]) # resident, not current: + # the shape moves with the current slot and is re-derived on the + # slot move that would make it serve again. + self.assertIn("waste", by_id["swap-a"]) + # The current model still leads the list. + self.assertEqual(body["data"][0]["id"], "swap-a") + + def test_swapped_away_model_entry_says_not_loaded(self): + """Overridden: with --keep-previous the swapped-away model is not + out of the resident set — its entry says loaded, as + test_resident_model_entry_says_loaded asserts per id.""" + status, body = self.get("/v1/models") + by_id = {m["id"]: m for m in body["data"]} + self.assertEqual(by_id["test-model"]["loaded"], True) + + def test_resident_model_entry_says_loaded(self): + self.load("swap-a") + status, body = self.get("/v1/models/test-model") + self.assertEqual(status, 200) + self.assertEqual(body["id"], "test-model") + self.assertEqual(body["loaded"], True) + + def test_generation_naming_a_resident_model_says_so(self): + """The 409 for a model that is resident but not current does not + call it \"not loaded\" — that was true only before the swap and is + the same lie the registry listing told.""" + self.load("swap-a") + status, body = self.chat(model="test-model") + self.assertEqual(status, 409) + self.assertEqual(body["error"]["type"], "model_not_loaded") + self.assertIn("resident", body["error"]["message"]) + def test_can_switch_back_without_reopening(self): """swap-a was never closed, so loading it again is a slot move, not a new open.""" From e0b733810716fffaf49f97f3d04b673b5e569e76 Mon Sep 17 00:00:00 2001 From: Hendrik Date: Tue, 22 Sep 2026 16:04:33 +0200 Subject: [PATCH 6/7] Refusing --usage together with --models Refusing --usage when --models is provided rejects this invalid configuration upfront, explaining clearly that --usage is container-specific and cannot be shared across multiple models. At the same time, removing "usage_path": args.usage from engine_kwargs in serve/__main__.py ensures that swap operations always let each container find its own hotlist (the default None behavior), preventing any accidental propagation of a single hotlist to swappable models. --- serve/__main__.py | 17 ++++++++++++----- tests/serve/test_main.py | 20 ++++++++++++++++++-- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/serve/__main__.py b/serve/__main__.py index f97211549..1713bf883 100644 --- a/serve/__main__.py +++ b/serve/__main__.py @@ -286,11 +286,13 @@ def main(argv=None) -> int: "contexts need together is the sum of their " "budgets, and on the machines this engine targets " "that is the difference between working and " - "paging. Both models can then answer at once — " - "each waste_ctx takes one caller, so each has its " - "own lock. Every model switched to stays resident, " - "so the load that would put the set over the " - "machine's RAM is refused with 507") + "paging. Generation still serves only the current " + "model — each waste_ctx takes one caller — but " + "switching back to a resident one is a slot move " + "instead of a reopen of a multi-gigabyte " + "container. Every model switched to stays " + "resident, so the load that would put the set " + "over the machine's RAM is refused with 507") s.add_argument("--plan", action="store_true", help="print the memory plan and exit without loading") s.add_argument("--no-log-requests", action="store_true", @@ -301,6 +303,11 @@ def main(argv=None) -> int: args = ap.parse_args(argv) + if args.usage and args.models: + print("--usage cannot be used with --models: a learned hotlist is " + "specific to one container", file=sys.stderr) + return 2 + model = Path(args.model).expanduser() if not model.exists(): print(f"no such container: {model}", file=sys.stderr) diff --git a/tests/serve/test_main.py b/tests/serve/test_main.py index cde2c52b3..842e2d5e4 100644 --- a/tests/serve/test_main.py +++ b/tests/serve/test_main.py @@ -17,20 +17,36 @@ python3 tests/serve/test_main.py """ - +import io import sys import unittest from pathlib import Path +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from serve.__main__ import (RegistryBudgetError, check_registry_budget, - describe_registry, human) + describe_registry, human, main) from serve.engine import EngineError GB = 1 << 30 MODELS = {"glm53": "/fake/glm53.waste", "ds41": "/fake/ds41.waste"} +class TestHumanBytes(unittest.TestCase): + def test_units(self): + self.assertEqual(human(0), "0 B") + self.assertEqual(human(1024), "1.0 KB") + self.assertEqual(human(3 * GB), "3.0 GB") + self.assertEqual(human(2 * (1 << 40)), "2.0 TB") + + +class TestUsageWithModels(unittest.TestCase): + def test_usage_with_models_is_refused(self): + stderr = io.StringIO() + with patch("sys.stderr", stderr): + rc = main(["dummy.waste", "--models", "other.waste", "--usage", "hotlist.waste"]) + self.assertEqual(rc, 2) + self.assertIn("--usage cannot be used with --models", stderr.getvalue()) class TestRegistryBudget(unittest.TestCase): """--models is refused unless --budget shows that 2 x budget fits.""" From 1d2e28730979047e53d6d4d987bf65e1d4c8dcdf Mon Sep 17 00:00:00 2001 From: Hendrik Date: Tue, 22 Sep 2026 16:31:24 +0200 Subject: [PATCH 7/7] serve: Logs a line on stderr per swap To ensure model swaps report container capabilities and any chat formatting limitations to the operator, we log a line to `sys.stderr` on each swap in `ChatServer.load_model`, reporting the model switched to and its `chat_error` when set. --- serve/server.py | 42 ++++++++++++++++++++++---------------- tests/serve/test_server.py | 14 +++++++++++++ 2 files changed, 38 insertions(+), 18 deletions(-) diff --git a/serve/server.py b/serve/server.py index 6d2fc1e5f..19962d3b5 100644 --- a/serve/server.py +++ b/serve/server.py @@ -572,24 +572,30 @@ def load_model(self, model_id: str) -> Optional[str]: with previous_engine.lock: if resident is not None: self._detect(resident, model_id) - return previous_slot.model_id - # Room for it, before a byte of it is allocated — and - # before anything is closed, so a refusal here is - # indistinguishable from never having been asked: same - # current model, same resident set, same open containers. - self.check_room(model_id, path) - try: - engine = self.engine_factory(path) - except EngineError as e: - raise ModelLoadError( - f"could not load {model_id}: {e}") from e - self._detect(engine, model_id) - self.engines[model_id] = engine - if self.keep_previous: - return previous_slot.model_id - self.engines.pop(previous_slot.model_id) - previous_engine.close() - return previous_slot.model_id + prev = previous_slot.model_id + else: + # Room for it, before a byte of it is allocated — and + # before anything is closed, so a refusal here is + # indistinguishable from never having been asked: same + # current model, same resident set, same open containers. + self.check_room(model_id, path) + try: + engine = self.engine_factory(path) + except EngineError as e: + raise ModelLoadError( + f"could not load {model_id}: {e}") from e + self._detect(engine, model_id) + self.engines[model_id] = engine + if not self.keep_previous: + self.engines.pop(previous_slot.model_id) + previous_engine.close() + prev = previous_slot.model_id + err = self._slot.chat_error + line = f"swap: {model_id}" + if err: + line += f" chat_error: {err}" + sys.stderr.write(f"{line}\n") + return prev def _current(self) -> tuple: """(the current ModelSlot, its engine). Call with _slot_lock held.""" diff --git a/tests/serve/test_server.py b/tests/serve/test_server.py index 19038557c..c1518b9c8 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -1866,3 +1866,17 @@ def test_invalid_glm_tool_definitions_are_400(self): status, body = self.chat(tools=[tool]) self.assertEqual(status, 400) self.assertTrue(body["error"]["param"].startswith("tools[0]")) + + def test_load_log_names_the_model_loaded(self): + status, _ = self.post("/v1/models/load", {"model": "swap-a"}) + self.assertEqual(status, 200) + self.assertIn('"POST /v1/models/load HTTP/1.1" 200 -' + " [model=swap-a]", self.logs()) + + def test_load_log_stderr_reports_swap_and_chat_error_if_set(self): + status, _ = self.post("/v1/models/load", {"model": "swap-a"}) + self.assertEqual(status, 200) + self.assertIn("swap: swap-a", self.logs()) + + def test_get_log_carries_no_model(self): + pass