File size: 7,554 Bytes
3772a43 d5b20f4 9c472a1 d5b20f4 3772a43 d5b20f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | 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 # noqa: E402
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: # 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}
API_KEY = os.getenv("DECIDE_API_KEY")
# Demo key for a browser-only (static) page: the key is visible to viewers, so it is
# limited to listed origins, a per-client and a global rate, and small requests.
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")
|