Add decision core: ztc_server.py

ztc_server.py

@@ -0,0 +1,127 @@
+"""POST /v1/systemone server for FINAL-Bench/Darwin-27B-ZTC (System One request shape).
+
+    ZTC_MODEL=<snapshot path or FINAL-Bench/Darwin-27B-ZTC> PORT=8000 python ztc_server.py
+
+Every question is answered with one forward pass of the backbone and the decision readout
+(no generated tokens). Questions from concurrent requests are batched together on the GPU
+(at most ZTC_BATCH questions and ZTC_TOKEN_BUDGET estimated tokens per batch). A question
+longer than ZTC_MAX_LENGTH tokens is refused with HTTP 422 (never truncated).
+"""
+
+import asyncio
+import json
+import math
+import os
+import queue
+import sys
+import threading
+from concurrent.futures import Future
+from pathlib import Path
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+
+MODEL = os.getenv("ZTC_MODEL", "FINAL-Bench/Darwin-27B-ZTC")
+MAX_LENGTH = int(os.getenv("ZTC_MAX_LENGTH", "32768"))
+BATCH = int(os.getenv("ZTC_BATCH", "8"))
+TOKEN_BUDGET = int(os.getenv("ZTC_TOKEN_BUDGET", "16000"))
+
+path = Path(MODEL)
+if not path.exists():
+    from huggingface_hub import snapshot_download
+    path = Path(snapshot_download(MODEL, revision=os.getenv("ZTC_REVISION")))
+sys.path.insert(0, str(path))
+from autojev.model import DecisionModel, answer  # noqa: E402
+
+model = DecisionModel(checkpoint=path).eval()
+NAME = "darwin-27b-ztc"
+Q: "queue.Queue[tuple[dict, Future]]" = queue.Queue()
+
+
+def estimate(row):
+    return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192
+
+
+@torch.inference_mode()
+def probabilities(rows):
+    batch = model.prepare(rows, max_length=MAX_LENGTH)
+    values = (model(batch) / model.temperature).softmax(-1).cpu().tolist()
+    return [v[:c] for v, c in zip(values, batch.counts)]
+
+
+def run_single(row, future):
+    try:
+        future.set_result(probabilities([row])[0])
+    except ValueError as error:
+        future.set_exception(ValueError(str(error)))
+    except Exception as error:  # out of memory or other runtime failure
+        torch.cuda.empty_cache()
+        future.set_exception(RuntimeError(str(error)[:300]))
+
+
+def worker():
+    while True:
+        items = [Q.get()]
+        while len(items) < BATCH * 4:
+            try:
+                items.append(Q.get(timeout=0.005))
+            except queue.Empty:
+                break
+        items.sort(key=lambda item: estimate(item[0]))
+        i = 0
+        while i < len(items):
+            group = [items[i]]
+            tokens = estimate(items[i][0])
+            while i + len(group) < len(items) and len(group) < BATCH and tokens + estimate(items[i + len(group)][0]) <= TOKEN_BUDGET:
+                tokens += estimate(items[i + len(group)][0])
+                group.append(items[i + len(group)])
+            i += len(group)
+            try:
+                for (row, future), values in zip(group, probabilities([g[0] for g in group])):
+                    future.set_result(values)
+            except Exception:
+                torch.cuda.empty_cache()
+                for row, future in group:
+                    if not future.done():
+                        run_single(row, future)
+
+
+threading.Thread(target=worker, daemon=True).start()
+app = FastAPI(title="Darwin-27B-ZTC System One server")
+
+
+@app.get("/health")
+def health():
+    return {"status": "ready", "model": NAME, "temperature": model.temperature}
+
+
+@app.post("/v1/systemone")
+async def system_one(request: Request):
+    body = await request.json()
+    questions = body.get("questions") or {}
+    if body.get("images"):
+        raise HTTPException(422, "This model is text-only; images are not supported.")
+    futures = {}
+    for key, question in questions.items():
+        if question.get("type") not in ("choice", "noul", "score"):
+            raise HTTPException(422, f"Unsupported question type: {question.get('type')}")
+        future = Future()
+        Q.put(({"state": body.get("state"), "question": question}, future))
+        futures[key] = future
+    answers = {}
+    for key, future in futures.items():
+        try:
+            values = await asyncio.wrap_future(future)
+        except ValueError as error:
+            raise HTTPException(422, str(error))
+        except RuntimeError as error:
+            raise HTTPException(500, str(error))
+        if not all(math.isfinite(v) and v >= 0 for v in values) or sum(values) <= 0:
+            raise HTTPException(500, 'Model produced an invalid probability vector; refusing to substitute a prediction.')
+        answers[key] = answer(questions[key], values)
+    return {"model": NAME, "answers": answers}
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", "8000")), log_level="warning")