| import asyncio |
| import hashlib |
| import hmac |
| import inspect |
| import logging |
| import os |
| import queue |
| import threading |
| import time |
| import uuid |
| from collections import OrderedDict |
| from concurrent.futures import Future |
| from contextlib import asynccontextmanager |
| from dataclasses import dataclass, field |
| from pathlib import Path |
|
|
| import torch |
| from fastapi import FastAPI, HTTPException, Request |
| from fastapi.exceptions import RequestValidationError |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.responses import JSONResponse |
|
|
| from models.autoports.cloudflare_clef.tt import encode as clef_encode |
| from models.autoports.cloudflare_clef.tt import head as clef_head |
| from models.autoports.cloudflare_clef.tt import precision_defaults |
| from models.autoports.cloudflare_clef.tt.api import SystemOneRequest, api_request, decode_media, to_record |
|
|
| DEFAULT_REPO = "Cloudflare/clef" |
| DEFAULT_REVISION = "2f3de3dd85f379784083b0814d997ab627200f0c" |
| DEFAULT_MODEL_NAMES = "clef" |
| HIDDEN = 5120 |
| ALIGN = 128 |
| TP = 2 |
| L1_SMALL_SIZE = 24576 |
| COMMAND_QUEUES = 2 |
| TRACE_REGION = 1 << 30 |
| ENCODE_MAX_LENGTH = 1 << 31 |
| API_KEY = os.environ.get("CLEF_API_KEY") |
| OPEN_PATHS = ("/health", "/v1/health") |
| PARENTS = {"1x4": ("FABRIC_1D", (1, 4)), "2x2": ("FABRIC_2D", (2, 2))} |
| MESH_DEVICE_SHAPES = {"P150X2": (1, 2), "P300": (1, 2), "P150X4": (1, 4), "P300X2": (1, 4), "QB2": (2, 2)} |
|
|
| log = logging.getLogger("clef.tt.server") |
| if not log.handlers: |
| handler = logging.StreamHandler() |
| handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s")) |
| log.addHandler(handler) |
| log.setLevel(logging.INFO) |
|
|
| WARMUP = { |
| "model": "clef", |
| "state": "Checkout has been failing for every customer for the last hour.", |
| "questions": { |
| "urgent": {"type": "noul", "instructions": "Is this support request urgent?"}, |
| "team": { |
| "type": "choice", |
| "instructions": "Which team should handle this request?", |
| "criteria": { |
| "billing": "Payments, invoices, and refunds", |
| "technical": "Outages, errors, and configuration", |
| "sales": "Plans and upgrades", |
| }, |
| }, |
| "severity": { |
| "type": "score", |
| "instructions": "How severe is the customer impact?", |
| "criteria": ["No impact", "Minor", "Major", "Critical"], |
| }, |
| }, |
| } |
|
|
|
|
| def flag(name, default="0"): |
| return os.environ.get(name, default) == "1" |
|
|
|
|
| def parse_mesh_shape(shape, mesh_device): |
| if shape: |
| r, c = shape.lower().split("x") |
| return int(r), int(c) |
| return MESH_DEVICE_SHAPES.get((mesh_device or "").upper(), (1, 2)) |
|
|
|
|
| def parse_offset(text): |
| r, c = (text or "0,0").split(",") |
| return int(r), int(c) |
|
|
|
|
| def resolve_model(value, revision=None): |
| if os.path.isdir(value): |
| return str(Path(value).resolve()) |
| from huggingface_hub import snapshot_download |
|
|
| repo, _, at_revision = value.partition("@") |
| offline = os.environ.get("HF_HUB_OFFLINE") == "1" |
| return snapshot_download(repo, revision=at_revision or revision or None, local_files_only=offline) |
|
|
|
|
| @dataclass |
| class Settings: |
| model: str |
| revision: str |
| snapshot: str |
| model_names: tuple |
| mesh_shape: tuple |
| parent_mesh: str |
| submesh_offset: tuple |
| fake: bool |
| fake_workers: int |
| trace_region: int |
| prefix_cache: int |
| max_state: int |
| max_tail: int |
| truncate: bool |
| traced: bool |
| planner: bool = True |
| allow_remote_images: bool = True |
| warmup: bool = True |
| extra: dict = field(default_factory=dict) |
|
|
| @classmethod |
| def from_env(cls): |
| model = os.environ.get("CLEF_MODEL") or os.environ.get("HF_MODEL") or DEFAULT_REPO |
| revision = os.environ.get("CLEF_REVISION", DEFAULT_REVISION) |
| snapshot = resolve_model(model, revision) |
| os.environ["CLEF_MODEL"] = snapshot |
| traced = flag("CLEF_TRACED", "0") |
| return cls( |
| model=model, |
| revision=revision, |
| snapshot=snapshot, |
| model_names=tuple( |
| n.strip() for n in os.environ.get("CLEF_MODEL_NAMES", DEFAULT_MODEL_NAMES).split(",") if n.strip() |
| ), |
| mesh_shape=parse_mesh_shape(os.environ.get("CLEF_MESH_SHAPE"), os.environ.get("MESH_DEVICE")), |
| parent_mesh=os.environ.get("CLEF_PARENT_MESH", ""), |
| submesh_offset=parse_offset(os.environ.get("CLEF_SUBMESH_OFFSET")), |
| fake=flag("CLEF_FAKE_ENGINE"), |
| fake_workers=int(os.environ.get("CLEF_FAKE_WORKERS", "1")), |
| trace_region=int(os.environ.get("CLEF_TRACE_REGION", str(TRACE_REGION if traced else 0))), |
| prefix_cache=int(os.environ.get("CLEF_PREFIX_CACHE", "4")), |
| max_state=int(os.environ.get("CLEF_MAX_STATE", "16384")), |
| max_tail=int(os.environ.get("CLEF_MAX_TAIL", "4096")), |
| truncate=flag("CLEF_TRUNCATE_STATES"), |
| traced=traced, |
| planner=flag("CLEF_PLANNER", "1"), |
| allow_remote_images=flag("CLEF_ALLOW_REMOTE_IMAGES", "1"), |
| warmup=flag("CLEF_WARMUP", "1"), |
| ) |
|
|
|
|
| def mesh_plan(mesh_shape, visible, parent="", offset=(0, 0)): |
| r, c = mesh_shape |
| if (r, c) == (1, 2): |
| if parent or visible > 2: |
| name = parent or "1x4" |
| if name not in PARENTS: |
| raise ValueError(f"CLEF_PARENT_MESH={name!r}; expected one of {sorted(PARENTS)}") |
| fabric, shape = PARENTS[name] |
| return {"fabric": fabric, "open": shape, "submeshes": [((1, 2), tuple(offset))], "parent": name} |
| return {"fabric": "FABRIC_1D", "open": (1, 2), "submeshes": [], "parent": None} |
| if (r, c) == (1, 4): |
| if parent: |
| raise ValueError("CLEF_PARENT_MESH is only supported for CLEF_MESH_SHAPE=1x2") |
| return { |
| "fabric": "FABRIC_1D", |
| "open": (1, 4), |
| "submeshes": [((1, 2), (0, 0)), ((1, 2), (0, 2))], |
| "parent": None, |
| } |
| if (r, c) == (2, 2): |
| if parent: |
| raise ValueError("CLEF_PARENT_MESH is only supported for CLEF_MESH_SHAPE=1x2") |
| return { |
| "fabric": "FABRIC_2D", |
| "open": (2, 2), |
| "submeshes": [((1, 2), (0, 0)), ((1, 2), (1, 0))], |
| "parent": None, |
| } |
| raise ValueError(f"unsupported mesh shape {r}x{c}; Clef runs TP={TP} groups on 1x2, 1x4 or 2x2") |
|
|
|
|
| def visible_chips(): |
| import ttnn |
|
|
| return len(ttnn.get_device_ids()) |
|
|
|
|
| def open_devices(settings): |
| import ttnn |
|
|
| visible = visible_chips() |
| plan = mesh_plan(settings.mesh_shape, visible, settings.parent_mesh, settings.submesh_offset) |
| ttnn.set_fabric_config(getattr(ttnn.FabricConfig, plan["fabric"])) |
| parent = ttnn.open_mesh_device( |
| mesh_shape=ttnn.MeshShape(*plan["open"]), |
| l1_small_size=L1_SMALL_SIZE, |
| num_command_queues=COMMAND_QUEUES, |
| trace_region_size=settings.trace_region, |
| ) |
| groups = [] |
| for shape, offset in plan["submeshes"]: |
| groups.append(parent.create_submesh(ttnn.MeshShape(*shape), offset=ttnn.MeshCoordinate(*offset))) |
| if not groups: |
| groups = [parent] |
| for group in groups: |
| group.enable_program_cache() |
| log.info( |
| "mesh: %d chip(s) visible, %s parent %s chips %s, %d TP group(s) %s", |
| visible, |
| plan["fabric"], |
| list(parent.shape), |
| list(parent.get_device_ids()), |
| len(groups), |
| [list(g.get_device_ids()) for g in groups], |
| ) |
|
|
| def close(): |
| for group in groups: |
| ttnn.synchronize_device(group) |
| for child in parent.get_submeshes(): |
| ttnn.close_mesh_device(child) |
| ttnn.close_mesh_device(parent) |
| ttnn.set_fabric_config(ttnn.FabricConfig.DISABLED) |
|
|
| return parent, groups, close, plan |
|
|
|
|
| def accepted_kwargs(fn, candidates): |
| params = inspect.signature(fn).parameters |
| if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()): |
| return dict(candidates) |
| return {k: v for k, v in candidates.items() if k in params} |
|
|
|
|
| def build_engine(mesh, settings): |
| from models.autoports.cloudflare_clef.tt.engine import ClefEngine |
|
|
| candidates = dict( |
| max_state_len=settings.max_state, |
| max_tail_len=settings.max_tail, |
| snapshot_slots=max(1, settings.prefix_cache), |
| traced=settings.traced, |
| planner=settings.planner, |
| ) |
| return ClefEngine(mesh, **accepted_kwargs(ClefEngine.__init__, candidates)) |
|
|
|
|
| def build_fake_engine(settings): |
| from models.autoports.cloudflare_clef.tt.fake_engine import FakeEngine |
|
|
| return FakeEngine( |
| max_state_len=settings.max_state, max_tail_len=settings.max_tail, snapshot_slots=max(1, settings.prefix_cache) |
| ) |
|
|
|
|
| def media_key(ids, media_digest): |
| key = clef_encode.cache_key(ids) |
| if media_digest is None: |
| return key |
| return hashlib.sha1(f"{key}:{media_digest}".encode()).hexdigest() |
|
|
|
|
| @dataclass |
| class Job: |
| ids: list |
| state_ids: list |
| tail_ids: list |
| encoded: object |
| media: dict |
| key: str |
| questions: dict |
|
|
|
|
| def supports_media(engine): |
| try: |
| return "media" in inspect.signature(engine.prefill_state).parameters |
| except (TypeError, ValueError): |
| return bool(getattr(engine, "supports_media", False)) |
|
|
|
|
| def supports_tail_len(engine): |
| try: |
| return "tail_len" in inspect.signature(engine.prefill_state).parameters |
| except (TypeError, ValueError): |
| return False |
|
|
|
|
| class Worker: |
| def __init__(self, wid, build, head, lm_rows, cache_size): |
| self.id = wid |
| self.build = build |
| self.head = head |
| self.lm_rows = lm_rows |
| self.cache_size = cache_size |
| self.queue = queue.Queue() |
| self.cache = OrderedDict() |
| self.free_slots = [] |
| self.hits = self.misses = self.requests = 0 |
| self.inflight = 0 |
| self.counter = threading.Lock() |
| self.ready = threading.Event() |
| self.error = None |
| self.engine = None |
| self.media_ok = False |
| self.tail_len_ok = False |
| self.traced = False |
| self.warm_grids = None |
| self.gdn_conv = None |
| self.planner = None |
| self.thread = threading.Thread(target=self._run, name=f"clef-worker-{wid}", daemon=True) |
| self.thread.start() |
|
|
| def load(self): |
| return self.queue.qsize() + self.inflight |
|
|
| def submit(self, job): |
| done = Future() |
| with self.counter: |
| self.inflight += 1 |
| self.queue.put((job, done)) |
| return done |
|
|
| def stop(self): |
| self.queue.put(None) |
| self.thread.join() |
|
|
| def _run(self): |
| try: |
| self.engine = self.build() |
| slots = getattr(self.engine, "snapshot_slots", 1) |
| self.cache_size = min(self.cache_size, slots) |
| self.free_slots = list(range(self.cache_size)) |
| self.media_ok = supports_media(self.engine) |
| self.tail_len_ok = supports_tail_len(self.engine) |
| self.traced = bool(getattr(self.engine, "traced", False)) |
| self.gdn_conv = getattr(self.engine, "gdn_conv_impl", None) |
| self.planner = getattr(self.engine, "planner", None) |
| if self.traced: |
| self.warm_grids = sorted(tuple(g) for g, _ in getattr(self.engine, "vision_warmed_grids", [])) |
| except Exception as error: |
| self.error = error |
| self.ready.set() |
| return |
| self.ready.set() |
| while True: |
| item = self.queue.get() |
| if item is None: |
| return |
| job, done = item |
| try: |
| done.set_result(self._serve(job)) |
| except Exception as error: |
| done.set_exception(error) |
| finally: |
| with self.counter: |
| self.inflight -= 1 |
|
|
| def _media_kwargs(self, job): |
| if job.media is None: |
| return {} |
| if not self.media_ok: |
| raise HTTPException(422, "this server build does not accept images or videos") |
| return {"media": job.media} |
|
|
| def _serve(self, job): |
| t0 = time.perf_counter() |
| ids = torch.tensor([job.ids], dtype=torch.long) |
| hit = False |
| if self.cache_size == 0: |
| hidden = self.engine.prefill_hidden(ids, slot=0, **self._media_kwargs(job)) |
| else: |
| entry = self.cache.get(job.key) |
| hit = entry is not None |
| if hit: |
| self.cache.move_to_end(job.key) |
| self.hits += 1 |
| handle, slot = entry |
| else: |
| self.misses += 1 |
| slot = self.free_slots.pop() if self.free_slots else self.cache.popitem(last=False)[1][1] |
| hint = {"tail_len": len(job.tail_ids)} if self.tail_len_ok else {} |
| try: |
| handle = self.engine.prefill_state( |
| torch.tensor([job.state_ids], dtype=torch.long), |
| slot=slot, |
| key=job.key, |
| **self._media_kwargs(job), |
| **hint, |
| ) |
| except ValueError as error: |
| self.free_slots.append(slot) |
| raise HTTPException(422, str(error)) |
| except Exception: |
| self.free_slots.append(slot) |
| raise |
| self.cache[job.key] = (handle, slot) |
| tail = self.engine.schema_hidden(handle, torch.tensor([job.tail_ids], dtype=torch.long)) |
| hidden = torch.cat([self.engine.prefix_hidden[slot], tail], dim=0) |
| probs = clef_head.probs_for_record(self.head, hidden, ids[0], job.encoded, self.lm_rows) |
| self.requests += 1 |
| return probs, hit, round((time.perf_counter() - t0) * 1000, 1) |
|
|
|
|
| class Server: |
| def __init__(self, settings): |
| self.settings = settings |
| self.snapshot = settings.snapshot |
| self.release = clef_encode.release_module(self.snapshot) |
| self.tok = clef_encode.load_tokenizer(self.snapshot) |
| self.processor = None |
| self.head = clef_head.load_head(self.snapshot) |
| self.parent = None |
| self.close_devices = lambda: None |
| self.plan = None |
| self.lock = threading.Lock() |
| self.workers = [] |
| self.stopping = threading.Event() |
| self.device_name = "none" |
| self.started = time.time() |
|
|
| @classmethod |
| def start(cls, settings): |
| server = cls(settings) |
| server.open() |
| if settings.warmup: |
| server.warmup() |
| return server |
|
|
| def processor_for_media(self): |
| if self.processor is None: |
| with self.lock: |
| if self.processor is None: |
| self.processor = clef_encode.load_processor(self.snapshot) |
| return self.processor |
|
|
| def _lm_rows(self): |
| from models.autoports.cloudflare_clef.tt.loader import LmHeadRows |
|
|
| return LmHeadRows(self.snapshot) |
|
|
| def open(self): |
| s = self.settings |
| if s.fake: |
| builds = [(f"fake{i}", lambda: build_fake_engine(s)) for i in range(max(1, s.fake_workers))] |
| self.device_name = f"fake x{len(builds)}" |
| else: |
| self.parent, groups, self.close_devices, self.plan = open_devices(s) |
| builds = [(str(list(g.get_device_ids())), lambda g=g: build_engine(g, s)) for g in groups] |
| self.device_name = f"ttnn {s.mesh_shape[0]}x{s.mesh_shape[1]}: {len(groups)} TP={TP} worker(s)" |
| self.workers = [ |
| Worker(i, build, self.head, self._lm_rows(), s.prefix_cache) for i, (_, build) in enumerate(builds) |
| ] |
| for w, (name, _) in zip(self.workers, builds): |
| w.ready.wait() |
| if w.error is not None: |
| self.close() |
| raise RuntimeError(f"worker {w.id} ({name}) failed to load") from w.error |
| log.info( |
| "worker %d ready on %s, prefix cache %d state(s), media=%s, mode=%s, gdn_conv=%s, planner=%s, " |
| "warm_grids=%s", |
| w.id, |
| name, |
| w.cache_size, |
| w.media_ok, |
| "traced" if w.traced else "eager", |
| w.gdn_conv, |
| w.planner, |
| w.warm_grids, |
| ) |
|
|
| def warmup(self): |
| req = SystemOneRequest.model_validate(WARMUP) |
| job, _ = self.encode(req) |
| for w in self.workers: |
| probs, _, ms = w.submit(job).result() |
| log.info("warmup worker=%d questions=%d latency_ms=%.1f", w.id, len(probs), ms) |
|
|
| def close(self): |
| if self.stopping.is_set(): |
| return |
| self.stopping.set() |
| for w in self.workers: |
| if w.ready.is_set() and w.error is None: |
| w.stop() |
| self.close_devices() |
|
|
| def _encode_record(self, record, processor, max_state_tokens=None): |
| try: |
| encoded = clef_encode.encode( |
| self.tok, record, processor=processor, max_length=ENCODE_MAX_LENGTH, max_state_tokens=max_state_tokens |
| ) |
| except ValueError as error: |
| raise HTTPException(422, str(error)) |
| try: |
| state_ids, tail_ids, _ = clef_encode.split_for_cache(encoded, self.tok, record) |
| except ValueError as error: |
| raise HTTPException(500, f"prefix split failed: {error}") |
| return encoded, state_ids, tail_ids |
|
|
| def encode(self, req): |
| s = self.settings |
| try: |
| images, videos, media_digest = decode_media(req) |
| except ValueError as error: |
| raise HTTPException(422, str(error)) |
| record = to_record(req, images, videos) |
| processor = self.processor_for_media() if (images or videos) else None |
| encoded, state_ids, tail_ids = self._encode_record(record, processor) |
| if encoded.media: |
| self.check_warm_grids(encoded.media) |
| S = len(state_ids) |
| stats = {"state_tokens": S, "state_tokens_used": S} |
| if S > s.max_state: |
| n_state = len(self.release._tokens(self.tok, self.release.render(req.state))) |
| fixed = S - n_state |
| keep = s.max_state - fixed |
| if not s.truncate or keep < 0: |
| raise HTTPException(422, overflow_message(S, fixed, s.max_state)) |
| encoded, state_ids, tail_ids = self._encode_record(record, processor, max_state_tokens=keep) |
| stats["state_tokens_used"] = len(state_ids) |
| if len(tail_ids) > s.max_tail: |
| raise HTTPException( |
| 422, |
| f"schema is {len(tail_ids):,} tokens, over the {s.max_tail:,}-token limit: " |
| "use fewer questions or shorter criteria, or split them across requests", |
| ) |
| key = media_key(state_ids, media_digest) |
| job = Job(list(encoded.input_ids), state_ids, tail_ids, encoded, encoded.media, key, req.question_dicts()) |
| stats["tokens"] = len(encoded.input_ids) |
| stats["images"] = len(images) |
| stats["videos"] = len(videos) |
| return job, stats |
|
|
| def warm_grids(self): |
| lists = [w.warm_grids for w in self.workers if w.warm_grids is not None] |
| if not lists: |
| return None |
| return sorted({grid for grids in lists for grid in grids}) |
|
|
| def check_warm_grids(self, media): |
| warm = self.warm_grids() |
| if warm is None: |
| return |
| rows = [] |
| for name in ("image_grid_thw", "video_grid_thw"): |
| grids = media.get(name) |
| if grids is not None: |
| rows.extend(torch.as_tensor(grids).reshape(-1, 3).tolist()) |
| if len(rows) > 1: |
| raise HTTPException(422, single_grid_message(len(rows))) |
| for grid in rows: |
| if tuple(grid) not in warm: |
| raise HTTPException(422, warm_grid_message(tuple(grid), warm)) |
|
|
| def pick_worker(self, key): |
| with self.lock: |
| for w in self.workers: |
| if key in w.cache: |
| return w |
| return min(self.workers, key=lambda w: (w.load(), w.id)) |
|
|
| def submit(self, req): |
| if self.stopping.is_set(): |
| raise HTTPException(503, "the server is stopping") |
| job, stats = self.encode(req) |
| worker = self.pick_worker(job.key) |
| done = Future() |
| inner = worker.submit(job) |
|
|
| def finish(f): |
| if f.exception() is not None: |
| done.set_exception(f.exception()) |
| return |
| probs, hit, ms = f.result() |
| stats.update(worker=worker.id, latency_ms=ms, prefix_cache_hit=hit) |
| log.info( |
| "worker=%d S=%d tail=%d questions=%d images=%d videos=%d latency_ms=%.1f cache_hit=%s", |
| worker.id, |
| stats["state_tokens_used"], |
| len(job.tail_ids), |
| len(probs), |
| stats["images"], |
| stats["videos"], |
| ms, |
| hit, |
| ) |
| done.set_result((probs, stats, job)) |
|
|
| inner.add_done_callback(finish) |
| return done |
|
|
| def as_request(self, rec): |
| return rec if isinstance(rec, SystemOneRequest) else SystemOneRequest.model_validate(api_request(rec)) |
|
|
| def predict(self, rec): |
| probs, stats, _ = self.submit(self.as_request(rec)).result() |
| return probs, stats |
|
|
| def answer(self, req): |
| return self.body(req, *self.submit(req).result()) |
|
|
| async def answer_async(self, req): |
| return self.body(req, *await asyncio.wrap_future(self.submit(req))) |
|
|
| def body(self, req, probs, stats, job): |
| answers = {qid: self.release.systemone_answer(job.questions[qid], dist) for qid, dist in probs.items()} |
| body = { |
| "model": req.model, |
| "answers": answers, |
| "usage": {"input_tokens": stats["tokens"], "output_tokens": 0}, |
| "latency_ms": stats["latency_ms"], |
| } |
| if self.settings.truncate: |
| body["usage"].update(state_tokens=stats["state_tokens"], state_tokens_used=stats["state_tokens_used"]) |
| body["truncated"] = stats["state_tokens"] > stats["state_tokens_used"] |
| return body |
|
|
| def card(self): |
| s = self.settings |
| workers = [ |
| { |
| "id": w.id, |
| "queued": w.load(), |
| "cached_states": len(w.cache), |
| "cache_size": w.cache_size, |
| "hits": w.hits, |
| "misses": w.misses, |
| "requests": w.requests, |
| "media": w.media_ok, |
| } |
| for w in self.workers |
| ] |
| backend = "fake engine" if s.fake else "ttnn, TP=2" |
| warm = self.warm_grids() |
| traced_media = None |
| if warm is not None: |
| traced_media = { |
| "warm_grids": [list(grid) for grid in warm], |
| "rule": "one image or video grid per request, and a grid (t, h, w) outside this list is refused " |
| "with 422; eager (CLEF_TRACED=0) accepts any grid and several images or videos per request", |
| } |
| return { |
| "description": f"Cloudflare/clef joint schema decisions on Tenstorrent ({backend})", |
| "weights": {"repo": DEFAULT_REPO, "revision": s.revision, "snapshot": self.snapshot}, |
| "device": self.device_name, |
| "mesh_shape": f"{s.mesh_shape[0]}x{s.mesh_shape[1]}", |
| "mesh_plan": self.plan, |
| "backend": "fake" if s.fake else "ttnn", |
| "precision": precision_defaults.active(), |
| "traced": s.traced, |
| "mode": "traced" if s.traced else "eager", |
| "gdn_conv": next((w.gdn_conv for w in self.workers if w.gdn_conv), None), |
| "prefix_planner": s.planner, |
| "traced_media": traced_media, |
| "max_state_tokens": s.max_state, |
| "max_schema_tokens": s.max_tail, |
| "truncate_states": s.truncate, |
| "remote_images": s.allow_remote_images, |
| "prefix_cache": { |
| "size": s.prefix_cache, |
| "hits": sum(w.hits for w in self.workers), |
| "misses": sum(w.misses for w in self.workers), |
| "cached_states": sum(len(w.cache) for w in self.workers), |
| }, |
| "workers": workers, |
| "requests": sum(w.requests for w in self.workers), |
| "uptime_s": round(time.time() - self.started, 1), |
| } |
|
|
|
|
| def overflow_message(S, fixed, max_state): |
| return ( |
| f"state is {S:,} tokens ({fixed} of them the system prompt and media placeholders), " |
| f"over the {max_state:,}-token limit: shorten the state or split it across requests; " |
| f"or start the server with CLEF_TRUNCATE_STATES=1 to read only its first {max_state:,} tokens " |
| "(responses then say truncated: true)" |
| ) |
|
|
|
|
| def warm_grid_message(grid, warm): |
| grids = "; ".join(",".join(str(x) for x in g) for g in warm) |
| return ( |
| f"image grid {grid} (t, h, w patches) is not in this traced server's warm list [{grids}]: " |
| "resize the image to a warmed grid, add the grid to CLEF_VISION_WARM_GRID at startup, " |
| "or serve with CLEF_TRACED=0 (the default), which accepts any grid" |
| ) |
|
|
|
|
| def single_grid_message(n_grids): |
| return ( |
| f"this request carries {n_grids} image or video grids; a traced server (CLEF_TRACED=1) takes one grid " |
| "per request, because the tower joins several images with a concat program compiled per image count " |
| "and no vision program may compile while traces are live: send one image (or one video) per request, " |
| "or serve with CLEF_TRACED=0 (the default), which accepts several" |
| ) |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(app): |
| settings = Settings.from_env() |
| log.info( |
| "starting: model=%s snapshot=%s fake=%s mesh=%s parent=%s offset=%s prefix_cache=%d max_state=%d " |
| "max_tail=%d truncate=%s traced=%s planner=%s trace_region=%d", |
| settings.model, |
| settings.snapshot, |
| settings.fake, |
| settings.mesh_shape, |
| settings.parent_mesh or "auto", |
| settings.submesh_offset, |
| settings.prefix_cache, |
| settings.max_state, |
| settings.max_tail, |
| settings.truncate, |
| settings.traced, |
| settings.planner, |
| settings.trace_region, |
| ) |
| app.state.server = await asyncio.to_thread(Server.start, settings) |
| server_ = app.state.server |
| log.info( |
| "ready: %d worker(s) on %s, mode=%s, gdn_conv=%s, planner=%s, warm_grids=%s", |
| len(server_.workers), |
| server_.device_name, |
| "traced" if settings.traced else "eager", |
| next((w.gdn_conv for w in server_.workers if w.gdn_conv), None), |
| settings.planner, |
| server_.warm_grids(), |
| ) |
| try: |
| yield |
| finally: |
| app.state.server.close() |
|
|
|
|
| app = FastAPI(title="clef-tt", lifespan=lifespan) |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| expose_headers=["x-typesafe-request-id", "server-timing"], |
| ) |
|
|
|
|
| def validation_detail(error): |
| errors = error.errors() |
| if not errors: |
| return "invalid request" |
| first = errors[0] |
| msg = str(first.get("msg", "invalid request")) |
| if msg.startswith("Value error, "): |
| return msg[len("Value error, ") :] |
| loc = ".".join(str(p) for p in first.get("loc", ()) if p != "body") |
| return f"{loc}: {msg}" if loc else msg |
|
|
|
|
| @app.exception_handler(RequestValidationError) |
| async def on_validation_error(request: Request, error: RequestValidationError): |
| return JSONResponse({"detail": validation_detail(error)}, 422) |
|
|
|
|
| @app.middleware("http") |
| async def typesafe(request, call_next): |
| started = time.perf_counter() |
| path = request.url.path |
| if ( |
| API_KEY |
| and path.startswith("/v1") |
| and path not in OPEN_PATHS |
| and not hmac.compare_digest(request.headers.get("authorization", ""), f"Bearer {API_KEY}") |
| ): |
| resp = JSONResponse( |
| {"detail": "missing or invalid API key; send Authorization: Bearer <CLEF_API_KEY>"}, |
| 401, |
| {"www-authenticate": "Bearer"}, |
| ) |
| else: |
| resp = await call_next(request) |
| resp.headers["x-typesafe-request-id"] = request.headers.get("x-typesafe-request-id") or uuid.uuid4().hex |
| resp.headers["server-timing"] = f"app;dur={(time.perf_counter() - started) * 1000:.1f}" |
| return resp |
|
|
|
|
| def server() -> Server: |
| return app.state.server |
|
|
|
|
| @app.post("/v1/systemone") |
| async def systemone(req: SystemOneRequest): |
| return await server().answer_async(req) |
|
|
|
|
| @app.get("/v1/models") |
| def models(): |
| card = server().card() |
| return {"models": [{"name": name, **card} for name in server().settings.model_names]} |
|
|
|
|
| @app.get("/health") |
| @app.get("/v1/health") |
| def health(): |
| s = server() |
| return {"status": "ok", "workers": len(s.workers), "queued": sum(w.load() for w in s.workers)} |
|
|
|
|
| def main(): |
| import argparse |
|
|
| import uvicorn |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--host", default="127.0.0.1") |
| ap.add_argument("--port", type=int, default=8008) |
| a = ap.parse_args() |
| uvicorn.run(app, host=a.host, port=a.port, lifespan="on") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
