[feat] Improve API: streaming server GpuPool + worker subprocess (#1257)
Co-authored-by: Junda (David) Su <90978028+Davids048@users.noreply.github.com> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com>
This commit is contained in:
William Lin
co-authored by
Junda Su
Matthew Noto
XOR-op
Zhang Peiyuan
Junda Su
Matthew Noto
XOR-op
Zhang Peiyuan
parent
140bd1a6cf
commit
eb0a41528a
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
@@ -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())
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user