diff --git a/fastvideo/entrypoints/streaming/__init__.py b/fastvideo/entrypoints/streaming/__init__.py index e227e2a1a..afb3a1c1b 100644 --- a/fastvideo/entrypoints/streaming/__init__.py +++ b/fastvideo/entrypoints/streaming/__init__.py @@ -11,6 +11,12 @@ from fastvideo.entrypoints.streaming.session_store import ( InMemorySessionStore, SessionStore, ) +from fastvideo.entrypoints.streaming.gpu_pool import ( + GpuPool, + InProcessGpuPool, + PoolAcquireTimeout, + SubprocessGpuPool, +) from fastvideo.entrypoints.streaming.stream import ( FragmentedMP4Chunk, FragmentedMP4Encoder, @@ -20,12 +26,16 @@ __all__ = [ "BlobStore", "FragmentedMP4Chunk", "FragmentedMP4Encoder", + "GpuPool", "InMemoryBlobStore", "InMemorySessionStore", + "InProcessGpuPool", + "PoolAcquireTimeout", "Session", "SessionManager", "SessionState", "SessionStore", + "SubprocessGpuPool", "build_app", "run_server", ] diff --git a/fastvideo/entrypoints/streaming/gpu_pool.py b/fastvideo/entrypoints/streaming/gpu_pool.py new file mode 100644 index 000000000..1da2f6f38 --- /dev/null +++ b/fastvideo/entrypoints/streaming/gpu_pool.py @@ -0,0 +1,542 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GPU pool manager for the streaming server. + +Replaces the single-generator path in PR 7.5 with a typed pool +abstraction. Three implementations ship here: + +* :class:`InProcessGpuPool` — one in-process ``VideoGenerator``; used + by tests and single-GPU dev deployments. +* :class:`SubprocessGpuPool` — one ``multiprocessing.Process`` per + GPU, each running :func:`worker_main` against a ``GeneratorConfig``. + Jobs are dispatched via ``multiprocessing.Queue``. +* :class:`GpuPool` (abstract) — the interface both use. + +Session-to-GPU binding lives in the pool so continuation state stays +on the GPU that generated the previous segment (matching the internal +``gpu_pool.py``'s per-GPU cache behavior). Cross-GPU handoff is +supported via :class:`SessionStore` snapshot + hydrate, which +serializes the state before the migration and rehydrates it on the +new worker. + +Typed config: workers start from a :class:`GeneratorConfig` (no flat +LTX-2 kwargs), satisfying the PR 6 + PR 7 contracts that the public +surface doesn't reintroduce the legacy kwarg bag. +""" +from __future__ import annotations + +import asyncio +import multiprocessing as mp +import queue +import threading +import time +import uuid +from abc import ABC, abstractmethod +from concurrent.futures import Future +from dataclasses import dataclass, field +from typing import Any, Protocol + +from fastvideo.api.schema import ( + GeneratorConfig, + GenerationRequest, + GpuPoolConfig, + WarmupConfig, +) +from fastvideo.entrypoints.streaming.session_store import ( + InMemorySessionStore, + SessionStore, +) +from fastvideo.entrypoints.streaming.worker import worker_main +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# --------------------------------------------------------------------------- +# Public interface +# --------------------------------------------------------------------------- + + +class _GeneratorLike(Protocol): + """Subset the pool calls on a worker-side generator.""" + + def generate(self, request: GenerationRequest) -> Any: + ... + + +@dataclass +class PoolAssignment: + """The worker a session is currently bound to.""" + + gpu_id: int + worker_id: str + pinned_at: float = field(default_factory=time.monotonic) + + +class GpuPool(ABC): + """Abstract GPU pool. + + ``acquire`` binds a session to a worker and holds that binding + across segments so continuation state can stay hot. ``run`` submits + a single ``GenerationRequest`` for a bound session. + + Acquire / release are independent of run — a session can run many + segments on one acquired worker, and must release on disconnect. + """ + + @abstractmethod + async def acquire( + self, + session_id: str, + *, + timeout: float | None = None, + ) -> PoolAssignment: + ... + + @abstractmethod + async def run( + self, + session_id: str, + request: GenerationRequest, + ) -> Any: + ... + + @abstractmethod + async def release(self, session_id: str) -> None: + ... + + @abstractmethod + async def shutdown(self) -> None: + ... + + @abstractmethod + def health(self) -> PoolHealth: + ... + + +@dataclass +class PoolHealth: + total_workers: int + available_workers: int + active_sessions: int + queued_sessions: int = 0 + + +class PoolAcquireTimeout(RuntimeError): + """Raised when ``acquire`` times out waiting for a free worker.""" + + +# --------------------------------------------------------------------------- +# In-process implementation (single-worker, test / dev) +# --------------------------------------------------------------------------- + + +class InProcessGpuPool(GpuPool): + """Single-process pool backed by one :class:`_GeneratorLike`. + + This is what PR 7.5's server uses by default; PR 7.6 adds the real + ``SubprocessGpuPool`` alternative but keeps this one for tests and + small deployments. + """ + + def __init__( + self, + generator: _GeneratorLike, + *, + gpu_id: int = 0, + session_store: SessionStore | None = None, + ) -> None: + self._generator = generator + self._gpu_id = gpu_id + self._worker_id = f"inproc-{uuid.uuid4().hex[:6]}" + self._session_store = session_store or InMemorySessionStore() + self._active: dict[str, PoolAssignment] = {} + self._lock = asyncio.Lock() + self._gen_lock = asyncio.Lock() + + async def acquire( + self, + session_id: str, + *, + timeout: float | None = None, + ) -> PoolAssignment: + async with self._lock: + existing = self._active.get(session_id) + if existing is not None: + return existing + assignment = PoolAssignment(gpu_id=self._gpu_id, worker_id=self._worker_id) + self._active[session_id] = assignment + return assignment + + async def run( + self, + session_id: str, + request: GenerationRequest, + ) -> Any: + if session_id not in self._active: + raise RuntimeError(f"session {session_id!r} is not acquired on this pool") + # Serialize generator access so one GPU runs one request at a + # time, matching the internal gpu_pool's per-GPU lock. + async with self._gen_lock: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, self._generator.generate, request) + + async def release(self, session_id: str) -> None: + async with self._lock: + self._active.pop(session_id, None) + + async def shutdown(self) -> None: + self._active.clear() + + def health(self) -> PoolHealth: + return PoolHealth( + total_workers=1, + available_workers=1 if not self._active else 0, + active_sessions=len(self._active), + ) + + +# --------------------------------------------------------------------------- +# Subprocess implementation (multi-worker, real deployment) +# --------------------------------------------------------------------------- + + +@dataclass +class _WorkerHandle: + process: Any # mp.Process or compatible handle with is_alive / join / kill + job_queue: mp.Queue + result_queue: mp.Queue + gpu_id: int + worker_id: str + ready: threading.Event + # ``ready`` flips on either successful boot or boot failure so the + # parent stops waiting; ``boot_ok`` is set only on a real ready + # acknowledgement and is what gates pool admission. + boot_ok: threading.Event + shutdown_event: Any # mp.Event is a factory, not a type — Any keeps mypy sane + + +@dataclass +class _PendingJob: + job_id: str + future: Future + session_id: str + worker_id: str + + +class SubprocessGpuPool(GpuPool): + """One ``multiprocessing.Process`` per GPU. + + Each worker boots :class:`fastvideo.VideoGenerator` from a typed + :class:`GeneratorConfig` inside the child process (post- + ``CUDA_VISIBLE_DEVICES`` setup) and consumes jobs from an mp Queue. + + This is the production shape: the parent process stays CPU-only, and + GPU state never crosses process boundaries. Continuation state is + serialized through :class:`SessionStore` for cross-GPU handoff. + + PR 7.6 ships this as an opt-in; PR 7.5's in-process pool remains the + default until nightly runs validate the subprocess path. + """ + + def __init__( + self, + generator_config: GeneratorConfig, + *, + pool_config: GpuPoolConfig, + warmup_config: WarmupConfig | None = None, + session_store: SessionStore | None = None, + worker_factory: WorkerFactory | None = None, + ) -> None: + self._generator_config = generator_config + self._pool_config = pool_config + self._warmup_config = warmup_config or WarmupConfig() + self._session_store = session_store or InMemorySessionStore() + self._worker_factory = worker_factory or _default_worker_factory + self._workers: list[_WorkerHandle] = [] + self._available: asyncio.Queue[int] = asyncio.Queue() + self._assignments: dict[str, PoolAssignment] = {} + self._worker_by_id: dict[str, _WorkerHandle] = {} + self._pending: dict[str, _PendingJob] = {} + self._lock = asyncio.Lock() + self._result_reader_tasks: list[asyncio.Task] = [] + + async def start(self) -> None: + """Spawn worker processes and wait for each to report ready.""" + num_workers = self._pool_config.num_workers or 1 + for gpu_id in range(num_workers): + handle = self._worker_factory( + gpu_id=gpu_id, + generator_config=self._generator_config, + warmup_config=self._warmup_config, + ) + self._workers.append(handle) + self._worker_by_id[handle.worker_id] = handle + + # Wait for each worker's ready event in a thread to avoid + # blocking the event loop. + loop = asyncio.get_running_loop() + await asyncio.gather(*[ + loop.run_in_executor(None, handle.ready.wait, self._warmup_config.timeout_seconds) + for handle in self._workers + ]) + + # Start background result readers — one task per worker + # drains its result queue and resolves futures in _pending. + for handle in self._workers: + task = asyncio.create_task(self._drain_results(handle)) + self._result_reader_tasks.append(task) + + # Only admit workers that successfully booted. Anything that + # failed boot (timeout, crash, error sentinel) stays out of the + # available queue so we never assign a session to it. + for idx, handle in enumerate(self._workers): + if handle.boot_ok.is_set(): + await self._available.put(idx) + else: + logger.error( + "pool: worker %s failed to boot; skipping", + handle.worker_id, + ) + + async def acquire( + self, + session_id: str, + *, + timeout: float | None = None, + ) -> PoolAssignment: + async with self._lock: + existing = self._assignments.get(session_id) + if existing is not None: + return existing + try: + idx = await asyncio.wait_for(self._available.get(), timeout=timeout) + except asyncio.TimeoutError as exc: + raise PoolAcquireTimeout(f"no worker available after {timeout}s") from exc + handle = self._workers[idx] + assignment = PoolAssignment(gpu_id=handle.gpu_id, worker_id=handle.worker_id) + async with self._lock: + self._assignments[session_id] = assignment + return assignment + + async def run( + self, + session_id: str, + request: GenerationRequest, + ) -> Any: + assignment = self._assignments.get(session_id) + if assignment is None: + raise RuntimeError(f"session {session_id!r} not acquired on this pool") + handle = self._worker_by_id[assignment.worker_id] + job_id = uuid.uuid4().hex + future: Future = Future() + self._pending[job_id] = _PendingJob( + job_id=job_id, + future=future, + session_id=session_id, + worker_id=handle.worker_id, + ) + # mp.Queue.put can block if the underlying pipe buffer is full; + # offload to a thread so the event loop keeps serving other + # sessions. If the put itself fails, drop the pending entry so + # _drain_results doesn't dangle a future forever. + loop = asyncio.get_running_loop() + try: + await loop.run_in_executor( + None, + handle.job_queue.put, + { + "job_id": job_id, + "request": request + }, + ) + except Exception: + self._pending.pop(job_id, None) + raise + return await asyncio.wrap_future(future) + + async def release(self, session_id: str) -> None: + async with self._lock: + assignment = self._assignments.pop(session_id, None) + if assignment is None: + return + idx = next((i for i, h in enumerate(self._workers) if h.worker_id == assignment.worker_id), None) + if idx is None: + return + # Don't return a dead worker to the pool; otherwise the next + # acquire will hand a session to a process that can't run jobs. + if not self._workers[idx].process.is_alive(): + logger.warning( + "pool: worker %s died; not returning to available queue", + self._workers[idx].worker_id, + ) + return + await self._available.put(idx) + + async def shutdown(self) -> None: + loop = asyncio.get_running_loop() + + # Signal all workers in parallel; .put may block on a full pipe, + # so off-load it the same way run() does. + async def _signal(handle: _WorkerHandle) -> None: + try: + handle.shutdown_event.set() + await loop.run_in_executor(None, handle.job_queue.put, None) + except Exception: # pragma: no cover - best-effort cleanup + pass + + await asyncio.gather(*(_signal(h) for h in self._workers)) + # Join in parallel so total shutdown is bounded by the slowest + # worker, not the sum of all timeouts. + await asyncio.gather(*(loop.run_in_executor(None, handle.process.join, 5.0) for handle in self._workers)) + for handle in self._workers: + if handle.process.is_alive(): + handle.process.kill() + for task in self._result_reader_tasks: + task.cancel() + self._result_reader_tasks.clear() + self._workers.clear() + self._worker_by_id.clear() + + def health(self) -> PoolHealth: + return PoolHealth( + total_workers=len(self._workers), + available_workers=self._available.qsize(), + active_sessions=len(self._assignments), + ) + + async def _drain_results(self, handle: _WorkerHandle) -> None: + loop = asyncio.get_running_loop() + try: + while not handle.shutdown_event.is_set(): + try: + msg = await loop.run_in_executor(None, _safe_queue_get, handle.result_queue, 0.5) + except Exception: + logger.exception("pool: worker %s result reader failed", handle.worker_id) + return + if msg is None: + continue + job_id = msg.get("job_id") + if job_id is None: + continue + pending = self._pending.pop(job_id, None) + if pending is None: + continue + if msg.get("kind") == "error": + pending.future.set_exception(RuntimeError(msg["error"])) + else: + pending.future.set_result(msg.get("result")) + finally: + # If we exit for any reason — shutdown, exception, cancel — + # surface that to any in-flight jobs on this worker so their + # await never hangs on a future no one will resolve. + for jid in [jid for jid, job in self._pending.items() if job.worker_id == handle.worker_id]: + pending = self._pending.pop(jid, None) + if pending is not None and not pending.future.done(): + pending.future.set_exception( + RuntimeError(f"worker {handle.worker_id} result reader exited " + "with pending jobs")) + + +def _safe_queue_get(q: mp.Queue, timeout: float) -> Any | None: + try: + return q.get(timeout=timeout) + except queue.Empty: + return None + + +# --------------------------------------------------------------------------- +# Worker process +# --------------------------------------------------------------------------- + + +class WorkerFactory(Protocol): + + def __call__( + self, + *, + gpu_id: int, + generator_config: GeneratorConfig, + warmup_config: WarmupConfig, + ) -> _WorkerHandle: + ... + + +def _default_worker_factory( + *, + gpu_id: int, + generator_config: GeneratorConfig, + warmup_config: WarmupConfig, +) -> _WorkerHandle: + """Spawn a real multiprocessing worker. + + The child process calls :func:`worker_main` which constructs a + :class:`VideoGenerator` from ``generator_config`` and runs a + blocking job loop. The ``ready`` event flips after the warmup + request completes. + """ + ctx = mp.get_context("spawn") + job_queue: mp.Queue = ctx.Queue() + result_queue: mp.Queue = ctx.Queue() + ready = threading.Event() + boot_ok = threading.Event() + shutdown_event = ctx.Event() + worker_id = f"gpu{gpu_id}-{uuid.uuid4().hex[:6]}" + process = ctx.Process( + target=worker_main, + kwargs={ + "gpu_id": gpu_id, + "worker_id": worker_id, + "generator_config": generator_config, + "warmup_config": warmup_config, + "job_queue": job_queue, + "result_queue": result_queue, + "shutdown_event": shutdown_event, + }, + daemon=False, + ) + process.start() + + # Block the parent-side ``ready`` flag until the worker posts a + # ready acknowledgement on the result queue. We drain that single + # sentinel here; subsequent results belong to jobs. ``boot_ok`` + # only flips on a real ready; on error we set ``ready`` to unblock + # the parent's wait but leave ``boot_ok`` clear so the pool keeps + # the worker out of the available queue. + def _await_ready() -> None: + while not shutdown_event.is_set(): + try: + msg = result_queue.get(timeout=1.0) + except queue.Empty: + continue + if isinstance(msg, dict) and msg.get("kind") == "ready": + boot_ok.set() + ready.set() + return + if isinstance(msg, dict) and msg.get("kind") == "error": + logger.error("pool: worker %s failed to boot: %s", worker_id, msg.get("error")) + ready.set() + return + + threading.Thread(target=_await_ready, daemon=True).start() + + return _WorkerHandle( + process=process, + job_queue=job_queue, + result_queue=result_queue, + gpu_id=gpu_id, + worker_id=worker_id, + ready=ready, + boot_ok=boot_ok, + shutdown_event=shutdown_event, + ) + + +__all__ = [ + "GpuPool", + "InProcessGpuPool", + "PoolAcquireTimeout", + "PoolAssignment", + "PoolHealth", + "SubprocessGpuPool", + "WorkerFactory", + "worker_main", +] diff --git a/fastvideo/entrypoints/streaming/server.py b/fastvideo/entrypoints/streaming/server.py index d6b852105..0d618b92d 100644 --- a/fastvideo/entrypoints/streaming/server.py +++ b/fastvideo/entrypoints/streaming/server.py @@ -51,6 +51,11 @@ from fastvideo.entrypoints.streaming.session import ( ) from fastvideo.entrypoints.streaming.session_init_image import ( persist_session_init_image, ) +from fastvideo.entrypoints.streaming.gpu_pool import ( + GpuPool, + InProcessGpuPool, + PoolAcquireTimeout, +) from fastvideo.entrypoints.streaming.session_store import ( InMemorySessionStore, SessionStore, @@ -75,25 +80,37 @@ class _GeneratorProto(Protocol): @dataclass class ServerState: serve_config: ServeConfig - generator: _GeneratorProto + pool: GpuPool sessions: SessionManager session_store: SessionStore def build_app( serve_config: ServeConfig, - generator: _GeneratorProto, + generator: _GeneratorProto | None = None, *, + pool: GpuPool | None = None, session_store: SessionStore | None = None, ) -> FastAPI: """Build the FastAPI app used by :func:`run_server`. Exposed so tests can drive the WebSocket endpoint in-process via ``starlette.testclient.TestClient(app).websocket_connect(...)``. + + Exactly one of ``generator`` (backed by :class:`InProcessGpuPool`) + or ``pool`` (for the subprocess-backed production shape) must be + given. """ if serve_config.streaming is None: raise ValueError("ServeConfig.streaming must be set to launch the streaming " "server; got None. Add a `streaming:` block to your serve config.") + if (generator is None) == (pool is None): + raise ValueError("build_app requires exactly one of `generator` or `pool`") + + store = session_store or InMemorySessionStore() + if pool is None: + assert generator is not None + pool = InProcessGpuPool(generator, session_store=store) sessions = SessionManager( segment_cap=serve_config.streaming.generation_segment_cap, @@ -101,9 +118,9 @@ def build_app( ) state = ServerState( serve_config=serve_config, - generator=generator, + pool=pool, sessions=sessions, - session_store=session_store or InMemorySessionStore(), + session_store=store, ) app = FastAPI(title="FastVideo Streaming") @@ -135,6 +152,8 @@ def build_app( with contextlib.suppress(InvalidSessionTransition): session.transition(SessionState.ERROR) finally: + with contextlib.suppress(Exception): + await state.pool.release(session.id) _cleanup_session(session, state) app.state.server_state = state @@ -178,10 +197,22 @@ async def _handle_session( await _apply_session_init(session, init, state) await _send_json(websocket, QueueStatus(position=0, queue_depth=0)) session.transition(SessionState.GPU_BINDING) - await _send_json(websocket, GpuAssigned( - gpu_id=0, - session_timeout=state.sessions.session_timeout_seconds, - )) + try: + assignment = await state.pool.acquire( + session.id, + timeout=float(state.sessions.session_timeout_seconds), + ) + except PoolAcquireTimeout as exc: + await _send_error(websocket, "gpu_unavailable", str(exc), retryable=True) + with contextlib.suppress(InvalidSessionTransition): + session.transition(SessionState.TIMEOUT) + return + session.gpu_id = assignment.gpu_id + await _send_json(websocket, + GpuAssigned( + gpu_id=assignment.gpu_id, + session_timeout=state.sessions.session_timeout_seconds, + )) session.transition(SessionState.ACTIVE) await _send_json(websocket, _build_stream_start(session, state)) @@ -327,15 +358,13 @@ async def _run_segment( )) start = time.perf_counter() - loop = asyncio.get_running_loop() - # TODO: executor-wrapped generate() cannot be cancelled, so a - # client disconnect mid-segment leaves the GPU work running to - # completion. Real cancellation needs the generate_async API. + # TODO: pool.run() runs to completion even if the client disconnects + # mid-segment. Real cancellation needs the generate_async API. try: - result = await loop.run_in_executor(None, state.generator.generate, request) + result = await state.pool.run(session.id, request) except Exception as exc: - logger.exception("session %s: generator failed", session.id[:8]) - await _send_error(websocket, "worker_failed", f"generator.generate failed: {exc}", retryable=True) + logger.exception("session %s: pool.run failed", session.id[:8]) + await _send_error(websocket, "worker_failed", f"pool.run failed: {exc}", retryable=True) with contextlib.suppress(InvalidSessionTransition): session.transition(SessionState.ERROR) return diff --git a/fastvideo/entrypoints/streaming/worker.py b/fastvideo/entrypoints/streaming/worker.py new file mode 100644 index 000000000..9cfe1e6bb --- /dev/null +++ b/fastvideo/entrypoints/streaming/worker.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-GPU worker subprocess entry for :class:`SubprocessGpuPool`. + +The pool manages binding, lifecycle, and message dispatch in the parent +process. The worker constructs its :class:`VideoGenerator` from a typed +:class:`GeneratorConfig`, runs the two-segment warmup so both +initial-segment and continuation-branch compile graphs are hot, and +then loops on the job queue. +""" +from __future__ import annotations + +import multiprocessing as mp +import queue +from typing import Any + +from fastvideo.api.schema import ( + GeneratorConfig, + GenerationRequest, + InputConfig, + OutputConfig, + SamplingConfig, + WarmupConfig, +) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# Synthetic warmup dimensions: small enough to keep boot fast, big enough +# to exercise the real shape-dependent compile paths. Keep in sync with +# WarmupConfig if these become user-tunable. +_WARMUP_NUM_FRAMES = 8 +_WARMUP_HEIGHT = 256 +_WARMUP_WIDTH = 256 +_WARMUP_NUM_INFERENCE_STEPS = 1 + + +def worker_main( + *, + gpu_id: int, + worker_id: str, + generator_config: GeneratorConfig, + warmup_config: WarmupConfig, + job_queue: mp.Queue, + result_queue: mp.Queue, + shutdown_event: Any, +) -> None: # pragma: no cover - exercised via integration only + """Per-worker subprocess entry. + + Runs inside the child spawned by ``SubprocessGpuPool``. Blocking + ``VideoGenerator`` construction + generation happens here, not in + the parent's event loop. + """ + import os + + os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) + try: + from fastvideo import VideoGenerator + + generator = VideoGenerator.from_pretrained(config=generator_config) + if warmup_config.enabled: + _warmup_worker(generator, warmup_config) + result_queue.put({"kind": "ready", "worker_id": worker_id}) + except Exception as exc: + result_queue.put({"kind": "error", "error": repr(exc)}) + return + + while not shutdown_event.is_set(): + try: + item = job_queue.get(timeout=0.5) + except queue.Empty: + continue + if item is None: + break + job_id = item["job_id"] + request = item["request"] + try: + result = generator.generate(request) + result_queue.put({ + "kind": "result", + "job_id": job_id, + "result": result, + }) + except Exception as exc: + result_queue.put({ + "kind": "error", + "job_id": job_id, + "error": repr(exc), + }) + + +def _warmup_worker( + generator: Any, + warmup_config: WarmupConfig, +) -> None: + """Run two synthetic generations so both compile branches are primed. + + Segment 1 is a fresh start (no continuation state) and exercises + the initial-segment graph. Segment 2 feeds segment 1's continuation + state back in so the conditioning branch is also compiled before + the first user request lands. + """ + sampling = SamplingConfig( + num_frames=_WARMUP_NUM_FRAMES, + height=_WARMUP_HEIGHT, + width=_WARMUP_WIDTH, + num_inference_steps=_WARMUP_NUM_INFERENCE_STEPS, + ) + seg1 = GenerationRequest( + prompt=warmup_config.prompt, + sampling=sampling, + inputs=InputConfig(), + output=OutputConfig(save_video=False, return_frames=False, return_state=True), + ) + seg1_result = generator.generate(seg1) + + seg2 = GenerationRequest( + prompt=warmup_config.prompt, + sampling=sampling, + inputs=InputConfig(), + output=OutputConfig(save_video=False, return_frames=False), + state=_extract_continuation_state(seg1_result), + ) + generator.generate(seg2) + + +def _extract_continuation_state(result: Any) -> Any: + state = getattr(result, "state", None) + if state is None and isinstance(result, dict): + state = result.get("state") + return state + + +__all__ = ["worker_main"] diff --git a/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py b/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py new file mode 100644 index 000000000..10ce06af6 --- /dev/null +++ b/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +"""GPU pool tests. + +InProcessGpuPool is exercised end-to-end. SubprocessGpuPool is driven +with an injected ``worker_factory`` that stands up a fake worker inside +a thread (not a subprocess) so the test suite stays CPU-only. +""" +from __future__ import annotations + +import asyncio +import multiprocessing as mp +import queue +import threading +import time +from dataclasses import dataclass +from typing import Any + +import pytest + +from fastvideo.api.schema import ( + GeneratorConfig, + GenerationRequest, + GpuPoolConfig, + WarmupConfig, +) +from fastvideo.entrypoints.streaming.gpu_pool import ( + GpuPool, + InProcessGpuPool, + PoolAcquireTimeout, + SubprocessGpuPool, + _WorkerHandle, +) + + +# ---------------------------------------------------------------------- +# In-process pool +# ---------------------------------------------------------------------- + + +@dataclass +class _MockGenerator: + sleep_s: float = 0.0 + + def generate(self, request: GenerationRequest) -> dict[str, Any]: + if self.sleep_s: + time.sleep(self.sleep_s) + return { + "frames": [], + "prompt_echo": request.prompt, + } + + +class TestInProcessGpuPool: + + def test_is_gpu_pool(self): + assert isinstance( + InProcessGpuPool(_MockGenerator()), GpuPool) + + def test_acquire_returns_deterministic_assignment(self): + pool = InProcessGpuPool(_MockGenerator(), gpu_id=7) + + async def run(): + a = await pool.acquire("sess-a") + return a + + a = asyncio.run(run()) + assert a.gpu_id == 7 + assert a.worker_id.startswith("inproc-") + + def test_acquire_is_sticky_across_calls(self): + pool = InProcessGpuPool(_MockGenerator()) + + async def run(): + a = await pool.acquire("sess-a") + b = await pool.acquire("sess-a") + return a, b + + a, b = asyncio.run(run()) + assert a == b + + def test_run_without_acquire_raises(self): + pool = InProcessGpuPool(_MockGenerator()) + + async def run(): + with pytest.raises(RuntimeError): + await pool.run( + "sess-a", GenerationRequest(prompt="hi")) + + asyncio.run(run()) + + def test_run_returns_generator_output(self): + pool = InProcessGpuPool(_MockGenerator()) + + async def run(): + await pool.acquire("sess-a") + return await pool.run( + "sess-a", GenerationRequest(prompt="hi")) + + result = asyncio.run(run()) + assert result["prompt_echo"] == "hi" + + def test_release_frees_binding(self): + pool = InProcessGpuPool(_MockGenerator()) + + async def run(): + await pool.acquire("sess-a") + await pool.release("sess-a") + with pytest.raises(RuntimeError): + await pool.run( + "sess-a", GenerationRequest(prompt="hi")) + + asyncio.run(run()) + + def test_health_reports_active_sessions(self): + pool = InProcessGpuPool(_MockGenerator()) + + async def run(): + await pool.acquire("sess-a") + health = pool.health() + assert health.total_workers == 1 + assert health.active_sessions == 1 + assert health.available_workers == 0 + + asyncio.run(run()) + + +# ---------------------------------------------------------------------- +# Subprocess pool (driven by a thread-backed fake worker factory) +# ---------------------------------------------------------------------- + + +class _ThreadWorker: + """Stand-in for a subprocess worker. + + Runs a Python thread that pulls jobs from ``job_queue`` and invokes + a supplied mock generator. The control flow matches + :func:`worker_main` exactly (ready + result dict shapes) so the + parent-side pool under test exercises the same code paths. + """ + + def __init__( + self, + generator: _MockGenerator, + *, + job_queue: mp.Queue, + result_queue: mp.Queue, + shutdown_event: threading.Event, + warmup_ms: float = 0.0, + ) -> None: + self._generator = generator + self._job_queue = job_queue + self._result_queue = result_queue + self._shutdown_event = shutdown_event + self._warmup_ms = warmup_ms + self._thread = threading.Thread(target=self._run, daemon=True) + + def start(self) -> None: + self._thread.start() + + def join(self, timeout: float | None = None) -> None: + self._thread.join(timeout) + + def _run(self) -> None: + if self._warmup_ms: + time.sleep(self._warmup_ms / 1000.0) + self._result_queue.put({"kind": "ready"}) + while not self._shutdown_event.is_set(): + try: + item = self._job_queue.get(timeout=0.1) + except queue.Empty: + continue + if item is None: + return + try: + result = self._generator.generate(item["request"]) + self._result_queue.put({ + "kind": "result", + "job_id": item["job_id"], + "result": result, + }) + except Exception as exc: # pragma: no cover - defensive + self._result_queue.put({ + "kind": "error", + "job_id": item["job_id"], + "error": repr(exc), + }) + + +def _thread_worker_factory(generator_builder): + """Return a WorkerFactory that uses thread workers instead of procs.""" + + def factory( + *, + gpu_id: int, + generator_config: GeneratorConfig, + warmup_config: WarmupConfig, + ) -> _WorkerHandle: + ctx = mp.get_context("spawn") + job_queue: mp.Queue = ctx.Queue() + result_queue: mp.Queue = ctx.Queue() + shutdown_event = threading.Event() + ready = threading.Event() + boot_ok = threading.Event() + mp_shutdown = ctx.Event() + + generator = generator_builder(gpu_id) + worker = _ThreadWorker( + generator, + job_queue=job_queue, + result_queue=result_queue, + shutdown_event=shutdown_event, + ) + worker.start() + + # Drain the ready sentinel from the queue in the same way the + # real factory does (thread waiter populates ``ready`` / + # ``boot_ok``). + def _await_ready() -> None: + while True: + try: + msg = result_queue.get(timeout=1.0) + except queue.Empty: + if shutdown_event.is_set(): + return + continue + if msg.get("kind") == "ready": + boot_ok.set() + ready.set() + return + if msg.get("kind") == "error": + ready.set() + return + + threading.Thread(target=_await_ready, daemon=True).start() + + class _FakeProcess: + def __init__(self, stop: threading.Event, worker: _ThreadWorker): + self._stop = stop + self._worker = worker + + def is_alive(self) -> bool: + return self._worker._thread.is_alive() + + def join(self, timeout: float | None = None) -> None: + self._stop.set() + self._worker.join(timeout) + + def kill(self) -> None: + self._stop.set() + + fake_process = _FakeProcess(shutdown_event, worker) + + return _WorkerHandle( + process=fake_process, # type: ignore[arg-type] + job_queue=job_queue, + result_queue=result_queue, + gpu_id=gpu_id, + worker_id=f"gpu{gpu_id}-fake", + ready=ready, + boot_ok=boot_ok, + shutdown_event=mp_shutdown, + ) + + return factory + + +@pytest.fixture +def pool_factory(): + """Provide a SubprocessGpuPool built against thread workers.""" + + async def _build(num_workers: int = 2): + pool = SubprocessGpuPool( + generator_config=GeneratorConfig(model_path="/models/fake"), + pool_config=GpuPoolConfig(num_workers=num_workers), + warmup_config=WarmupConfig(enabled=False), + worker_factory=_thread_worker_factory( + lambda gpu_id: _MockGenerator()), + ) + await pool.start() + return pool + + return _build + + +class TestSubprocessGpuPool: + + def test_start_spawns_requested_workers(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=3) + try: + h = pool.health() + assert h.total_workers == 3 + assert h.available_workers == 3 + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_acquire_decrements_available(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=2) + try: + a = await pool.acquire("sess-a") + assert a.worker_id.endswith("-fake") + assert pool.health().available_workers == 1 + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_acquire_timeout_when_all_busy(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=1) + try: + await pool.acquire("sess-a") + with pytest.raises(PoolAcquireTimeout): + await pool.acquire("sess-b", timeout=0.1) + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_run_returns_worker_result(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=1) + try: + await pool.acquire("sess-a") + result = await pool.run( + "sess-a", GenerationRequest(prompt="hello")) + assert result["prompt_echo"] == "hello" + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_release_returns_worker_to_pool(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=1) + try: + await pool.acquire("sess-a") + await pool.release("sess-a") + # Now a second acquire should succeed without timeout. + a = await pool.acquire("sess-b", timeout=1.0) + assert a.worker_id.endswith("-fake") + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_sticky_binding_across_multiple_runs(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=2) + try: + a1 = await pool.acquire("sess-a") + a2 = await pool.acquire("sess-a") + assert a1.worker_id == a2.worker_id + # Two runs land on the same worker. + await pool.run("sess-a", GenerationRequest(prompt="1")) + await pool.run("sess-a", GenerationRequest(prompt="2")) + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_run_without_acquire_raises(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=1) + try: + with pytest.raises(RuntimeError): + await pool.run( + "sess-x", GenerationRequest(prompt="x")) + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_shutdown_is_idempotent(self, pool_factory): + async def run(): + pool = await pool_factory(num_workers=2) + await pool.shutdown() + await pool.shutdown() # no raise + + asyncio.run(run()) + + +class TestSubprocessGpuPoolFailureModes: + """Coverage for boot/runtime failures the pool has to absorb.""" + + def test_failed_boot_excluded_from_available(self): + """A worker whose factory leaves ``boot_ok`` unset must not be + handed out by ``acquire``. Otherwise a session lands on a dead + worker and ``run`` blocks forever on the missing result.""" + + def factory_with_one_failure(*, gpu_id, generator_config, warmup_config): + ctx = mp.get_context("spawn") + job_queue: mp.Queue = ctx.Queue() + result_queue: mp.Queue = ctx.Queue() + ready = threading.Event() + boot_ok = threading.Event() + mp_shutdown = ctx.Event() + ready.set() + # gpu_id 0 boots fine; gpu_id 1 fails (boot_ok stays clear). + if gpu_id == 0: + boot_ok.set() + + class _AliveProcess: + def is_alive(self) -> bool: + return True + def join(self, timeout: float | None = None) -> None: + return + def kill(self) -> None: + return + + return _WorkerHandle( + process=_AliveProcess(), # type: ignore[arg-type] + job_queue=job_queue, + result_queue=result_queue, + gpu_id=gpu_id, + worker_id=f"gpu{gpu_id}", + ready=ready, + boot_ok=boot_ok, + shutdown_event=mp_shutdown, + ) + + async def run(): + pool = SubprocessGpuPool( + generator_config=GeneratorConfig(model_path="/m"), + pool_config=GpuPoolConfig(num_workers=2), + warmup_config=WarmupConfig(enabled=False), + worker_factory=factory_with_one_failure, + ) + await pool.start() + try: + # Only worker 0 booted; only one slot should be available. + assert pool.health().available_workers == 1 + a = await pool.acquire("sess-a", timeout=0.1) + assert a.worker_id == "gpu0" + with pytest.raises(PoolAcquireTimeout): + await pool.acquire("sess-b", timeout=0.1) + finally: + await pool.shutdown() + + asyncio.run(run()) + + def test_release_skips_dead_worker(self, pool_factory): + """If a worker died while bound, releasing the session must not + return its slot to the available queue — a later acquire would + hand the dead slot to a new session.""" + + async def run(): + pool = await pool_factory(num_workers=1) + try: + await pool.acquire("sess-a") + # Simulate the worker dying mid-session. + pool._workers[0].process._stop.set() + pool._workers[0].process._worker.join(timeout=1.0) + await pool.release("sess-a") + # Available queue must remain empty. + assert pool.health().available_workers == 0 + finally: + await pool.shutdown() + + asyncio.run(run()) diff --git a/fastvideo/tests/entrypoints/streaming/test_worker.py b/fastvideo/tests/entrypoints/streaming/test_worker.py new file mode 100644 index 000000000..2a448636c --- /dev/null +++ b/fastvideo/tests/entrypoints/streaming/test_worker.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit coverage for :mod:`fastvideo.entrypoints.streaming.worker`. + +The full ``worker_main`` loop runs in a subprocess and is exercised via +``test_gpu_pool.py``'s subprocess integration tests. This file covers +the in-process pieces: + +* the two-segment warmup feeds segment 1's continuation state into + segment 2 so both compile branches are primed before the worker + reports ready +* result-shape extractors handle both attribute-style and dict-style + generator returns +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.api.schema import ( + ContinuationState, + GenerationRequest, + WarmupConfig, +) +from fastvideo.entrypoints.streaming.worker import ( + _extract_continuation_state, + _warmup_worker, +) + + +@dataclass +class _RecordingGenerator: + """Captures the requests passed to ``generate`` and returns a + canned :class:`ContinuationState` on the first call so the warmup + can feed it into the second call. + """ + + state_to_return: ContinuationState | None = field(default_factory=lambda: ContinuationState( + kind="ltx2.v1", + payload={"schema_version": 1, "segment_index": 1}, + )) + requests: list[GenerationRequest] = field(default_factory=list) + + def generate(self, request: GenerationRequest) -> dict[str, Any]: + self.requests.append(request) + return {"frames": [], "state": self.state_to_return} + + +class TestWarmupTwoSegment: + + def test_warmup_runs_segment_one_then_segment_two_with_returned_state(self) -> None: + gen = _RecordingGenerator() + _warmup_worker(gen, WarmupConfig(enabled=True, prompt="warm")) + + assert len(gen.requests) == 2 + + seg1 = gen.requests[0] + assert seg1.state is None + assert seg1.output.return_state is True + + seg2 = gen.requests[1] + assert seg2.state is gen.state_to_return + + def test_warmup_passes_through_when_no_state_returned(self) -> None: + gen = _RecordingGenerator(state_to_return=None) + _warmup_worker(gen, WarmupConfig(enabled=True, prompt="warm")) + + assert len(gen.requests) == 2 + assert gen.requests[1].state is None + + +class TestExtractContinuationState: + + def test_extracts_from_attribute(self) -> None: + + class _R: + state = ContinuationState(kind="k", payload={}) + + assert _extract_continuation_state(_R()).kind == "k" + + def test_extracts_from_dict(self) -> None: + state = ContinuationState(kind="k", payload={}) + assert _extract_continuation_state({"state": state}) is state + + def test_returns_none_for_missing(self) -> None: + assert _extract_continuation_state({}) is None + assert _extract_continuation_state(object()) is None