"""HTTP server. POST /decide {"context": str, "schema": {...}} -> typed JSON decisions (all questions packed in one row) schema: {question: {"type": "choice", "options": [...]} | {"type": "bool"} | {"type": "scale", "legend": [...]}}; a missing or malformed schema is a 422 naming this form (decider.infer.Decider._check_schema) POST /v1/systemone {"state": str|object|array, "questions": {id: {...}}} -> the TypeSafe/Jev wire format (decider.systemone): Choice (up to 255 described options), Score, Noul; every question is scored in its own row, so answers are independent of each other ("independent": false packs them behind one copy of the state instead). GET /v1/models /health /stats uvicorn decider.serve:app --host 0.0.0.0 --port 8000 (env: DECIDER_MODEL and the variables below) Execution (1.1.1; docs/SERVING.md has the design and the measurements): * decider.engine_v2.EngineV2: CUDA graphs keyed on (batch bucket, padded length bucket) only. The whole grid is captured during start-up and the engine is then sealed, so no request pays a graph capture or a torch.compile. Rows longer than the last length bucket (8,192 tokens) run eager, in chunks bounded by DECIDER_GRAPH_TOKEN_BUDGET padded tokens. * one GPU thread runs every forward; tokenisation runs on a CPU pool; rows waiting at the same moment are partitioned by decider.batching.plan_batches, which pads a shorter row into a longer row's bucket when that costs less than a second forward (cost model: DECIDER_MERGE_OVERHEAD_TOKENS + rows * padded length). When a collection held rows from more than one live request the next one waits up to DECIDER_BATCH_ADAPTIVE_WAIT_MS for more rows, unless its own first row took more than 50 ms to arrive (a quiet queue drops the window again); after a single-request collection it does not wait. * an independent request with more than one question over a state of at least DECIDER_SHARED_MIN_TOKENS tokens runs the state once and forks its cache per question (EngineV2.score_shared, decider.shared_prefix). The fork is made in chunks that fit DECIDER_SHARED_FORK_GB, so the peak memory does not grow with the question count. The cuDNN SDPA backend is off (decider.engine.set_attention_backend_policy), which is what makes that path agree with the full forward. * the schema cache (questions-first layout, prefix cached per question set) is honoured exactly as before: on when decider_config.json has "schema_first": true, or with DECIDER_SCHEMA_CACHE=1 on a model trained for that layout. It is the one path that does GPU work after start-up that is not a graph replay: the first request with a new schema runs its prefix and captures one graph per (schema, batch bucket, state bucket) as in 1.0.x (/stats -> schema_cache.captures). Its requests are bounded and admitted like the others, on prefix plus suffix tokens, before any GPU preparation. * prompt layout (1.2.0): a model whose decider_config.json has "layout": "chat" (a chat-trained checkpoint) is read in the chat layout (decider.prompt.build_chat: the tokenizer's chat template around one user turn, thinking off, "Answer: (" in the assistant turn) on every route, including the shared-prefix fork and the schema cache. Every other model is read in the plain layout, with the same token ids as 1.1.x. A config naming an unknown layout stops start-up with a ValueError. * requests are bounded before they reach the GPU: HTTP 413 when a row, the request or the expanded question count exceeds the limits below, HTTP 503 when the outstanding work exceeds DECIDER_MAX_QUEUE_ROWS. Variables (default): DECIDER_DEVICE (auto: cuda, else mps, else cpu) DECIDER_COMPILE (0) DECIDER_FP8 (0) DECIDER_SHARED (1) DECIDER_SHARED_MIN_TOKENS (768) DECIDER_SHARED_FORK_GB (8) DECIDER_MERGE_OVERHEAD_TOKENS (512) DECIDER_BATCH_ADAPTIVE_WAIT_MS (2) DECIDER_MAX_BATCH (32) DECIDER_BATCH_WAIT_MS (0; DECIDER_MAX_WAIT_MS is an alias) DECIDER_MAX_STATE_TOKENS (32768) DECIDER_T_BUCKETS DECIDER_B_BUCKETS DECIDER_GRAPH_TOKEN_BUDGET (32768) DECIDER_WARMUP (1) DECIDER_TOKENIZE_THREADS (8) DECIDER_MAX_ROWS (1024) DECIDER_MAX_ROW_TOKENS (DECIDER_MAX_STATE_TOKENS + 4096) DECIDER_MAX_REQUEST_TOKENS (1048576) DECIDER_MAX_QUEUE_ROWS (4096) DECIDER_TEMPERATURE DECIDER_SCHEMA_CACHE (0) DECIDER_SCHEMA_MIN_SEEN (2) DECIDER_SCHEMAS Temperatures (1.4.0, decider.temperature): decider_config.json "temperature", and optionally "temperature_by_type" {"choice": T, "noul": T, "score": T} (a /decide "bool" field is "noul", a "scale" field is "score"; a missing type uses "temperature"). Every row carries the answer type of each of its slots, so a batch that mixes requests and types still applies each answer's own temperature. DECIDER_TEMPERATURE replaces "temperature" and switches the by-type map off. /health and the ready line report the temperature every answer type gets. """ import asyncio, json, os, time from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException from pydantic import BaseModel from decider import systemone as S1 from decider import temperature as TT from decider.batching import DEFAULT_MERGE_OVERHEAD_TOKENS, plan_batches from decider.prompt import build, MAX_OPTIONS, resolve_layout, chat_template from decider.prompt_fast import build_rows, unique_tokens def _env_int(name, default): return int(os.environ.get(name, default)) MODEL = os.environ.get("DECIDER_MODEL", "runs/r3_v2/model") MAX_BATCH = _env_int("DECIDER_MAX_BATCH", 32) BATCH_WAIT_MS = float(os.environ.get("DECIDER_BATCH_WAIT_MS", os.environ.get("DECIDER_MAX_WAIT_MS", "0"))) ADAPTIVE_WAIT_MS = float(os.environ.get("DECIDER_BATCH_ADAPTIVE_WAIT_MS", "2")) # window after a collection that held >1 request; 0 disables ADAPTIVE_IDLE_RESET_MS = 50.0 # an idler queue than this drops the adaptive window again MERGE_OVERHEAD_TOKENS = _env_int("DECIDER_MERGE_OVERHEAD_TOKENS", DEFAULT_MERGE_OVERHEAD_TOKENS) MAX_STATE_TOKENS = _env_int("DECIDER_MAX_STATE_TOKENS", 32768) SHARED = os.environ.get("DECIDER_SHARED", "1") == "1" SHARED_MIN_TOKENS = _env_int("DECIDER_SHARED_MIN_TOKENS", 768) DEVICE = os.environ.get("DECIDER_DEVICE", "auto") # auto | cuda[:i] | mps | cpu COMPILE = os.environ.get("DECIDER_COMPILE", "0") == "1" FP8 = os.environ.get("DECIDER_FP8", "0") == "1" WARMUP = os.environ.get("DECIDER_WARMUP", "1") == "1" TOKENIZE_THREADS = _env_int("DECIDER_TOKENIZE_THREADS", 8) GRAPH_TOKEN_BUDGET = _env_int("DECIDER_GRAPH_TOKEN_BUDGET", 32768) # request bounds, checked after tokenisation and before anything is queued MAX_ROWS = _env_int("DECIDER_MAX_ROWS", 1024) # scoring rows per request (questions, expanded score levels) MAX_ROW_TOKENS = _env_int("DECIDER_MAX_ROW_TOKENS", MAX_STATE_TOKENS + 4096) # tokens in one row: truncated state + question block MAX_REQUEST_TOKENS = _env_int("DECIDER_MAX_REQUEST_TOKENS", 1 << 20) # sum of row lengths of one request MAX_QUEUE_ROWS = _env_int("DECIDER_MAX_QUEUE_ROWS", 4096) # rows admitted and not yet scored, over all requests DECIDE_MAX_CTX_TOKENS = 1536 # /decide context cap, unchanged from 1.0.x MODEL_NAME = "decider"; TEMP = 1.0; TEMP_SCHEMA = 1.0; TEMP_BY_TYPE = {}; TEMP_SCHEMA_BY_TYPE = {}; RELEASE_DATE = "2026-09-17"; ISOLATED = False; NEUTRALIZE_NONE = True SCHEMA_FIRST = False; LAYOUT = "plain"; CHAT = None; se = None; squeue = None; schemas = {}; seen = {} eng = None; queue = None; gpu = None; cpu = None; batcher_task = None; schema_task = None outstanding = 0; REQ_SEQ = 0 stats = dict(requests=0, batches=0, decisions=0, rows=0, shared_prefix_requests=0, errors=0, rejected_too_large=0, rejected_overloaded=0, batch_hist={}, bucket_hist={}) def _ints(name): v = os.environ.get(name) return [int(x) for x in v.split(",")] if v else None def load_config(path): """decider_config.json from a model folder or a Hub repository (the way decider.infer.Decider resolves it).""" try: if os.path.isdir(path): return json.load(open(os.path.join(path, "decider_config.json"))) from huggingface_hub import hf_hub_download return json.load(open(hf_hub_download(path, "decider_config.json"))) except Exception: return {} # ---- request preparation (CPU) ------------------------------------------- class _NoShuffle: def shuffle(self, x): pass def sample(self, xs, k): return xs[:k] def prepare(tok, state, questions, independent, isolated=False, max_state_tokens=32768, chat=None): """Render, plan the rows, tokenize the state once. -> (rqs, index, items, ctx_len). The rows are those `prompt.build` produces for the same request (tests/test_prompt_fast.py, tests/test_serve_prepare.py); chat: the ChatTemplate of a chat-layout model, None for the plain layout.""" ctx = S1.render_state(state) rqs = {k: S1.render_question(v) for k, v in questions.items()} flat, index = S1.plan_rows(rqs, isolated and independent) pairs = [(r["question"], list(r["options"])) for r in flat] rows = [[p] for p in pairs] if independent else [pairs] items, ctx_len = build_rows(tok, ctx, rows, max_ctx_tokens=max_state_tokens, chat=chat) types = S1.row_types(rqs, index) # answer type per slot (decider.temperature) for it, ts in zip(items, [[t] for t in types] if independent else [types]): it["types"] = ts return rqs, index, items, ctx_len def _prepare_s1(state, questions, independent): return prepare(eng.tok, state, questions, independent, ISOLATED, MAX_STATE_TOKENS, chat=CHAT) def _prepare_decide(context, schema): from decider.infer import Decider, Example, Q, neutralize_options qs = Decider._schema_to_questions(schema) for q in qs: if NEUTRALIZE_NONE: q["options"], q["_back"] = neutralize_options(q["options"]) ex = Example(context, [Q(q["question"], list(q["options"]), 0) for q in qs]) it = build(ex, eng.tok, _NoShuffle(), max_options=MAX_OPTIONS, max_ctx_tokens=DECIDE_MAX_CTX_TOKENS, chat=CHAT) it["types"] = [q["_type"] for q in qs] # bool -> noul, scale -> score (decider.temperature) return qs, it def _format_decide(schema, qs, probs): o = {} for (qtext, spec), q, p in zip(schema.items(), qs, probs): p = p[:len(q["options"])].tolist(); t = spec.get("type", "choice"); j = max(range(len(p)), key=p.__getitem__) back = q.get("_back", {}); names = [back.get(x, x) for x in q["options"]] if t == "bool": o[qtext] = {"noul": round(p[1], 4), "type": "noul"} elif t == "choice": o[qtext] = {"choice": names[j], "confidence": round(p[j], 4), "type": "choice", "probabilities": {k: round(v, 4) for k, v in zip(names, p)}} else: keys = q["_keys"]; score = sum(float(k) * pi for k, pi in zip(keys, p)) o[qtext] = {"score": round(score, 2), "confidence": round(p[j], 4), "type": "scale", "legend": q["_legend"], "probabilities": {str(keys[i]): round(pi, 4) for i, pi in enumerate(p)}} return o # ---- limits --------------------------------------------------------------- def check_size(row_lengths, max_rows=None, max_row_tokens=None, max_request_tokens=None): """Raise HTTPException(413) when a request exceeds the row count, per-row token or total token limit.""" max_rows = MAX_ROWS if max_rows is None else max_rows max_row_tokens = MAX_ROW_TOKENS if max_row_tokens is None else max_row_tokens max_request_tokens = MAX_REQUEST_TOKENS if max_request_tokens is None else max_request_tokens n, total, longest = len(row_lengths), sum(row_lengths), max(row_lengths, default=0) if n > max_rows: msg = f"too many questions: the request expands to {n} scoring rows, the limit is {max_rows} (DECIDER_MAX_ROWS)" elif longest > max_row_tokens: msg = f"too many tokens: one row has {longest} tokens, the limit is {max_row_tokens} per row (DECIDER_MAX_ROW_TOKENS)" elif total > max_request_tokens: msg = f"too many tokens: the request has {total} tokens over {n} rows, the limit is {max_request_tokens} (DECIDER_MAX_REQUEST_TOKENS)" else: return stats["rejected_too_large"] += 1 raise HTTPException(413, msg) def _admit(n): """Reserve n rows of outstanding work or raise HTTPException(503).""" global outstanding if outstanding + n > MAX_QUEUE_ROWS: stats["rejected_overloaded"] += 1 raise HTTPException(503, f"server busy: {outstanding} rows queued, the limit is {MAX_QUEUE_ROWS} (DECIDER_MAX_QUEUE_ROWS); retry later") outstanding += n def _release(n): global outstanding outstanding -= n # ---- GPU work ---------------------------------------------------------------- def _score_items(items): return eng.score_items(items, temperature=TT.for_items(TEMP, TEMP_BY_TYPE, items)) # the scalar TEMP without a by-type map def _score_shared(items): return eng.score_shared(items, temperature=TT.for_items(TEMP, TEMP_BY_TYPE, items)) async def _collect(q, wait_ms=None, adaptive_ms=0.0, idle_reset_ms=None): """Take what is already queued and go. `wait_ms` (DECIDER_BATCH_WAIT_MS) is the unconditional collection window. `adaptive_ms` is the extra window the batcher asks for after a collection that held more than one live request; it is dropped again when the first row of this collection took longer than `idle_reset_ms` to arrive, so an isolated request after a quiet period never waits. """ wait = BATCH_WAIT_MS if wait_ms is None else wait_ms reset = ADAPTIVE_IDLE_RESET_MS if idle_reset_ms is None else idle_reset_ms t0 = time.monotonic() batch = [await q.get()] if adaptive_ms > 0 and (time.monotonic() - t0) * 1000 <= reset: wait = max(wait, adaptive_ms) deadline = time.monotonic() + wait / 1000 while len(batch) < MAX_BATCH: try: batch.append(q.get_nowait()) except asyncio.QueueEmpty: timeout = deadline - time.monotonic() if timeout <= 0: break try: batch.append(await asyncio.wait_for(q.get(), timeout)) except asyncio.TimeoutError: break return batch def _bucketed(n): """False for a row longer than the engine's last captured length bucket. Those run eager at a request-specific shape, so they are grouped by exact padded length as in 1.1.0 instead of being padded into another row's bucket.""" t_bucket = getattr(eng, "t_bucket", None) return True if t_bucket is None else t_bucket(n) is not None def adaptive_ms(batch): """The extra collection window the next collection may use: DECIDER_BATCH_ADAPTIVE_WAIT_MS when this collection held live rows from more than one request, 0 otherwise. Rows whose future is already done (a cancelled or failed request) are not counted: they are not evidence that requests are overlapping.""" if ADAPTIVE_WAIT_MS <= 0: return 0.0 return ADAPTIVE_WAIT_MS if len({rid for fut, _, rid in batch if not fut.done()}) > 1 else 0.0 async def batcher(): """One forward per planned group. Rows queued at the same moment are partitioned by decider.batching.plan_batches: a shorter row is padded into a longer row's bucket when that costs less than running a second forward.""" loop = asyncio.get_running_loop() extra = 0.0 # the adaptive window the previous collection earned while True: batch = await _collect(queue, BATCH_WAIT_MS, extra) extra = adaptive_ms(batch) groups = plan_batches([len(it["ids"]) for _, it, _ in batch], eng.pad_len, eng.max_rows, MAX_BATCH, MERGE_OVERHEAD_TOKENS, _bucketed) for T, idx in groups: part = [batch[i] for i in idx] stats["bucket_hist"][T] = stats["bucket_hist"].get(T, 0) + len(part) try: probs = await loop.run_in_executor(gpu, _score_items, [it for _, it, _ in part]) for (fut, _, _), p in zip(part, probs): if not fut.done(): fut.set_result(p) except Exception as e: for fut, _, _ in part: if not fut.done(): fut.set_exception(e) stats["batches"] += 1 stats["batch_hist"][len(part)] = stats["batch_hist"].get(len(part), 0) + 1 # ---- schema cache (questions-first layout; opt-in) -------------------------------- def _schema_key(questions, independent): return (json.dumps(questions, sort_keys=True, ensure_ascii=False), independent) class _SQ: def __init__(self, text, options): self.text, self.options = text, options def _plan_schema(questions, independent, state): """CPU part of a schema-cache request: prefix token lengths (from the cached handle when the schema is known, otherwise tokenised here, without touching the GPU) and the suffix row (context plus answer slots). -> (tps, (suffix ids, slots)). Raises ValueError for an invalid question.""" from decider.prompt import schema_prefix_ids, schema_suffix_ids rqs = {k: S1.render_question(v) for k, v in questions.items()}; rows, _ = S1.plan_rows(rqs, ISOLATED and independent) cached = schemas.get(_schema_key(questions, independent)) if cached is not None: tps = list(cached[1].tps) else: qs = [_SQ(x["question"], list(x["options"])) for x in rows] tps = [len(schema_prefix_ids(eng.tok, g, chat=CHAT)) for g in ([[q] for q in qs] if independent else [qs])] row = schema_suffix_ids(eng.tok, S1.render_state(state), 1 if independent else len(rows), MAX_STATE_TOKENS, chat=CHAT) return tps, row def _schema_handle(questions, independent, compile=False): """Compile (or look up) the question schema: its prefix is run once, requests then only run the state. GPU thread.""" key = _schema_key(questions, independent) if key not in schemas: rqs = {k: S1.render_question(v) for k, v in questions.items()}; rows, index = S1.plan_rows(rqs, ISOLATED and independent) if len(schemas) >= 128: old = next(iter(schemas)); hid = schemas.pop(old)[1].id for k in [k for k in se.graphs if k[0] == hid]: del se.graphs[k] h = se.prepare(rows, independent=independent, compile=compile) h.types = S1.row_types(rqs, index) # answer type per schema row (decider.temperature) schemas[key] = (rqs, h, index) return schemas[key] def _worth_caching(questions, independent): """A schema gets a cached prefix and CUDA graphs from its second request on.""" key = _schema_key(questions, independent) if key in schemas: return True if len(seen) > 50000: seen.clear() seen[key] = seen.get(key, 0) + 1 return seen[key] >= _env_int("DECIDER_SCHEMA_MIN_SEEN", 2) def _score_schema(h, rows): return se.score_rows(h, rows, temperature=TT.for_types(TEMP_SCHEMA, TEMP_SCHEMA_BY_TYPE, h.types)) async def schema_batcher(): loop = asyncio.get_running_loop() while True: batch = await _collect(squeue) groups = {} for fut, h, row in batch: groups.setdefault((h.id, se.bucket(len(row[0]))), (h, []))[1].append((fut, row)) for h, items in groups.values(): step = max(1, MAX_BATCH // h.P) for i in range(0, len(items), step): chunk = items[i:i + step] try: probs = await loop.run_in_executor(gpu, _score_schema, h, [c for _, c in chunk]) for (fut, _), p in zip(chunk, probs): if not fut.done(): fut.set_result(p) except Exception as e: for fut, _ in chunk: if not fut.done(): fut.set_exception(e) stats["schema_batches"] = stats.get("schema_batches", 0) + 1 # ---- start-up / shutdown ------------------------------------------------------- def apply_config(cfg): global MODEL_NAME, TEMP, TEMP_SCHEMA, TEMP_BY_TYPE, TEMP_SCHEMA_BY_TYPE, RELEASE_DATE, ISOLATED, NEUTRALIZE_NONE, SCHEMA_FIRST, LAYOUT LAYOUT = resolve_layout(cfg) # ValueError for an unknown layout, before the engine is built NEUTRALIZE_NONE = bool(cfg.get("neutralize_none", True)) MODEL_NAME = "decider-" + str(cfg.get("version", "dev")) # ValueError for a temperature that is not finite and > 0 or an unknown by-type key, before the engine is built. # DECIDER_TEMPERATURE replaces "temperature" and switches "temperature_by_type" off (decider.temperature). (TEMP, TEMP_BY_TYPE), (TEMP_SCHEMA, TEMP_SCHEMA_BY_TYPE) = TT.from_config(cfg, os.environ.get("DECIDER_TEMPERATURE")) RELEASE_DATE = str(cfg.get("release_date", RELEASE_DATE)) ISOLATED = bool(cfg.get("isolated_levels", False)) trained = bool(cfg.get("schema_first", False) or cfg.get("schema_first_trained", False)) SCHEMA_FIRST = trained and (bool(cfg.get("schema_first", False)) or os.environ.get("DECIDER_SCHEMA_CACHE", "0") == "1") def start_workers(loop=None): """Queues, executors and batcher tasks. Needs `eng` set; a test can set a stand-in engine and call this directly.""" global queue, gpu, cpu, batcher_task, squeue, schema_task loop = loop or asyncio.get_running_loop() if gpu is None: gpu = ThreadPoolExecutor(max_workers=1, thread_name_prefix="gpu") if cpu is None: cpu = ThreadPoolExecutor(max_workers=TOKENIZE_THREADS, thread_name_prefix="tok") queue = asyncio.Queue(); batcher_task = loop.create_task(batcher()) if SCHEMA_FIRST: squeue = asyncio.Queue(); schema_task = loop.create_task(schema_batcher()) def resolve_device(requested=None): """The device and dtype the server runs on. `auto` picks as decider.infer.Decider does: CUDA, else MPS, else CPU; float16 on MPS, bfloat16 elsewhere. An explicit device that is not available, or FP8 / torch.compile off CUDA, is a start-up error that says so, instead of the torch assertion a CUDA call raises on a build without CUDA.""" import torch req = (DEVICE if requested is None else requested).strip().lower() if req in ("", "auto"): dev = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") else: dev = req kind = dev.split(":")[0] if kind == "cuda" and not torch.cuda.is_available(): raise RuntimeError(f"DECIDER_DEVICE={req}, but torch.cuda.is_available() is False on this machine. " "Set DECIDER_DEVICE=mps or cpu (or leave it at auto).") if kind == "mps" and not torch.backends.mps.is_available(): raise RuntimeError(f"DECIDER_DEVICE={req}, but torch.backends.mps.is_available() is False on this machine.") if kind not in ("cuda", "mps", "cpu"): raise RuntimeError(f"DECIDER_DEVICE={req}: expected auto, cuda, cuda:, mps or cpu.") if not dev.startswith("cuda") and (FP8 or COMPILE): raise RuntimeError(f"DECIDER_FP8 and DECIDER_COMPILE need CUDA; the server is starting on {dev}. Unset them.") return dev, (torch.float16 if dev.startswith("mps") else torch.bfloat16) async def _start(): global eng, se, gpu, CHAT from decider.engine_v2 import EngineV2 apply_config(load_config(MODEL)) gpu = ThreadPoolExecutor(max_workers=1, thread_name_prefix="gpu") loop = asyncio.get_running_loop() dev, dtype = resolve_device() eng = await loop.run_in_executor(gpu, lambda: EngineV2( MODEL, device=dev, dtype=dtype, compile=COMPILE, fp8=FP8, max_ctx_tokens=MAX_STATE_TOKENS, t_buckets=_ints("DECIDER_T_BUCKETS"), b_buckets=_ints("DECIDER_B_BUCKETS"), token_budget=GRAPH_TOKEN_BUDGET)) CHAT = chat_template(eng.tok) if LAYOUT == "chat" else None print("[serve] engine", dict({k: v for k, v in eng.cfg.items() if k not in ("t_buckets", "b_buckets")}, device=dev, layout=LAYOUT), flush=True) if SCHEMA_FIRST: from decider.schema_engine import SchemaEngine se = SchemaEngine(eng, chat=CHAT); print("[serve] schema cache on", flush=True) pre = os.environ.get("DECIDER_SCHEMAS") # JSON: [{"questions": {...}, "independent": true, "batch_sizes": [1, 8, 32], "state_tokens": [64, 256]}] for spec in (json.load(open(pre)) if pre else []): _, h, _ = await loop.run_in_executor(gpu, _schema_handle, spec["questions"], spec.get("independent", True), COMPILE) t = await loop.run_in_executor(gpu, se.warmup, h, spec.get("batch_sizes", (1, 8, 32)), spec.get("state_tokens", (64, 128, 256))) print(f"[serve] preloaded schema with {h.nq} rows, prefix {sum(h.tps)} tokens, graphs ready in {t:.0f}s", flush=True) if WARMUP and eng.use_graphs: # off CUDA there are no graphs to capture; every request runs eager t = await loop.run_in_executor(gpu, lambda: eng.warmup(log=lambda s: print(s, flush=True))) print(f"[serve] captured {len(eng.graphs)} graphs in {t:.0f}s", flush=True) eng.seal() print("[serve] ready", json.dumps(dict(model=MODEL_NAME, layout=LAYOUT, temperature=TEMP, temperature_by_type=TT.effective(TEMP, TEMP_BY_TYPE), **({"temperature_schema_first_by_type": TT.effective(TEMP_SCHEMA, TEMP_SCHEMA_BY_TYPE)} if SCHEMA_FIRST else {}), isolated_levels=ISOLATED, schema_first=SCHEMA_FIRST, shared=SHARED, graphs=len(eng.graphs), limits=dict( max_rows=MAX_ROWS, max_row_tokens=MAX_ROW_TOKENS, max_request_tokens=MAX_REQUEST_TOKENS, max_queue_rows=MAX_QUEUE_ROWS))), flush=True) start_workers(loop) def _stop(): global gpu, cpu, batcher_task, schema_task for t in (batcher_task, schema_task): if t is not None: t.cancel() for ex in (gpu, cpu): if ex is not None: ex.shutdown(wait=False, cancel_futures=True) gpu = cpu = batcher_task = schema_task = None @asynccontextmanager async def lifespan(app): await _start() yield _stop() app = FastAPI(title="decider", lifespan=lifespan) # ---- routes ------------------------------------------------------------------- class Req(BaseModel): context: str schema_: object = None # any JSON value: Decider._check_schema gives the 422 (pydantic's echoes NaN/Infinity and cannot be serialised) model_config = {"populate_by_name": True} def __init__(self, **kw): if "schema" in kw: kw["schema_"] = kw.pop("schema") super().__init__(**kw) class S1Req(BaseModel): state: object questions: dict model: str | None = None independent: bool = True layout: str | None = None # "state_first" forces the uncached layout on a schema-first model def _alive(): return batcher_task is not None and not batcher_task.done() and (schema_task is None or not schema_task.done()) async def _queued(items): """Queue one request's rows. Every row carries the request's id, which is what the batcher's adaptive wait reads.""" global REQ_SEQ loop = asyncio.get_running_loop(); futs = [] REQ_SEQ += 1; rid = REQ_SEQ for it in items: f = loop.create_future(); futs.append(f); queue.put_nowait((f, it, rid)) return await asyncio.gather(*futs) @app.post("/decide") async def decide(r: Req): loop = asyncio.get_running_loop() try: qs, it = await loop.run_in_executor(cpu, _prepare_decide, r.context, r.schema_) except (ValueError, KeyError) as e: stats["errors"] += 1 raise HTTPException(422, str(e)) if not qs: # empty schema: nothing to score (the 1.1.2 answer, without a forward) stats["requests"] += 1; return {} check_size([len(it["ids"])]) _admit(1) try: probs = (await _queued([it]))[0] except Exception: stats["errors"] += 1; raise finally: _release(1) stats["requests"] += 1; stats["decisions"] += len(qs); stats["rows"] += 1 return _format_decide(r.schema_, qs, probs) @app.post("/v1/systemone") async def systemone(r: S1Req): loop = asyncio.get_running_loop() if SCHEMA_FIRST and r.questions and r.layout != "state_first" and _worth_caching(r.questions, r.independent): try: # CPU: rows, prefix lengths, suffix ids -> the complete cost before any GPU work tps, row = await loop.run_in_executor(cpu, _plan_schema, r.questions, r.independent, r.state) except ValueError as e: stats["errors"] += 1 raise HTTPException(422, str(e)) check_size([tp + len(row[0]) for tp in tps]) _admit(len(tps)) try: rqs, h, index = await loop.run_in_executor(gpu, _schema_handle, r.questions, r.independent) fut = loop.create_future(); await squeue.put((fut, h, row)); p = await fut except Exception: stats["errors"] += 1; raise finally: _release(len(tps)) stats["requests"] += 1; stats["decisions"] += len(rqs); stats["schema_requests"] = stats.get("schema_requests", 0) + 1 return {"model": MODEL_NAME, "answers": S1.assemble(rqs, index, [pk.tolist() for pk in p]), "usage": {"input_tokens": len(row[0]) * h.P, "cached_tokens": sum(h.tps), "output_tokens": 0}} try: rqs, index, items, ctx_len = await loop.run_in_executor(cpu, _prepare_s1, r.state, r.questions, r.independent) except ValueError as e: stats["errors"] += 1 raise HTTPException(422, str(e)) check_size([len(it["ids"]) for it in items]) _admit(len(items)) try: if SHARED and len(items) > 1 and min(len(it["ids"]) for it in items) >= SHARED_MIN_TOKENS: res = await loop.run_in_executor(gpu, _score_shared, items) # long state: run it once, fork the cache per question stats["shared_prefix_requests"] += 1 else: res = await _queued(items) except Exception: stats["errors"] += 1; raise finally: _release(len(items)) probs = [p for ps in res for p in ps] # one prob row per question, request order stats["requests"] += 1; stats["decisions"] += len(rqs); stats["rows"] += len(items) return {"model": MODEL_NAME, "answers": S1.assemble(rqs, index, [p.tolist() for p in probs]), "usage": {"input_tokens": unique_tokens(items, ctx_len), "output_tokens": 0}} @app.get("/v1/models") async def models(): return {"models": [{"name": MODEL_NAME, "description": "decider: one-pass typed decisions with calibrated probabilities", "release_date": RELEASE_DATE}]} @app.get("/health") async def health(): return {"ok": eng is not None and bool(getattr(eng, "sealed", True)) and _alive(), "model": MODEL, "device": str(getattr(eng, "dev", "")) if eng is not None else None, "layout": LAYOUT, "temperature": TEMP, "temperature_by_type": TT.effective(TEMP, TEMP_BY_TYPE), **({"temperature_schema_first_by_type": TT.effective(TEMP_SCHEMA, TEMP_SCHEMA_BY_TYPE)} if SCHEMA_FIRST else {})} @app.get("/stats") async def get_stats(): return dict(stats, outstanding_rows=outstanding, engine=eng.stats if eng else None, graphs=len(eng.graphs) if eng else 0, sealed=bool(eng and getattr(eng, "sealed", False)), schema_cache=dict(se.stats, schemas=len(schemas)) if se else None, limits=dict(max_rows=MAX_ROWS, max_row_tokens=MAX_ROW_TOKENS, max_request_tokens=MAX_REQUEST_TOKENS, max_queue_rows=MAX_QUEUE_ROWS))