SolitaryThinker
58223c0c41
Answers "what's the right way to name/register custom pipelines like T2I→I2V": treat a Workflow like a card — a stable namespaced id in the same servable namespace, declared dependencies, and a two-level registry. No new concepts; mirrors how cards are registered (and vllm-omni's pipeline_registry). - program/workflow.py: Workflow gains `requires` (the cards it composes, derived from stages) and `validate(engine)` (fail-fast if a required card is absent, P7). New WorkflowRegistry: declarative workflow_id -> builder catalog for out-of-tree/ad hoc use. - runtime/engine.py: `_workflows` registry + register_workflow (validates deps, rejects id collision with a model_id) + serves(); engine.run routes a request whose model_id is a workflow to workflow.run — addressable exactly like a model. Single-model hot path untouched. - runtime/async_engine.py + serving/server.py: serves() and /models include workflows (discoverable as servables). - models/__init__.py: declarative `_WORKFLOWS` catalog (the cross-model analog of _BUILDERS) + register_workflows() helper; build_image_video_engine now registers the workflow too. Adding a custom pipeline = one catalog line. - Naming convention: dotted/namespaced workflow_id (`image_video.t2i_i2v`), distinct from kebab model ids, collision-checked. Renamed from `t2i_then_i2v`. - tests (+5, 12 total in the file): addressable by id, requires/validate, id collision, registry catalog, register_workflows helper. Full suite 114 passed. - designv4 §9.6: the naming & registration convention documented.
218 lines
9.5 KiB
Python
218 lines
9.5 KiB
Python
"""AsyncEngine — the serving runtime (design_v3 §6; plan M3).
|
|
|
|
A real request queue, lifecycle state machine, live streaming, and step-level concurrency over the
|
|
same step scheduler. Each request runs as an asyncio task that ticks its runner one step at a time
|
|
and yields the event loop (`await asyncio.sleep(0)`), so concurrent requests interleave at step
|
|
granularity — and a per-request stall backs off instead of busy-spinning. Drives BOTH inline
|
|
``ProgramRunner`` and ``DisaggregatedRunner`` (capacity-aware pool dispatch) through one driver.
|
|
|
|
Cancellation is common-path (design_v3 §6.4): ``cancel(request_id)`` trips the runner's CancelScope,
|
|
which raises at the next step boundary and delivers partial artifacts + a structured ``cancelled``.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any, AsyncIterator
|
|
|
|
from ..request.artifacts import Output
|
|
from ..request.cancel import Cancelled
|
|
from ..request.streams import OmniEvent
|
|
from .disaggregated import DisaggregatedRunner
|
|
from .engine import Engine
|
|
|
|
_SENTINEL = object()
|
|
|
|
|
|
def _safe_put(q: "asyncio.Queue", item: Any) -> None:
|
|
"""Non-blocking put on an unbounded queue — safe to call from finally/cancellation paths."""
|
|
try:
|
|
q.put_nowait(item)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class RequestState:
|
|
WAITING = "waiting"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
CANCELLED = "cancelled"
|
|
FAILED = "failed"
|
|
|
|
|
|
class AsyncEngine:
|
|
def __init__(self, engine: Engine | None = None, *, max_stall_ticks: int = 200_000,
|
|
max_history: int = 512):
|
|
self.engine = engine if engine is not None else Engine()
|
|
self._disagg: dict[str, tuple] = {} # model_id -> (PoolSet, Program)
|
|
self._events: dict[str, asyncio.Queue] = {}
|
|
self._states: dict[str, str] = {}
|
|
self._results: dict[str, Output] = {}
|
|
self._runners: dict[str, Any] = {}
|
|
self._order: list[str] = [] # submission order, for bounded eviction
|
|
self.max_stall_ticks = max_stall_ticks
|
|
self.max_history = max_history
|
|
|
|
def _evict(self) -> None:
|
|
"""Bound retained per-request state to the most recent ``max_history`` terminal requests."""
|
|
terminal = (RequestState.COMPLETED, RequestState.CANCELLED, RequestState.FAILED)
|
|
i = 0
|
|
while len(self._order) > self.max_history and i < len(self._order):
|
|
rid = self._order[i]
|
|
if self._states.get(rid) in terminal:
|
|
self._order.pop(i)
|
|
self._events.pop(rid, None)
|
|
self._states.pop(rid, None)
|
|
self._results.pop(rid, None)
|
|
self._runners.pop(rid, None)
|
|
else:
|
|
i += 1
|
|
|
|
# --- registration ------------------------------------------------------- #
|
|
def register(self, model_id: str, instance: Any, program: Any) -> None:
|
|
self.engine.register(model_id, instance, program)
|
|
|
|
def register_disaggregated(self, model_id: str, pools: Any, program: Any) -> None:
|
|
self._disagg[model_id] = (pools, program)
|
|
|
|
def register_workflow(self, workflow: Any) -> Any:
|
|
return self.engine.register_workflow(workflow)
|
|
|
|
def is_disaggregated(self, model_id: str) -> bool:
|
|
return model_id in self._disagg
|
|
|
|
def _make_runner(self, request: Any) -> Any:
|
|
if request.model_id in self._disagg:
|
|
pools, program = self._disagg[request.model_id]
|
|
return DisaggregatedRunner(pools, program, request,
|
|
observers=self.engine.observers, interceptors=self.engine.interceptors)
|
|
return self.engine._make_runner(request)
|
|
|
|
# --- lifecycle / control ------------------------------------------------ #
|
|
def state(self, request_id: str) -> str | None:
|
|
return self._states.get(request_id)
|
|
|
|
def cancel(self, request_id: str) -> bool:
|
|
runner = self._runners.get(request_id)
|
|
if runner is not None and getattr(runner, "cancel_scope", None) is not None:
|
|
runner.cancel_scope.cancel()
|
|
return True
|
|
return False
|
|
|
|
def result(self, request_id: str) -> Output | None:
|
|
return self._results.get(request_id)
|
|
|
|
@property
|
|
def in_flight(self) -> int:
|
|
return sum(1 for s in self._states.values() if s == RequestState.RUNNING)
|
|
|
|
@property
|
|
def queue_depth(self) -> int:
|
|
return sum(1 for s in self._states.values() if s == RequestState.WAITING)
|
|
|
|
def serves(self, model_id: str) -> bool:
|
|
return (model_id in self.engine._registry or model_id in self._disagg
|
|
or model_id in self.engine._workflows)
|
|
|
|
# --- the driver: one asyncio task per request --------------------------- #
|
|
async def _run(self, request: Any) -> None:
|
|
rid = request.request_id
|
|
evq = self._events[rid]
|
|
runner = None
|
|
stream_cursor = 0
|
|
stall = 0
|
|
try:
|
|
runner = self._make_runner(request)
|
|
self._runners[rid] = runner
|
|
self._states[rid] = RequestState.RUNNING
|
|
await evq.put(OmniEvent("request.start", rid, payload={"task": request.task.value}))
|
|
while not runner.done:
|
|
runner.cancel_scope.check() # common-path cancellation (step boundary)
|
|
p0 = runner._progress
|
|
runner.tick()
|
|
# surface any newly emitted media chunks as live events
|
|
events = runner.stream.events
|
|
while stream_cursor < len(events):
|
|
await evq.put(events[stream_cursor])
|
|
stream_cursor += 1
|
|
if not runner.done and runner._progress == p0:
|
|
stall += 1
|
|
if stall > self.max_stall_ticks:
|
|
raise RuntimeError(f"request {rid} stalled (state={getattr(runner,'state','?')})")
|
|
await asyncio.sleep(0.0005) # backpressure backoff (capacity wait)
|
|
else:
|
|
stall = 0
|
|
await asyncio.sleep(0) # cooperative yield → step interleaving
|
|
out = runner.output()
|
|
self._results[rid] = out
|
|
for name in out.artifacts:
|
|
await evq.put(OmniEvent("artifact.ready", rid, payload={"name": name}))
|
|
self._states[rid] = RequestState.COMPLETED
|
|
await evq.put(OmniEvent("request.complete", rid, payload={"metrics": dict(out.metrics)}))
|
|
except Cancelled:
|
|
self._states[rid] = RequestState.CANCELLED
|
|
if runner is not None and hasattr(runner, "output"):
|
|
try:
|
|
partial = runner.output()
|
|
partial.error = "cancelled"
|
|
self._results[rid] = partial
|
|
except Exception:
|
|
pass
|
|
_safe_put(evq, OmniEvent("request.cancelled", rid))
|
|
except asyncio.CancelledError: # driver task cancelled (e.g. client gone)
|
|
self._states[rid] = RequestState.CANCELLED
|
|
raise
|
|
except Exception as e: # noqa: BLE001 — request-fatal isolation: one request's error
|
|
self._states[rid] = RequestState.FAILED
|
|
self._results[rid] = Output(request_id=rid, error=str(e))
|
|
_safe_put(evq, OmniEvent("request.failed", rid, payload={"error": str(e)}))
|
|
finally:
|
|
# release any pool this request occupied — a cancel/error mid-loop must NOT leak capacity
|
|
if runner is not None and hasattr(runner, "close"):
|
|
try:
|
|
runner.close()
|
|
except Exception:
|
|
pass
|
|
_safe_put(evq, _SENTINEL)
|
|
self._evict()
|
|
|
|
async def submit(self, request: Any) -> AsyncIterator[OmniEvent]:
|
|
"""Submit a request and stream its events live (request.start → media.chunk* → artifact.ready*
|
|
→ request.complete | request.cancelled | request.failed). Cancels the driver if the consumer
|
|
abandons the stream (e.g. client disconnect), so no orphaned compute keeps running."""
|
|
rid = request.request_id
|
|
if rid in self._events and self._states.get(rid) in (RequestState.WAITING, RequestState.RUNNING):
|
|
raise ValueError(f"request_id {rid!r} is already in flight")
|
|
self._events[rid] = asyncio.Queue()
|
|
self._states[rid] = RequestState.WAITING
|
|
self._order.append(rid)
|
|
task = asyncio.create_task(self._run(request))
|
|
evq = self._events[rid]
|
|
try:
|
|
while True:
|
|
ev = await evq.get()
|
|
if ev is _SENTINEL:
|
|
break
|
|
yield ev
|
|
finally:
|
|
if not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except (asyncio.CancelledError, Exception): # noqa: BLE001
|
|
pass
|
|
|
|
async def generate(self, request: Any) -> Output:
|
|
"""Non-streaming: drive to completion, return the Output (the offline path over the queue)."""
|
|
async for _ in self.submit(request):
|
|
pass
|
|
return self._results[request.request_id]
|
|
|
|
async def generate_many(self, requests: list[Any]) -> dict[str, Output]:
|
|
"""Run requests concurrently (they interleave at step granularity through the event loop)."""
|
|
await asyncio.gather(*(self._drain(r) for r in requests))
|
|
return {r.request_id: self._results[r.request_id] for r in requests}
|
|
|
|
async def _drain(self, request: Any) -> None:
|
|
async for _ in self.submit(request):
|
|
pass
|