SolitaryThinker
8240970ea6
Builds the full serving layer the design files specify, instead of deferring it to Dynamo: - transport/ (§7.3): pluggable Connectors (in-proc zero-copy / SHM-fake copy) with chunk_ready readiness (vllm-omni) AND credit-based flow control (sglang-omni Relay); KVConnector protocol shape; TransferManifest. - runtime/ (§6, §13; plan M3/M4): AsyncEngine — request queue, lifecycle state machine (waiting→running→completed/cancelled/failed), live AsyncIterator[OmniEvent] streaming, common-path cancellation, step-level concurrency. RolePool + DisaggregatedRunner (encoder→denoiser→decoder, capacity-aware dispatch, cross-pool transfers via connectors); disaggregated output is bit-identical to inline. No-progress detection (no busy-spin). - deploy/ (§14, §6.3.5-6): DeploymentCard; OUR OWN LocalFleet (discovery, health/drain, least-loaded / cost-model / sticky-affinity routing) so we never rely on Dynamo; DynamoWorkerAdapter + FakeDynamoRuntime export the SAME card + cost model so Dynamo CAN front us — one object, two consumers. - serving/ (§6.3.5, §12): framework-free stdlib-asyncio OpenAI server (our own version of the vllm-omni pattern): /v1/chat/completions (SSE), /v1/images/generations, /v1/videos (async job+poll) + /v1/videos/sync, /v1/models, /health, /metrics. A thin shim over the STEP-scheduled engine — the runtime-visible loop scheduler vllm-omni's request-scheduled opaque DIFFUSION stage lacks. 15 serving tests (real-socket HTTP+SSE via stdlib asyncio, disagg==inline, fleet routing, Dynamo contract, cancellation); 76 tests pass total (pytest + zero-dep runner). ~7900 LOC.
118 lines
4.6 KiB
Python
118 lines
4.6 KiB
Python
"""OpenAI-shaped request schemas (design_v3 §12; the vllm-omni protocol, framework-free).
|
|
|
|
Plain dataclasses parsed from JSON (no pydantic dep). Each translates to a typed ``OmniRequest`` at
|
|
the server boundary — "the request is the only currency crossing the product boundary" (§12).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from ..request import DiffusionParams, OutputSpec, SamplingParams, TaskType, make_request
|
|
|
|
|
|
def _wh(size: str, default=(832, 480)) -> tuple[int, int]:
|
|
try:
|
|
w, h = size.lower().split("x")
|
|
return int(w), int(h)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _task_for(model_id: str, kind: str) -> TaskType:
|
|
m = model_id.lower()
|
|
if kind == "image":
|
|
return TaskType.T2I if ("bagel" in m or "image" in m) else TaskType.T2V
|
|
if kind == "video":
|
|
return TaskType.T2V
|
|
if kind == "chat":
|
|
if "bagel" in m:
|
|
return TaskType.T2I
|
|
if "cosmos" in m or "omni" in m:
|
|
return TaskType.T2V
|
|
return TaskType.T2V
|
|
return TaskType.T2V
|
|
|
|
|
|
@dataclass
|
|
class ImageGenerationRequest:
|
|
prompt: str
|
|
model: str
|
|
n: int = 1
|
|
size: str = "832x480"
|
|
num_inference_steps: int = 4
|
|
guidance_scale: float = 5.0
|
|
seed: int = 0
|
|
|
|
@classmethod
|
|
def from_json(cls, d: dict) -> "ImageGenerationRequest":
|
|
return cls(prompt=d.get("prompt", ""), model=d.get("model", ""), n=int(d.get("n", 1)),
|
|
size=d.get("size", "832x480"), num_inference_steps=int(d.get("num_inference_steps", 4)),
|
|
guidance_scale=float(d.get("guidance_scale", 5.0)), seed=int(d.get("seed", 0)))
|
|
|
|
def to_omni(self):
|
|
w, h = _wh(self.size)
|
|
return make_request(_task_for(self.model, "image"), self.model, self.prompt,
|
|
diffusion=DiffusionParams(num_steps=self.num_inference_steps,
|
|
guidance_scale=self.guidance_scale,
|
|
height=h, width=w, seed=self.seed))
|
|
|
|
|
|
@dataclass
|
|
class VideoGenerationRequest:
|
|
prompt: str
|
|
model: str
|
|
seconds: float = 2.0
|
|
size: str = "832x480"
|
|
num_frames: int = 81
|
|
num_inference_steps: int = 4
|
|
guidance_scale: float = 5.0
|
|
seed: int = 0
|
|
stream: dict = field(default_factory=dict)
|
|
|
|
@classmethod
|
|
def from_json(cls, d: dict) -> "VideoGenerationRequest":
|
|
return cls(prompt=d.get("prompt", ""), model=d.get("model", ""),
|
|
seconds=float(d.get("seconds", 2.0)), size=d.get("size", "832x480"),
|
|
num_frames=int(d.get("num_frames", 81)),
|
|
num_inference_steps=int(d.get("num_inference_steps", 4)),
|
|
guidance_scale=float(d.get("guidance_scale", 5.0)), seed=int(d.get("seed", 0)),
|
|
stream={"video": True} if d.get("stream") else {})
|
|
|
|
def to_omni(self):
|
|
w, h = _wh(self.size)
|
|
return make_request(TaskType.T2V, self.model, self.prompt,
|
|
diffusion=DiffusionParams(num_steps=self.num_inference_steps,
|
|
guidance_scale=self.guidance_scale, height=h, width=w,
|
|
num_frames=self.num_frames, seed=self.seed),
|
|
outputs=OutputSpec(modalities=frozenset({"video"}), stream=self.stream))
|
|
|
|
|
|
@dataclass
|
|
class ChatCompletionRequest:
|
|
model: str
|
|
messages: list[dict]
|
|
stream: bool = False
|
|
max_tokens: int = 6
|
|
num_inference_steps: int = 4
|
|
seed: int = 0
|
|
|
|
@classmethod
|
|
def from_json(cls, d: dict) -> "ChatCompletionRequest":
|
|
return cls(model=d.get("model", ""), messages=d.get("messages", []),
|
|
stream=bool(d.get("stream", False)), max_tokens=int(d.get("max_tokens", 6)),
|
|
num_inference_steps=int(d.get("num_inference_steps", 4)), seed=int(d.get("seed", 0)))
|
|
|
|
def prompt(self) -> str:
|
|
for m in reversed(self.messages):
|
|
if m.get("role") == "user":
|
|
c = m.get("content", "")
|
|
return c if isinstance(c, str) else " ".join(p.get("text", "") for p in c if isinstance(p, dict))
|
|
return ""
|
|
|
|
def to_omni(self):
|
|
return make_request(_task_for(self.model, "chat"), self.model, self.prompt(),
|
|
sampling=SamplingParams(max_tokens=self.max_tokens, seed=self.seed),
|
|
diffusion=DiffusionParams(num_steps=self.num_inference_steps, seed=self.seed),
|
|
outputs=OutputSpec(stream={"video": True} if self.stream else {}))
|