hao-ai-lab-FastVideo/fastvideo/tests/entrypoints/streaming/test_gpu_pool.py

# 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())