SolitaryThinker
71caa58ba9
Review confirmed the core bets (concurrent disaggregation is bit-identical, design conformance holds, Dynamo genuinely optional, engine stays step-scheduled). Fixes for the untested failure paths: - HIGH: pool capacity (RolePool.in_flight) no longer leaks when a disaggregated request is cancelled or errors mid-occupancy — DisaggregatedRunner.close() releases the occupied pool and AsyncEngine._run calls it in a finally (a cancel on a capacity-1 denoiser no longer bricks the pool). - credit flow-control: cross-pool transfer wraps acquire/release in try/finally (no credit leak on a failing transfer); slot is re-homed only on a successful fetch. - AsyncEngine: duplicate in-flight request_id is rejected (was a deadlock); submit() cancels the driver task when the consumer abandons the stream (client disconnect → no orphaned compute); bounded per-request history (no unbounded _events/_states/_results/_runners growth). - cancellation is common-path on the OFFLINE path too (cancel check at the top of every runner.tick()). - HTTP server: read timeout (slowloris guard → 408), body-size cap (→ 413), invalid Content-Length (→ 400), explicit StreamReader limit, and aclose() of the SSE generator on client disconnect. - build_deployment_card no longer aliases one mutable CostModel across replica cards (dataclasses.replace), so online calibration of one worker's cost doesn't mutate another's. - video-job tasks tracked (not fire-and-forget); server.close() cancels/drains them; jobs dict bounded; fleet affinity map bounded. 6 regression tests added for these paths. 82 tests pass (pytest + zero-dep runner).
211 lines
9.7 KiB
Python
211 lines
9.7 KiB
Python
"""OmniOpenAIServer — our own OpenAI-compatible worker server (design.md §6.3.5; 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; the runtime-visible loop scheduler is the substance.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import itertools
|
|
import time
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from ..request.artifacts import Output
|
|
from ..request.streams import StreamChunk
|
|
from .http import HttpServer, Request, Response
|
|
|
|
|
|
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)
|
|
from .protocol import ChatCompletionRequest, ImageGenerationRequest, VideoGenerationRequest
|
|
|
|
_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 = "mini-fastvideo", 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)
|
|
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):
|
|
try:
|
|
await t
|
|
except (asyncio.CancelledError, Exception): # noqa: BLE001
|
|
pass
|
|
await self.http.close()
|