| import hashlib |
| import time |
| """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 |
| import time |
| from collections import defaultdict, deque |
|
|
| from fastapi import FastAPI, HTTPException, Request |
| from fastapi.middleware.cors import CORSMiddleware |
|
|
| 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 |
|
|
| model = DecisionModel(checkpoint=path).eval() |
| NAME = os.getenv("DECIDE_NAME", "darwin-35b-ztc-vision") |
| Q: "queue.Queue[tuple[dict, Future]]" = queue.Queue() |
|
|
|
|
| def estimate(row): |
| return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192 + 1200 * len(row.get("images", [])) |
|
|
|
|
| @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: |
| 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} |
|
|
|
|
| API_KEY = os.getenv("DECIDE_API_KEY") |
| |
| |
| DEMO_KEY = os.getenv("DECIDE_DEMO_KEY") |
| DEMO_ORIGINS = [o for o in os.getenv("DECIDE_DEMO_ORIGINS", "").split(",") if o] |
| DEMO_PER_MIN = int(os.getenv("DECIDE_DEMO_PER_MIN", "40")) |
| DEMO_GLOBAL_PER_MIN = int(os.getenv("DECIDE_DEMO_GLOBAL_PER_MIN", "240")) |
| HITS = defaultdict(deque) |
| if DEMO_ORIGINS: |
| app.add_middleware(CORSMiddleware, allow_origins=DEMO_ORIGINS, allow_methods=["POST", "GET", "OPTIONS"], |
| allow_headers=["Authorization", "Content-Type"], max_age=600) |
|
|
|
|
| def demo_check(request, body): |
| if DEMO_ORIGINS and request.headers.get("origin", "") not in DEMO_ORIGINS: |
| raise HTTPException(403, "This key only works from the demo page.") |
| now = time.time() |
| client = request.headers.get("cf-connecting-ip") or (request.client.host if request.client else "?") |
| for name, limit in ((client, DEMO_PER_MIN), ("*", DEMO_GLOBAL_PER_MIN)): |
| q = HITS[name] |
| while q and now - q[0] > 60: |
| q.popleft() |
| if len(q) >= limit: |
| raise HTTPException(429, "Too many requests. Wait a minute and try again.") |
| for name in (client, "*"): |
| HITS[name].append(now) |
| images = body.get("images") or [] |
| if len(images) > 4 or any(len(i) > 4_000_000 for i in images): |
| raise HTTPException(413, "Up to 4 images, each under about 3 MB.") |
| if len(body.get("questions") or {}) > 4: |
| raise HTTPException(413, "Up to 4 questions per request.") |
|
|
|
|
| @app.post("/v1/systemone") |
| async def system_one(request: Request): |
| auth = request.headers.get("authorization", "") |
| demo = bool(DEMO_KEY) and auth == "Bearer " + DEMO_KEY |
| if API_KEY and auth != "Bearer " + API_KEY and not demo: |
| raise HTTPException(401, "Invalid API key.") |
| body = await request.json() |
| if demo: |
| demo_check(request, body) |
| questions = body.get("questions") or {} |
| images = body.get("images") or [] |
| if images and model.text_only: |
| raise HTTPException(422, "This model is text-only; images are not supported.") |
| if len(images) > 4: |
| raise HTTPException(422, "At most 4 images per request.") |
| 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() |
| row = {"state": body.get("state"), "question": question} |
| if images: |
| row["images"] = images |
| Q.put((row, 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)) |
| bad = not all(math.isfinite(v) and v >= 0 for v in values) or sum(values) <= 0 |
| flat = (not bad) and len(values) > 1 and max(values) - min(values) < 1e-9 |
| if bad or flat: |
| key_hash = hashlib.sha256(json.dumps([body.get('state'), questions[key], len(images)], sort_keys=True, default=str).encode()).hexdigest()[:16] |
| with open(os.getenv('INVALID_LOG', 'invalid_outputs.jsonl'), 'a') as fh: |
| fh.write(json.dumps({'time': time.time(), 'request': key_hash, 'question': key, 'kind': 'invalid' if bad else 'uniform', 'values': [repr(v) for v in values]}) + '\n') |
| if bad: |
| 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") |
|
|
