Download source/src/deg/server.py from gcoderw/DEG-5B-E2B-1019: direct link, hf CLI and curl.
- Browser
- Download file 21.6 kB
-
https://huggingface.co/gcoderw/DEG-5B-E2B-1019/resolve/main/source/src/deg/server.py
- Command line
-
hf download hf://gcoderw/DEG-5B-E2B-1019/source/src/deg/server.py
-
curl -L -o server.py https://huggingface.co/gcoderw/DEG-5B-E2B-1019/resolve/main/source/src/deg/server.py
21.6 kB
| """deg HTTP server: OpenAI Decisions, TypeSafe Jev SystemOne and deg fill on one model. | |
| Routes (shapes in :mod:`deg.api`): | |
| * ``POST /v1/decisions``: OpenAI Decisions API (stock ``openai`` >= 3.26 works with | |
| ``OpenAI(base_url="http://host:8790/v1")``); deg extensions ``extract`` / ``generate``. | |
| * ``POST /v1/systemone``: TypeSafe Jev / SystemOne (and Clef); same extensions. | |
| * ``POST /v1/fill``: deg's schema fill (SPEC.md §9). | |
| * ``GET /v1/models``, ``GET /health``. | |
| Every request becomes one :func:`deg.infer.fill` job (all of a request's questions are | |
| fields of one schema: the input is read once). Jobs go through a :class:`Batcher`: one GPU | |
| worker; concurrent jobs that arrive within ``window_s`` of the first are taken together (up | |
| to ``max_batch``) and handed to ``fill_many`` when the runtime provides one, else run back | |
| to back on the worker. A failing job fails only itself. More than ``max_pending`` queued | |
| jobs -> 429. | |
| Auth: when ``DEG_API_KEY`` is set, requests need ``Authorization: Bearer <key>`` (401 in | |
| OpenAI shape on ``/v1/decisions`` and ``/v1/fill``; 403 in Jev shape on ``/v1/systemone``, | |
| as api.typesafe.ai answers). | |
| Telemetry (OpenTelemetry API; exporters are the deployer's choice). Cardinality: | |
| * span ``deg.request``: ``deg.route`` in {decisions, systemone, fill, models, health, | |
| other} (6), ``deg.status`` in {ok, client_error, unauthorized, overloaded, | |
| server_error} (5), ``deg.questions`` and ``deg.batch_size`` (integers on spans only, | |
| never metric dimensions); | |
| * histogram ``deg.server.duration`` (s) by route x status: <= 30 series; | |
| * histogram ``deg.server.batch_size``: no attributes; | |
| * counter ``deg.server.questions`` by route (2) x question_type (5): <= 10 series; | |
| * counter ``deg.server.tokens`` by route (3) x direction {input, output}: <= 6 series. | |
| No request content, question names, model names or client identifiers are recorded. | |
| ``python -m deg.server --ckpt /data/deg/runs/run-1/step-7000 --port 8790`` | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import dataclasses | |
| import json | |
| import math | |
| import os | |
| import queue | |
| import sys | |
| import threading | |
| import time | |
| from collections.abc import Callable | |
| from concurrent.futures import Future | |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer | |
| from typing import Any | |
| from opentelemetry import metrics, trace | |
| from deg.api import ( | |
| JEV_COMPAT_MODELS, | |
| OPENAI_COMPAT_MODELS, | |
| ApiError, | |
| Job, | |
| decisions_response, | |
| parse_decisions, | |
| parse_systemone, | |
| systemone_response, | |
| ) | |
| from deg.request import ( | |
| MODEL_ALIASES, | |
| FillOptions, | |
| RequestError, | |
| generate_arg, | |
| parse_fill_request, | |
| ) | |
| from deg.response import MODEL_ID, FillError | |
| READ_PIECE_BYTES = 1 << 20 # socket read granularity, not a size limit | |
| ROUTES = {"/v1/decisions": "decisions", "/v1/systemone": "systemone", "/v1/fill": "fill", | |
| "/v1/models": "models", "/health": "health"} | |
| FillFn = Callable[[Any, dict[str, Any], FillOptions], dict[str, Any]] | |
| FillManyFn = Callable[[list[tuple[Any, dict[str, Any], FillOptions]]], list[Any]] | |
| _tracer = trace.get_tracer("deg.server") | |
| _meter = metrics.get_meter("deg.server") | |
| _duration = _meter.create_histogram("deg.server.duration", unit="s", | |
| description="request latency by route and status") | |
| _batch_size = _meter.create_histogram("deg.server.batch_size", unit="{job}", | |
| description="fill jobs taken together by the GPU worker") | |
| _questions = _meter.create_counter("deg.server.questions", unit="{question}", | |
| description="questions answered, by route and type") | |
| _tokens = _meter.create_counter("deg.server.tokens", unit="{token}", | |
| description="model tokens, by route and direction") | |
| class Overloaded(Exception): | |
| """Too many jobs queued (429) or the model ran out of memory (503 / 529).""" | |
| def __init__(self, message: str, status: int) -> None: | |
| super().__init__(message) | |
| self.status = status | |
| class Batcher: | |
| """One GPU worker; jobs arriving within ``window_s`` of the first are taken together.""" | |
| def __init__(self, fill: FillFn, *, fill_many: FillManyFn | None = None, | |
| window_s: float = 0.005, max_batch: int = 16, max_pending: int = 256) -> None: | |
| self.fill, self.fill_many = fill, fill_many | |
| self.window, self.max_batch, self.max_pending = window_s, max_batch, max_pending | |
| self.q: queue.Queue = queue.Queue() | |
| self.lock = threading.Lock() | |
| self.counts = {"jobs": 0, "batches": 0, "max_seen_batch": 0, "rejected": 0} | |
| threading.Thread(target=self._loop, daemon=True, name="deg-gpu").start() | |
| def stats(self) -> dict[str, Any]: | |
| with self.lock: | |
| out: dict[str, Any] = dict(self.counts) | |
| out["mean_batch"] = round(out["jobs"] / out["batches"], 2) if out["batches"] else None | |
| out["queued"] = self.q.qsize() | |
| return out | |
| def submit(self, context: Any, schema: dict[str, Any], options: FillOptions) -> Future: | |
| if self.q.qsize() >= self.max_pending: | |
| with self.lock: | |
| self.counts["rejected"] += 1 | |
| raise Overloaded("too many requests queued", 429) | |
| fut: Future = Future() | |
| self.q.put(((context, schema, options), fut)) | |
| return fut | |
| def _take(self) -> list: | |
| batch = [self.q.get()] | |
| deadline = time.perf_counter() + self.window | |
| while len(batch) < self.max_batch: | |
| left = deadline - time.perf_counter() | |
| if left <= 0: | |
| break | |
| try: | |
| batch.append(self.q.get(timeout=left)) | |
| except queue.Empty: | |
| break | |
| return batch | |
| def _run(self, batch: list) -> None: | |
| jobs = [job for job, _ in batch] | |
| if self.fill_many is not None and len(jobs) > 1: | |
| try: | |
| results = self.fill_many(jobs) | |
| except Exception: # noqa: BLE001 -- isolated: rerun each job alone below | |
| results = None | |
| if results is not None and len(results) == len(jobs): | |
| for (_, fut), res in zip(batch, results): | |
| if isinstance(res, BaseException): | |
| fut.set_exception(res) | |
| else: | |
| fut.set_result(res) | |
| return | |
| for job, fut in batch: | |
| try: | |
| fut.set_result(self.fill(*job)) | |
| except Exception as exc: # noqa: BLE001 -- delivered to that request only | |
| fut.set_exception(exc) | |
| def _loop(self) -> None: | |
| while True: | |
| batch = self._take() | |
| with self.lock: | |
| self.counts["jobs"] += len(batch) | |
| self.counts["batches"] += 1 | |
| self.counts["max_seen_batch"] = max(self.counts["max_seen_batch"], len(batch)) | |
| _batch_size.record(len(batch)) | |
| self._run(batch) | |
| def _trace_null_margin(options: FillOptions) -> None: | |
| """A request-set null margin on the request span (a span attribute only: no metric | |
| dimension, so its unbounded float values add no metric cardinality).""" | |
| if options.null_margin is not None: | |
| trace.get_current_span().set_attribute("deg.request.null_margin", | |
| float(options.null_margin)) | |
| class Engine: | |
| """Routes -> parsed jobs -> batcher -> API responses.""" | |
| def __init__(self, batcher: Batcher, *, model_id: str = MODEL_ID, | |
| openai_models: frozenset[str] = MODEL_ALIASES | OPENAI_COMPAT_MODELS, | |
| jev_models: frozenset[str] = MODEL_ALIASES | JEV_COMPAT_MODELS) -> None: | |
| self.batcher, self.model_id = batcher, model_id | |
| self.openai_models, self.jev_models = openai_models, jev_models | |
| def _fill(self, job: Job) -> dict[str, Any] | None: | |
| if job.schema is None: | |
| return None | |
| _trace_null_margin(job.options) | |
| return self.batcher.submit(job.context, job.schema, job.options).result() | |
| def decisions(self, payload: Any) -> tuple[dict[str, Any], Job]: | |
| job = parse_decisions(payload, self.openai_models) | |
| return decisions_response(job, self._fill(job), self.model_id), job | |
| def systemone(self, payload: Any) -> tuple[dict[str, Any], Job]: | |
| job = parse_systemone(payload, self.jev_models) | |
| return systemone_response(job, self._fill(job), self.model_id), job | |
| def fill(self, payload: Any) -> dict[str, Any]: | |
| req = parse_fill_request(payload) | |
| _trace_null_margin(req.options) | |
| return self.batcher.submit(req.context, payload["schema"], req.options).result() | |
| def models(self) -> dict[str, Any]: | |
| names = sorted(MODEL_ALIASES | self.openai_models | self.jev_models) | |
| return {"object": "list", "data": [ | |
| {"id": n, "object": "model", "created": 0, "owned_by": "deg", | |
| **({"root": self.model_id} if n != self.model_id else {})} for n in names]} | |
| class BodyError(ValueError): | |
| pass | |
| def read_body(stream: Any, headers: Any) -> bytes: | |
| """Whole body, Content-Length or chunked, read in bounded pieces (no size ceiling).""" | |
| body = bytearray() | |
| if "chunked" in (headers.get("transfer-encoding") or "").lower(): | |
| while True: | |
| line = stream.readline() | |
| try: | |
| size = int(line.split(b";", 1)[0].strip(), 16) | |
| except ValueError as error: | |
| raise BodyError("bad chunk size") from error | |
| if size == 0: | |
| while stream.readline() not in (b"\r\n", b"\n", b""): | |
| pass | |
| return bytes(body) | |
| while size: | |
| piece = stream.read(min(size, READ_PIECE_BYTES)) | |
| if not piece: | |
| raise BodyError("truncated chunk") | |
| body += piece | |
| size -= len(piece) | |
| stream.readline() | |
| try: | |
| remaining = int(headers.get("content-length") or 0) | |
| except ValueError as error: | |
| raise BodyError("bad content-length") from error | |
| while remaining > 0: | |
| piece = stream.read(min(remaining, READ_PIECE_BYTES)) | |
| if not piece: | |
| raise BodyError("truncated body") | |
| body += piece | |
| remaining -= len(piece) | |
| return bytes(body) | |
| def _route_error(route: str, status: int, message: str) -> tuple[int, dict[str, Any]]: | |
| """Error body in the route's own convention.""" | |
| if route == "systemone": | |
| kind = {401: "authentication_error", 403: "authentication_error", | |
| 429: "rate_limit_error", 529: "overloaded_error"}.get(status, "internal_error") | |
| if status < 500 and status not in (401, 403, 429): | |
| kind = "api_usage_error" | |
| return status, {"detail": {"error_type": kind, "message": message}} | |
| if route == "fill": | |
| kind = {401: "unauthorized", 404: "not_found", 422: "invalid_request", | |
| 429: "rate_limited", 503: "overloaded"}.get(status, "server_error") | |
| return status, {"error": {"type": kind, "status": status, "message": message}} | |
| kind = {401: "invalid_request_error", 404: "invalid_request_error", | |
| 400: "invalid_request_error", 429: "requests", 503: "server_error"}.get( | |
| status, "server_error") | |
| code = {401: "invalid_api_key", 429: "rate_limit_exceeded"}.get(status) | |
| return status, {"error": {"message": message, "type": kind, "param": None, "code": code}} | |
| def make_handler(engine: Engine, *, api_key: str | None = None): | |
| class Handler(BaseHTTPRequestHandler): | |
| protocol_version = "HTTP/1.1" | |
| def log_message(self, *args: Any) -> None: # no request content in logs | |
| return | |
| def _send(self, status: int, body: Any) -> None: | |
| data = json.dumps(body, ensure_ascii=False).encode("utf-8") | |
| self.send_response(status) | |
| self.send_header("content-type", "application/json") | |
| self.send_header("content-length", str(len(data))) | |
| self.end_headers() | |
| self.wfile.write(data) | |
| def _authorized(self) -> bool: | |
| return not api_key or self.headers.get("authorization", "") == f"Bearer {api_key}" | |
| def _handle(self, method: str) -> None: | |
| path = self.path.split("?", 1)[0].rstrip("/") | |
| route = ROUTES.get(path, "other") | |
| t0 = time.perf_counter() | |
| with _tracer.start_as_current_span("deg.request") as span: | |
| span.set_attribute("deg.route", route) | |
| status, body, outcome = self._dispatch(method, route, span) | |
| span.set_attribute("deg.status", outcome) | |
| _duration.record(time.perf_counter() - t0, {"route": route, "status": outcome}) | |
| self._send(status, body) | |
| def _dispatch(self, method: str, route: str, span: Any | |
| ) -> tuple[int, Any, str]: | |
| if not self._authorized(): | |
| if route == "systemone": | |
| return 403, {"detail": { | |
| "error_type": "authentication_error", | |
| "message": "Must supply an API key! Check your request and try " | |
| "again."}}, "unauthorized" | |
| return (*_route_error(route, 401, "Incorrect API key provided."), | |
| "unauthorized") | |
| if method == "GET": | |
| if route == "health": | |
| return 200, {"ok": True, "model": engine.model_id, | |
| "batcher": engine.batcher.stats()}, "ok" | |
| if route == "models": | |
| return 200, engine.models(), "ok" | |
| return (*_route_error("fill", 404, "unknown path"), "client_error") | |
| if route not in ("decisions", "systemone", "fill"): | |
| return (*_route_error("fill", 404, "unknown path"), "client_error") | |
| try: | |
| raw = read_body(self.rfile, self.headers) | |
| payload = json.loads(raw) if raw else None | |
| except (BodyError, json.JSONDecodeError, UnicodeDecodeError) as exc: | |
| self.close_connection = isinstance(exc, BodyError) | |
| status = 400 if route != "fill" else 422 | |
| if route == "systemone": | |
| status = 422 | |
| return (*_route_error(route, status, "body is not valid JSON"), | |
| "client_error") | |
| try: | |
| if route == "fill": | |
| body = engine.fill(payload) | |
| _tokens.add(body["usage"]["prompt_tokens"], | |
| {"route": route, "direction": "input"}) | |
| _tokens.add(body["usage"]["generated_tokens"], | |
| {"route": route, "direction": "output"}) | |
| return 200, body, "ok" | |
| body, job = (engine.decisions if route == "decisions" | |
| else engine.systemone)(payload) | |
| span.set_attribute("deg.questions", len(job.questions)) | |
| for kind in job.question_types: | |
| _questions.add(1, {"route": route, "question_type": kind}) | |
| usage = body["usage"] | |
| _tokens.add(usage["input_tokens"], {"route": route, "direction": "input"}) | |
| _tokens.add(usage["output_tokens"], {"route": route, "direction": "output"}) | |
| return 200, body, "ok" | |
| except ApiError as exc: | |
| return exc.status, exc.body, "client_error" | |
| except RequestError as exc: | |
| return exc.status, exc.body(), "client_error" | |
| except Overloaded as exc: | |
| return (*_route_error(route, exc.status, str(exc)), "overloaded") | |
| except FillError as exc: | |
| span.record_exception(exc) | |
| if route == "fill": | |
| return 500, {"error": {"type": "fill_error", "status": 500, | |
| "detail": exc.errors}}, "server_error" | |
| return (*_route_error(route, 500, f"fill failed: {exc}"), "server_error") | |
| except RuntimeError as exc: | |
| span.record_exception(exc) | |
| if "out of memory" in str(exc).lower(): | |
| return (*_route_error(route, 529 if route == "systemone" else 503, | |
| "model overloaded"), "overloaded") | |
| return (*_route_error(route, 500, f"{type(exc).__name__}: {exc}"), | |
| "server_error") | |
| except Exception as exc: # noqa: BLE001 -- reported, never a silent answer | |
| span.record_exception(exc) | |
| return (*_route_error(route, 500, f"{type(exc).__name__}: {exc}"), | |
| "server_error") | |
| def do_GET(self) -> None: # noqa: N802 | |
| self._handle("GET") | |
| def do_POST(self) -> None: # noqa: N802 | |
| self._handle("POST") | |
| return Handler | |
| class DegHTTPServer(ThreadingHTTPServer): | |
| request_queue_size = 256 | |
| daemon_threads = True | |
| def make_server(engine: Engine, host: str, port: int, *, api_key: str | None = None | |
| ) -> DegHTTPServer: | |
| return DegHTTPServer((host, port), make_handler(engine, api_key=api_key)) | |
| def model_fill(ckpt: str | None, device: str, generate: str | None = None, | |
| null_margin: float = 0.0) -> tuple[FillFn, FillManyFn | None]: | |
| """``fill`` (and ``fill_many`` when the runtime exposes one) bound to a loaded model. | |
| ``generate="base"``: x-generate fields decode with the frozen base weights (hybrid). | |
| ``null_margin``: the extract null margin of requests that do not set one.""" | |
| from deg.model import load_model | |
| return bind_model(load_model(ckpt, device=device, generate=generate), | |
| null_margin=null_margin) | |
| def with_null_margin(options: FillOptions, default: float) -> FillOptions: | |
| """``options`` with the server's default null margin when the request set none.""" | |
| if options.null_margin is not None: | |
| return options | |
| return dataclasses.replace(options, null_margin=float(default)) | |
| def bind_model(model: Any, null_margin: float = 0.0) -> tuple[FillFn, FillManyFn | None]: | |
| """``fill`` and ``fill_many`` (auto-detected on :mod:`deg.infer`) bound to ``model``; | |
| requests without ``options.null_margin`` get ``null_margin``.""" | |
| from deg import infer | |
| def fill(context: Any, schema: dict[str, Any], options: FillOptions) -> dict[str, Any]: | |
| return infer.fill(model, context, schema, with_null_margin(options, null_margin)) | |
| many = getattr(infer, "fill_many", None) | |
| fill_many = (lambda jobs: many(model, [(c, s, with_null_margin(o, null_margin)) | |
| for c, s, o in jobs])) if callable(many) else None | |
| return fill, fill_many | |
| def main(argv: list[str] | None = None) -> None: | |
| ap = argparse.ArgumentParser(description="deg HTTP server (decisions, systemone, fill)") | |
| ap.add_argument("--ckpt", help="deg checkpoint directory (omit: base model, untrained " | |
| "heads -- for smoke tests only)") | |
| ap.add_argument("--host", default="127.0.0.1") | |
| ap.add_argument("--port", type=int, default=8790) | |
| ap.add_argument("--device", default="cuda") | |
| ap.add_argument("--batch-window-ms", type=float, default=5.0) | |
| ap.add_argument("--max-batch", type=int, default=16) | |
| ap.add_argument("--max-pending", type=int, default=256) | |
| ap.add_argument("--generate", type=generate_arg, default=None, | |
| help="x-generate decoder: deg's own weights or the frozen base " | |
| "Qwen3.5 (hybrid; +~1 GB), or base+lora:<dir> = that hybrid with a " | |
| "GENERATE LoRA (deg.gentrain) merged into the base copy. " | |
| "Default: the checkpoint's generate_backbone, else self") | |
| ap.add_argument("--null-margin", type=float, default=0.0, | |
| help="extract null margin in nats for requests that do not set " | |
| "options.null_margin: a nullable/optional extract leaf is null when " | |
| "best span log-prob < null log-prob + margin (> 0 = fewer made-up " | |
| "answers, more omissions; 0 = argmax). See SPEC.md 12.1") | |
| ap.add_argument("--no-compat-models", action="store_true", | |
| help="accept only deg model names (not gpt-6-luna / jev-latest)") | |
| a = ap.parse_args(argv) | |
| if a.ckpt is None: | |
| print("warning: no --ckpt; serving the base model with untrained heads", | |
| file=sys.stderr, flush=True) | |
| if not math.isfinite(a.null_margin): | |
| ap.error("--null-margin must be finite") | |
| fill, fill_many = model_fill(a.ckpt, a.device, a.generate, a.null_margin) | |
| batcher = Batcher(fill, fill_many=fill_many, window_s=a.batch_window_ms / 1000, | |
| max_batch=a.max_batch, max_pending=a.max_pending) | |
| engine = (Engine(batcher, openai_models=MODEL_ALIASES, jev_models=MODEL_ALIASES) | |
| if a.no_compat_models else Engine(batcher)) | |
| server = make_server(engine, a.host, a.port, api_key=os.environ.get("DEG_API_KEY") or None) | |
| print(json.dumps({"serving": f"http://{a.host}:{a.port}", "model": MODEL_ID, | |
| "ckpt": a.ckpt, "generate": a.generate, "null_margin": a.null_margin, | |
| "fill_many": fill_many is not None}), flush=True) | |
| server.serve_forever() | |
| if __name__ == "__main__": | |
| main() | |