SolitaryThinker
21e59268d8
Sweep all of v2/ through pre-commit (was previously only run on changed files). Fixes surfaced across untouched modules: - ruff: isinstance-tuple -> X | Y (UP038), try/except/pass -> contextlib.suppress (SIM105), negated-return (SIM103), %-format -> f-string (UP031). - mypy: add annotations for no-untyped-call + var-annotated across recipes, training methods, torch/toy backends, and serving. - yapf reflow of the SF-Wan KV-cache call sites (semantics unchanged). yapf/ruff/codespell/mypy all pass; v2 tests 237 passed / 5 pre-existing (bucket_c torch-absent on CPU venv).
278 lines
10 KiB
Python
278 lines
10 KiB
Python
"""OmniOpenAIServer — OpenAI-compatible worker server (vllm-omni pattern).
|
|
|
|
Fronts ONE AsyncEngine (the per-pool worker surface). Routes OpenAI-shaped endpoints to typed
|
|
OmniRequests, streams chat via SSE, and serves video as an async job (POST returns a job id;
|
|
GET polls) plus a /sync variant — mirroring vllm-omni's `/v1/videos` + `/v1/videos/sync`. The fleet
|
|
(deploy/fleet.py) or Dynamo routes across multiple such servers/engines above this layer.
|
|
|
|
Unlike vllm-omni's request-scheduled opaque DIFFUSION stage, the engine underneath is step-scheduled;
|
|
the endpoint is a thin shim over the loop scheduler.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import contextlib
|
|
import itertools
|
|
import time
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from v2.request.artifacts import Output
|
|
from v2.request.streams import StreamChunk
|
|
from v2.serving.http import HttpServer, Request, Response
|
|
from v2.serving.protocol import ChatCompletionRequest, ImageGenerationRequest, VideoGenerationRequest
|
|
|
|
|
|
def _json_default(o: Any) -> Any:
|
|
"""Make SSE event payloads JSON-safe (media.chunk carries a StreamChunk with an array)."""
|
|
if isinstance(o, StreamChunk):
|
|
shape = list(np.asarray(o.data).shape) if o.data is not None else None
|
|
return {
|
|
"stream_id": o.stream_id,
|
|
"modality": o.modality,
|
|
"seq": o.seq,
|
|
"preview": o.preview,
|
|
"final": o.final,
|
|
"shape": shape
|
|
}
|
|
if isinstance(o, np.ndarray):
|
|
return {"shape": list(o.shape), "dtype": str(o.dtype)}
|
|
return str(o)
|
|
|
|
|
|
_job_ctr = itertools.count(1)
|
|
_chat_ctr = itertools.count(1)
|
|
|
|
|
|
def _payload(art: Any) -> Any:
|
|
for attr in ("text", "frames", "samples", "tensor", "latent", "token_ids"):
|
|
if hasattr(art, attr) and getattr(art, attr) is not None:
|
|
return getattr(art, attr)
|
|
return None
|
|
|
|
|
|
def _ser_artifacts(out: Output) -> dict:
|
|
res: dict = {}
|
|
for name, art in out.artifacts.items():
|
|
p = _payload(art)
|
|
if isinstance(p, str):
|
|
res[name] = {"type": "text", "text": p}
|
|
elif p is None:
|
|
res[name] = {"type": "empty"}
|
|
else:
|
|
a = np.asarray(p)
|
|
res[name] = {
|
|
"type": "tensor",
|
|
"shape": list(a.shape),
|
|
"dtype": str(a.dtype),
|
|
"b64_json": base64.b64encode(np.ascontiguousarray(a).tobytes()).decode()
|
|
}
|
|
return res
|
|
|
|
|
|
class OmniOpenAIServer:
|
|
|
|
def __init__(self, engine: Any, *, engine_id: str = "v2", max_jobs: int = 512):
|
|
self.engine = engine # an AsyncEngine
|
|
self.engine_id = engine_id
|
|
self.jobs: dict[str, dict] = {} # video job store
|
|
self.http = HttpServer(self.dispatch)
|
|
self.served = 0
|
|
self.max_jobs = max_jobs
|
|
self._tasks: set = set() # tracked video-job tasks (no fire-and-forget leak)
|
|
|
|
def _evict_jobs(self) -> None:
|
|
terminal = ("completed", "failed")
|
|
for jid in list(self.jobs):
|
|
if len(self.jobs) <= self.max_jobs:
|
|
break
|
|
if self.jobs[jid].get("status") in terminal:
|
|
self.jobs.pop(jid, None)
|
|
|
|
# --- routing ------------------------------------------------------------ #
|
|
async def dispatch(self, req: Request) -> Response:
|
|
m, p = req.method, req.path
|
|
if m == "GET" and p == "/health":
|
|
return Response.json(self._health())
|
|
if m == "GET" and p == "/metrics":
|
|
return Response.json(self._metrics())
|
|
if m == "GET" and p == "/v1/models":
|
|
return Response.json(self._models())
|
|
if m == "POST" and p == "/v1/images/generations":
|
|
return await self._images(req)
|
|
if m == "POST" and p == "/v1/videos":
|
|
return await self._video_create(req)
|
|
if m == "POST" and p == "/v1/videos/sync":
|
|
return await self._video_sync(req)
|
|
if m == "GET" and p.startswith("/v1/videos/"):
|
|
return self._video_get(p.rsplit("/", 1)[-1])
|
|
if m == "POST" and p == "/v1/chat/completions":
|
|
return await self._chat(req)
|
|
return Response.json({"error": f"no route for {m} {p}"}, status=404)
|
|
|
|
# --- metadata endpoints ------------------------------------------------- #
|
|
def _health(self) -> dict:
|
|
return {
|
|
"status": "healthy",
|
|
"engine_id": self.engine_id,
|
|
"in_flight": self.engine.in_flight,
|
|
"queue_depth": self.engine.queue_depth
|
|
}
|
|
|
|
def _metrics(self) -> dict:
|
|
sm = self.engine.engine.metrics
|
|
adm = self.engine.engine.admission.metrics
|
|
return {
|
|
"served": self.served,
|
|
"in_flight": self.engine.in_flight,
|
|
"stepped_units": sm.stepped_units,
|
|
"batches": sm.batches,
|
|
"admitted": adm.admitted,
|
|
"deferred": adm.deferred,
|
|
"gpu_seconds": round(adm.gpu_seconds, 6),
|
|
"by_kind": dict(adm.by_kind)
|
|
}
|
|
|
|
def _models(self) -> dict:
|
|
ids = (list(self.engine.engine._registry) + list(self.engine._disagg) + list(self.engine.engine._workflows)
|
|
) # workflows are servables too
|
|
return {"object": "list", "data": [{"id": m, "object": "model"} for m in ids]}
|
|
|
|
# --- images ------------------------------------------------------------- #
|
|
async def _images(self, req: Request) -> Response:
|
|
ig = ImageGenerationRequest.from_json(req.json())
|
|
out = await self.engine.generate(ig.to_omni())
|
|
self.served += 1
|
|
arts = _ser_artifacts(out)
|
|
primary = arts.get("image") or arts.get("video") or arts.get("latents")
|
|
return Response.json({
|
|
"created": int(time.time()),
|
|
"model": ig.model,
|
|
"data": [primary] * max(1, ig.n),
|
|
"omni": arts,
|
|
"metrics": {
|
|
k: round(v, 6)
|
|
for k, v in out.metrics.items()
|
|
}
|
|
})
|
|
|
|
# --- video: async job + poll + sync ------------------------------------- #
|
|
async def _video_create(self, req: Request) -> Response:
|
|
vg = VideoGenerationRequest.from_json(req.json())
|
|
job_id = f"video-{next(_job_ctr)}"
|
|
self.jobs[job_id] = {"id": job_id, "status": "queued", "model": vg.model}
|
|
t = asyncio.create_task(self._run_video_job(job_id, vg))
|
|
self._tasks.add(t)
|
|
t.add_done_callback(self._tasks.discard) # tracked, not fire-and-forget
|
|
self._evict_jobs()
|
|
return Response.json({"id": job_id, "status": "queued", "model": vg.model}, status=202)
|
|
|
|
async def _run_video_job(self, job_id: str, vg: VideoGenerationRequest) -> None:
|
|
self.jobs[job_id]["status"] = "running"
|
|
try:
|
|
out = await self.engine.generate(vg.to_omni())
|
|
self.served += 1
|
|
self.jobs[job_id].update(status="completed",
|
|
result=_ser_artifacts(out),
|
|
metrics={
|
|
k: round(v, 6)
|
|
for k, v in out.metrics.items()
|
|
})
|
|
except Exception as e: # noqa: BLE001
|
|
self.jobs[job_id].update(status="failed", error=str(e))
|
|
|
|
def _video_get(self, job_id: str) -> Response:
|
|
job = self.jobs.get(job_id)
|
|
if job is None:
|
|
return Response.json({"error": f"no such job {job_id}"}, status=404)
|
|
return Response.json(job)
|
|
|
|
async def _video_sync(self, req: Request) -> Response:
|
|
vg = VideoGenerationRequest.from_json(req.json())
|
|
out = await self.engine.generate(vg.to_omni())
|
|
self.served += 1
|
|
return Response.json({
|
|
"model": vg.model,
|
|
"data": _ser_artifacts(out),
|
|
"metrics": {
|
|
k: round(v, 6)
|
|
for k, v in out.metrics.items()
|
|
}
|
|
})
|
|
|
|
# --- chat: SSE stream or JSON ------------------------------------------- #
|
|
async def _chat(self, req: Request) -> Response:
|
|
cc = ChatCompletionRequest.from_json(req.json())
|
|
omni = cc.to_omni()
|
|
if cc.stream:
|
|
return Response.sse(self._chat_sse(omni, cc.model))
|
|
out = await self.engine.generate(omni)
|
|
self.served += 1
|
|
arts = _ser_artifacts(out)
|
|
content = arts.get("text", {}).get("text", "") if "text" in arts else ""
|
|
return Response.json({
|
|
"id":
|
|
f"chatcmpl-{next(_chat_ctr)}",
|
|
"object":
|
|
"chat.completion",
|
|
"created":
|
|
int(time.time()),
|
|
"model":
|
|
cc.model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": content
|
|
}
|
|
}],
|
|
"omni":
|
|
arts,
|
|
"metrics": {
|
|
k: round(v, 6)
|
|
for k, v in out.metrics.items()
|
|
}
|
|
})
|
|
|
|
async def _chat_sse(self, omni: Any, model: str):
|
|
import json
|
|
cid = f"chatcmpl-{next(_chat_ctr)}"
|
|
async for ev in self.engine.submit(omni): # live OmniEvents → SSE chunks
|
|
chunk = {"id": cid, "object": "chat.completion.chunk", "model": model, "event": ev.type, "data": ev.payload}
|
|
yield f"data: {json.dumps(chunk, default=_json_default)}\n\n"
|
|
self.served += 1
|
|
out = self.engine.result(omni.request_id)
|
|
arts = _ser_artifacts(out) if out is not None else {}
|
|
content = arts.get("text", {}).get("text", "") if "text" in arts else ""
|
|
final = {
|
|
"id": cid,
|
|
"object": "chat.completion.chunk",
|
|
"model": model,
|
|
"choices": [{
|
|
"index": 0,
|
|
"delta": {
|
|
"role": "assistant",
|
|
"content": content
|
|
},
|
|
"finish_reason": "stop"
|
|
}],
|
|
"omni": arts
|
|
}
|
|
yield f"data: {json.dumps(final, default=_json_default)}\n\n"
|
|
yield "data: [DONE]\n\n"
|
|
|
|
# --- lifecycle ---------------------------------------------------------- #
|
|
async def serve(self, host: str = "127.0.0.1", port: int = 0) -> tuple[str, int]:
|
|
return await self.http.start(host, port)
|
|
|
|
async def close(self) -> None:
|
|
for t in list(self._tasks): # cancel + drain in-flight video jobs on shutdown
|
|
t.cancel()
|
|
for t in list(self._tasks):
|
|
with contextlib.suppress(asyncio.CancelledError, Exception):
|
|
await t
|
|
await self.http.close()
|