diff --git a/mini_fastvideo/deploy/card.py b/mini_fastvideo/deploy/card.py index ac3e5ba27..373fba43a 100644 --- a/mini_fastvideo/deploy/card.py +++ b/mini_fastvideo/deploy/card.py @@ -8,6 +8,7 @@ This is the contract our OWN fleet (``deploy/fleet.py``) consumes AND the Dynamo """ from __future__ import annotations +import dataclasses from dataclasses import dataclass, field from typing import Any @@ -59,7 +60,9 @@ def build_deployment_card(engine_id: str, model_cards: list, *, max_concurrent: if cost_model is None: for lp in c.loops.values(): if lp.step_cost_model is not None: - cost_model = lp.step_cost_model + # own copy so online calibration of one card's cost doesn't alias another replica's + cost_model = dataclasses.replace(lp.step_cost_model, + coefficients=dict(lp.step_cost_model.coefficients)) break return DeploymentCard( engine_id=engine_id, model_cards=ids, capabilities=frozenset(caps), diff --git a/mini_fastvideo/deploy/fleet.py b/mini_fastvideo/deploy/fleet.py index 740fae34c..b399f4fcf 100644 --- a/mini_fastvideo/deploy/fleet.py +++ b/mini_fastvideo/deploy/fleet.py @@ -53,11 +53,12 @@ class Worker: class LocalFleet: - def __init__(self, policy: str = "least_loaded"): + def __init__(self, policy: str = "least_loaded", *, max_affinity: int = 100_000): assert policy in ("least_loaded", "cost", "affinity") self.policy = policy + self.max_affinity = max_affinity self.workers: dict[str, Worker] = {} - self._affinity: dict[str, str] = {} # affinity key -> worker_id (sticky) + self._affinity: dict[str, str] = {} # affinity key -> worker_id (sticky), FIFO-bounded # --- discovery / health (what Dynamo's registry + planner would do) ------ # def register(self, worker_id: str, engine: Any, card: DeploymentCard) -> Worker: @@ -89,6 +90,8 @@ class LocalFleet: if wid in self.workers and self.workers[wid] in cands: return self.workers[wid] chosen = min(cands, key=lambda w: w.load) # cold key → least-loaded, then pin + if len(self._affinity) >= self.max_affinity: # FIFO-bound the sticky map + self._affinity.pop(next(iter(self._affinity)), None) self._affinity[key] = chosen.worker_id return chosen if self.policy == "cost": diff --git a/mini_fastvideo/runtime/async_engine.py b/mini_fastvideo/runtime/async_engine.py index 13c6b41bf..36e003cd3 100644 --- a/mini_fastvideo/runtime/async_engine.py +++ b/mini_fastvideo/runtime/async_engine.py @@ -23,6 +23,14 @@ 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" @@ -32,14 +40,32 @@ class RequestState: class AsyncEngine: - def __init__(self, engine: Engine | None = None, *, max_stall_ticks: int = 200_000): + 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: @@ -87,13 +113,14 @@ class AsyncEngine: async def _run(self, request: Any) -> None: rid = request.request_id evq = self._events[rid] - 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})) + 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 @@ -119,24 +146,41 @@ class AsyncEngine: await evq.put(OmniEvent("request.complete", rid, payload={"metrics": dict(out.metrics)})) except Cancelled: self._states[rid] = RequestState.CANCELLED - partial = runner.output() if hasattr(runner, "output") else None - if partial is not None: - partial.error = "cancelled" - self._results[rid] = partial - await evq.put(OmniEvent("request.cancelled", rid)) + 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)) - await evq.put(OmniEvent("request.failed", rid, payload={"error": str(e)})) + _safe_put(evq, OmniEvent("request.failed", rid, payload={"error": str(e)})) finally: - await evq.put(_SENTINEL) + # 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).""" + → 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: @@ -146,7 +190,12 @@ class AsyncEngine: break yield ev finally: - await task + 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).""" diff --git a/mini_fastvideo/runtime/disaggregated.py b/mini_fastvideo/runtime/disaggregated.py index 2f06cdfb0..af58cca50 100644 --- a/mini_fastvideo/runtime/disaggregated.py +++ b/mini_fastvideo/runtime/disaggregated.py @@ -71,17 +71,19 @@ class DisaggregatedRunner: if holder_id and holder_id != pool.pool_id and r in self.slots: src = self.pools.by_id[holder_id] got = src.connector.acquire_credit() # credit-based flow control - manifest = TransferManifest(keys=(r,), producer_id=holder_id, consumer_id=pool.pool_id, - src_location=src.connector.name, dst_location=pool.connector.name, - nbytes=payload_nbytes(self.slots.get(r))) - src.connector.put(r, self.slots.get(r), manifest) # producer publishes (chunk_ready) - if src.connector.chunk_ready(r): # readiness signal - self.slots[r] = src.connector.take(r) # consumer fetches - if got: - src.connector.release_credit() - self.slot_pool[r] = pool.pool_id - self.transfers += 1 - self.metrics["transfers"] = self.metrics.get("transfers", 0) + 1 + try: + manifest = TransferManifest(keys=(r,), producer_id=holder_id, consumer_id=pool.pool_id, + src_location=src.connector.name, dst_location=pool.connector.name, + nbytes=payload_nbytes(self.slots.get(r))) + src.connector.put(r, self.slots.get(r), manifest) # producer publishes (chunk_ready) + if src.connector.chunk_ready(r): # readiness signal + self.slots[r] = src.connector.take(r) # consumer fetches + self.slot_pool[r] = pool.pool_id # re-home ONLY on a successful fetch + self.transfers += 1 + self.metrics["transfers"] = self.metrics.get("transfers", 0) + 1 + finally: + if got: # never leak a credit on an error path + src.connector.release_credit() def _commit_node_writes(self, node, pool) -> None: for w in node.writes: @@ -90,6 +92,7 @@ class DisaggregatedRunner: def tick(self) -> bool: if self.done: return True + self.cancel_scope.check() # common-path cancellation (online + offline, all nodes) nodes = self.program.active_nodes(self.request) while True: if self.node_idx >= len(nodes): @@ -156,6 +159,11 @@ class DisaggregatedRunner: else: stuck = 0 + def close(self) -> None: + """Release any pool this request is occupying (idempotent). Called on terminal exit so a + cancel/error mid-loop never leaks pool capacity (the §6.4 failure-isolation guarantee).""" + self._leave_current() + def output(self) -> Output: artifacts = {name: _to_artifact(name, self.slots.get(slot), producer=slot) for name, slot in self.program.output_artifacts.items()} diff --git a/mini_fastvideo/runtime/engine.py b/mini_fastvideo/runtime/engine.py index 098e43bc6..a96d5de4c 100644 --- a/mini_fastvideo/runtime/engine.py +++ b/mini_fastvideo/runtime/engine.py @@ -84,6 +84,7 @@ class ProgramRunner: def tick(self) -> bool: if self.done: return True + self.cancel_scope.check() # common-path cancellation (offline + all-component programs) nodes = self.program.nodes while True: if self.node_idx >= len(nodes): diff --git a/mini_fastvideo/serving/http.py b/mini_fastvideo/serving/http.py index 4bec65af4..d6e688bd6 100644 --- a/mini_fastvideo/serving/http.py +++ b/mini_fastvideo/serving/http.py @@ -42,23 +42,41 @@ class Response: _STATUS = {200: "OK", 201: "Created", 202: "Accepted", 400: "Bad Request", - 404: "Not Found", 500: "Internal Server Error"} + 404: "Not Found", 408: "Request Timeout", 413: "Payload Too Large", + 500: "Internal Server Error"} Handler = Callable[[Request], Awaitable[Response]] -async def _read_request(reader: asyncio.StreamReader) -> Request: +class _BadRequest(Exception): + pass + + +class _TooLarge(Exception): + pass + + +async def _read_request(reader: asyncio.StreamReader, max_body: int) -> Request: head = await reader.readuntil(b"\r\n\r\n") lines = head.decode("latin1").split("\r\n") - method, target, _ = lines[0].split(" ", 2) + parts = lines[0].split(" ", 2) + if len(parts) < 2: + raise _BadRequest("malformed request line") + method, target = parts[0], parts[1] path, _, qs = target.partition("?") query = dict(kv.split("=", 1) for kv in qs.split("&") if "=" in kv) - headers = {} + headers: dict[str, str] = {} for ln in lines[1:]: if ":" in ln: k, v = ln.split(":", 1) headers[k.strip().lower()] = v.strip() - clen = int(headers.get("content-length", "0") or "0") + raw = headers.get("content-length", "0") or "0" + try: + clen = int(raw) + except ValueError: + raise _BadRequest(f"invalid Content-Length {raw!r}") + if clen < 0 or clen > max_body: + raise _TooLarge(f"Content-Length {clen} exceeds limit {max_body}") body = await reader.readexactly(clen) if clen else b"" return Request(method, path, headers, body, query) @@ -84,26 +102,43 @@ async def _write(writer: asyncio.StreamWriter, resp: Response) -> None: class HttpServer: - def __init__(self, dispatch: Handler): + def __init__(self, dispatch: Handler, *, read_timeout: float = 30.0, + max_body_bytes: int = 64 << 20, header_limit: int = 1 << 20): self.dispatch = dispatch + self.read_timeout = read_timeout # slowloris guard + self.max_body_bytes = max_body_bytes # body-size cap → 413 + self.header_limit = header_limit # StreamReader buffer for the header block self._server: asyncio.AbstractServer | None = None async def _conn(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: try: - req = await _read_request(reader) + req = await asyncio.wait_for(_read_request(reader, self.max_body_bytes), self.read_timeout) resp = await self.dispatch(req) + except asyncio.TimeoutError: + resp = Response.json({"error": "request read timeout"}, status=408) + except _BadRequest as e: + resp = Response.json({"error": str(e)}, status=400) + except _TooLarge as e: + resp = Response.json({"error": str(e)}, status=413) + except (asyncio.IncompleteReadError, asyncio.LimitOverrunError): + resp = Response.json({"error": "malformed or oversized request"}, status=400) except Exception as e: # noqa: BLE001 — never let one connection kill the server resp = Response.json({"error": str(e)}, status=500) try: await _write(writer, resp) except Exception: + if resp.stream is not None: # client gone mid-SSE → close the generator + try: + await resp.stream.aclose() # → GeneratorExit → submit() cancels the driver + except Exception: + pass try: writer.close() except Exception: pass async def start(self, host: str = "127.0.0.1", port: int = 0) -> tuple[str, int]: - self._server = await asyncio.start_server(self._conn, host, port) + self._server = await asyncio.start_server(self._conn, host, port, limit=self.header_limit) sock = self._server.sockets[0].getsockname() return sock[0], sock[1] diff --git a/mini_fastvideo/serving/server.py b/mini_fastvideo/serving/server.py index 59669a4e5..0b761e4dc 100644 --- a/mini_fastvideo/serving/server.py +++ b/mini_fastvideo/serving/server.py @@ -61,12 +61,22 @@ def _ser_artifacts(out: Output) -> dict: class OmniOpenAIServer: - def __init__(self, engine: Any, *, engine_id: str = "mini-fastvideo"): + 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: @@ -122,7 +132,10 @@ class OmniOpenAIServer: 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} - asyncio.create_task(self._run_video_job(job_id, vg)) + 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: @@ -187,4 +200,11 @@ class OmniOpenAIServer: 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() diff --git a/mini_fastvideo/tests/test_serving_fixes.py b/mini_fastvideo/tests/test_serving_fixes.py new file mode 100644 index 000000000..efcac34c1 --- /dev/null +++ b/mini_fastvideo/tests/test_serving_fixes.py @@ -0,0 +1,127 @@ +"""Regression tests for the serving adversarial-review findings (previously-untested failure paths).""" +from __future__ import annotations + +import asyncio + +from mini_fastvideo.card import CostModel +from mini_fastvideo.deploy import build_deployment_card +from mini_fastvideo.models import build_default_engine, build_wan21_card, build_wan_t2v_program +from mini_fastvideo.request import Cancelled, DiffusionParams, OutputSpec, TaskType, make_request +from mini_fastvideo.runtime import AsyncEngine, PoolSet, wan_t2v_disaggregated +from mini_fastvideo.transport import InProcConnector + + +def _disagg_engine(): + ae = AsyncEngine(build_default_engine()) + card = build_wan21_card() + pools = PoolSet(wan_t2v_disaggregated(), card) + pools.warmup() + ae.register_disaggregated("wd", pools, build_wan_t2v_program()) + return ae, pools + + +def test_pool_capacity_released_on_cancel_mid_disaggregation(): + """HIGH: cancelling mid-denoise must not brick the capacity-1 denoiser pool.""" + async def run(): + ae, pools = _disagg_engine() + req = make_request(TaskType.T2V, "wd", "x", + diffusion=DiffusionParams(num_steps=80, seed=1), + outputs=OutputSpec(stream={"video": True})) # emit per-step → observe occupancy + cancelled = False + async for ev in ae.submit(req): + if not cancelled and pools.by_id["den"].in_flight >= 1: + ae.cancel(req.request_id) + cancelled = True + assert cancelled and pools.by_id["den"].in_flight == 0 # capacity returned, not leaked + # the pool is NOT bricked — a fresh request still runs + out = await ae.generate(make_request(TaskType.T2V, "wd", "y", diffusion=DiffusionParams(num_steps=3, seed=2))) + assert out.artifacts["video"].frames is not None + asyncio.run(run()) + + +def test_credit_and_capacity_released_on_transfer_failure(): + """A failing cross-pool transfer releases its credit AND the occupied pool (no leak).""" + async def run(): + ae, pools = _disagg_engine() + + class Boom(InProcConnector): + def put(self, key, value, manifest=None): + raise RuntimeError("transfer boom") + + pools.by_id["enc"].connector = Boom("enc") + out = await ae.generate(make_request(TaskType.T2V, "wd", "x", diffusion=DiffusionParams(num_steps=3, seed=1))) + assert out.error is not None # request-fatal isolation + enc = pools.by_id["enc"].connector + assert enc.credits.available == enc.credits.capacity # credit refunded on the error path + assert pools.by_id["den"].in_flight == 0 # pool capacity refunded too + asyncio.run(run()) + + +def test_duplicate_request_id_rejected_not_deadlocked(): + async def run(): + ae = AsyncEngine(build_default_engine()) + r1 = make_request(TaskType.T2V, "wan2.1-1.3b", "x", + diffusion=DiffusionParams(num_steps=30, seed=1), request_id="dup") + r2 = make_request(TaskType.T2V, "wan2.1-1.3b", "y", + diffusion=DiffusionParams(num_steps=30, seed=2), request_id="dup") + g1 = ae.submit(r1) + await g1.__anext__() # r1 now in flight as "dup" + raised = False + try: + await ae.submit(r2).__anext__() + except ValueError: + raised = True + assert raised, "duplicate in-flight request_id must be rejected, not silently deadlock" + async for _ in g1: + pass + asyncio.run(run()) + + +def test_deployment_cards_do_not_alias_cost_model(): + card = build_wan21_card() + a = build_deployment_card("A", [card]) + b = build_deployment_card("B", [card]) + assert a.cost_model is not b.cost_model # independent calibration state + a.cost_model.per_unit_seconds = 999.0 + assert b.cost_model.per_unit_seconds != 999.0 + + +def test_offline_runner_honors_cancellation(): + """Cancellation is common-path even on the offline (non-async) path.""" + eng = build_default_engine() + runner = eng._make_runner(make_request(TaskType.T2V, "wan2.1-1.3b", "x", + diffusion=DiffusionParams(num_steps=4, seed=1))) + runner.cancel_scope.cancel() + raised = False + try: + runner.run_to_completion() + except Cancelled: + raised = True + assert raised, "offline runner must honor cancel, not run to completion" + + +def test_http_rejects_oversized_and_malformed_requests(): + from mini_fastvideo.runtime import AsyncEngine as AE + from mini_fastvideo.serving import OmniOpenAIServer + + async def status_of(host, port, raw: bytes) -> int: + r, w = await asyncio.open_connection(host, port) + w.write(raw) + await w.drain() + data = await r.read() + w.close() + return int(data.split(b" ")[1]) + + async def run(): + s = OmniOpenAIServer(AE(build_default_engine())) + host, port = await s.serve(port=0) + # 413: Content-Length far beyond the body cap (rejected before reading the body) + st = await status_of(host, port, + b"POST /v1/images/generations HTTP/1.1\r\nHost: x\r\nContent-Length: 999999999\r\n\r\n") + assert st == 413 + # 400: non-numeric Content-Length + st = await status_of(host, port, + b"POST /v1/images/generations HTTP/1.1\r\nHost: x\r\nContent-Length: abc\r\n\r\n") + assert st == 400 + await s.close() + asyncio.run(run())