diff --git a/CHANGELOG.md b/CHANGELOG.md index cd11ede41..d21e00e00 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -97,6 +97,38 @@ No ABI move: `src/waste.h` changes only its version macros. ### 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 + 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`. + - **DeepSeek-V4.1 without `--reclaim`: the Engram index was never written** (#71, #72 — found and fixed by @helenkwok converting the release on an EPYC 7713). Only the `--reclaim` path called @@ -125,6 +157,48 @@ No ABI move: `src/waste.h` changes only its version macros. ### Measured and not adopted +- **`--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 + this one may not skip". Now neither may. `CI_K2_ORACLE_STRICT=1` turns + every skip path — missing template, missing jinja2, unresolved markers — + into a failure, and the job asserts all **seven** checks ran, so it + cannot be green by having done nothing. Vendoring the template is what + made it possible: no download, no weights, no machine-local `~/models`. - **The i8mm trunk kernel as DeepSeek-V4.1's default.** 3.87–3.97 tok/s becomes 4.52–4.56. Of 2,080 routing decisions it flips one, a tie the reference itself cannot resolve, and the 201 others that differ all diff --git a/docs/SERVE.md b/docs/SERVE.md index e5f478195..3796be8aa 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: `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 | @@ -318,6 +319,91 @@ 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 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 + + 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, 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 + 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 — 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 @@ -380,6 +466,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 @@ -429,7 +529,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; @@ -444,5 +546,15 @@ 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). 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; 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 061355db5..1713bf883 100644 --- a/serve/__main__.py +++ b/serve/__main__.py @@ -25,18 +25,137 @@ 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) -from .server import serve # noqa: E402 + plan_memory, usable_ram) +from .server import ModelLoadError, serve # noqa: E402 POLICIES = {"lfru": CACHE_LFRU, "lru": CACHE_LRU} -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" +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 + + +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: @@ -82,6 +201,16 @@ 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 --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 — 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", @@ -95,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)") @@ -140,17 +271,73 @@ 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. 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 " + "contexts need together is the sum of their " + "budgets, and on the machines this engine targets " + "that is the difference between working and " + "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", + 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) + 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) 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) @@ -167,6 +354,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( @@ -217,7 +408,23 @@ 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=registry, + keep_previous=args.keep_previous, + usable_ram=usable, + 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) @@ -257,6 +464,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'})") @@ -272,7 +493,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..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 -------------------------------------------------- @@ -558,9 +574,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 1f83ac849..607bcbf1f 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,11 +41,12 @@ import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from typing import Optional +from typing import Callable, NamedTuple, Optional 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" @@ -53,6 +57,63 @@ 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 ModelSlot(NamedTuple): + """Everything one container makes true, published as one object. + + 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. + """ + + 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 ------------------------------- + # + # 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 factory, + # run again on every load, and it keeps the old comment because every + # word of it still holds. + + @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. + """ + default_thinking = start_thinking class ChatServer(ThreadingHTTPServer): """Threaded HTTP, one engine, one lock.""" @@ -85,9 +146,9 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, self._tmp = tmpdir or tempfile.mkdtemp(prefix="waste-serve-") self.tmpdir = self._tmp 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. # @@ -108,16 +169,15 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, # 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 @@ -125,21 +185,21 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, # 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: @@ -147,18 +207,19 @@ def __init__(self, addr, handler, *, engine: Engine, model_id: str, # 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. @@ -185,6 +246,406 @@ def new_parser(self, thinking: bool, tools=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, + 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() + # 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 ---------------------------------------------- + + 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. 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") + 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", + 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). + + 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 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. + 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].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_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. + resident = self.engines.get(model_id) + with previous_engine.lock: + if resident is not None: + self._detect(resident, 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.""" + slot = self._slot + return slot, slot.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 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. @@ -207,7 +668,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() @@ -272,6 +747,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": @@ -291,6 +767,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(): @@ -301,6 +778,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") @@ -324,37 +803,92 @@ 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. `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=mid in srv.engines) + 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)) + # 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. + + 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 ----------------------------------------------------------- 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 + # 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. # 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", @@ -377,39 +911,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} @@ -457,7 +999,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): @@ -465,23 +1007,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") @@ -535,8 +1077,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 @@ -571,7 +1113,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") @@ -590,6 +1132,14 @@ 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) prompt_text = body.get("prompt") if isinstance(prompt_text, list): if len(prompt_text) != 1 or not isinstance(prompt_text[0], str): @@ -599,15 +1149,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) @@ -622,11 +1173,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: @@ -637,7 +1188,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)), @@ -648,8 +1199,27 @@ 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, + 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 + 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). + + 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. if ":" in host: @@ -658,7 +1228,11 @@ 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, + usable_ram=usable_ram, memory_plan=memory_plan) 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_main.py b/tests/serve/test_main.py new file mode 100644 index 000000000..842e2d5e4 --- /dev/null +++ b/tests/serve/test_main.py @@ -0,0 +1,160 @@ +# 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 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, 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.""" + + 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 cd1441481..c1518b9c8 100644 --- a/tests/serve/test_server.py +++ b/tests/serve/test_server.py @@ -17,7 +17,10 @@ """ import json +import contextlib +import dataclasses import http.client +import io import shutil import socket import sys @@ -27,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)) @@ -35,6 +39,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 +48,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,6 +411,25 @@ def test_client_disconnect_stops_generation(self): class TestValidation(ServerTestCase): def test_missing_messages(self): + status, body = self.post("/v1/chat/completions", {}) + self.assertEqual(status, 400) + self.assertEqual(body["error"]["param"], "messages") + + 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, 400) self.assertEqual(body["error"]["param"], "messages") @@ -819,6 +845,730 @@ 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): + # 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"}]}) + 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_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) + 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_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) + 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_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.""" + 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 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 +# 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.""" @@ -1116,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