diff --git a/fastvideo/entrypoints/cli/main.py b/fastvideo/entrypoints/cli/main.py index 3d753089d..a7a66f677 100644 --- a/fastvideo/entrypoints/cli/main.py +++ b/fastvideo/entrypoints/cli/main.py @@ -3,6 +3,8 @@ from fastvideo.entrypoints.cli.cli_types import CLISubcommand from fastvideo.entrypoints.cli.generate import cmd_init as generate_cmd_init from fastvideo.utils import FlexibleArgumentParser +from fastvideo.entrypoints.cli.router_serve import ( + cmd_init as router_serve_cmd_init, ) from fastvideo.entrypoints.cli.serve import cmd_init as serve_cmd_init from fastvideo.entrypoints.cli.bench import cmd_init as bench_cmd_init @@ -12,6 +14,7 @@ def cmd_init() -> list[CLISubcommand]: commands = [] commands.extend(generate_cmd_init()) commands.extend(serve_cmd_init()) + commands.extend(router_serve_cmd_init()) commands.extend(bench_cmd_init()) return commands diff --git a/fastvideo/entrypoints/cli/router_serve.py b/fastvideo/entrypoints/cli/router_serve.py new file mode 100644 index 000000000..0bd37ee6b --- /dev/null +++ b/fastvideo/entrypoints/cli/router_serve.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""``fastvideo router-serve`` CLI subcommand. + +Launches the streaming router from a YAML config. Separate from +``fastvideo serve`` because the router is an orthogonal process: it +fronts one or more running servers rather than hosting a generator +itself. +""" +from __future__ import annotations + +import argparse +import os +from typing import cast + +from fastvideo.api.parser import load_raw_config +from fastvideo.entrypoints.cli.cli_types import CLISubcommand +from fastvideo.entrypoints.streaming.router.config import ( + ReplicaEndpoint, + RouterConfig, +) +from fastvideo.logger import init_logger +from fastvideo.utils import FlexibleArgumentParser + +logger = init_logger(__name__) + + +class RouterServeSubcommand(CLISubcommand): + """Start the multi-replica WebSocket router.""" + + def __init__(self) -> None: + self.name = "router-serve" + super().__init__() + + def cmd(self, args: argparse.Namespace) -> None: + config = _load_router_config(args.config) + logger.info( + "router listening on %s:%d (%d replicas, %d primary)", + config.host, + config.port, + len(config.replicas), + sum(1 for r in config.replicas if r.primary), + ) + from fastvideo.entrypoints.streaming.router.main import run_router + + run_router(config) + + def validate(self, args: argparse.Namespace) -> None: + if not args.config: + raise ValueError("fastvideo router-serve requires --config PATH") + if not os.path.exists(args.config): + raise ValueError(f"Router config file not found: {args.config}") + + def subparser_init( + self, + subparsers: argparse._SubParsersAction, + ) -> FlexibleArgumentParser: + parser = subparsers.add_parser( + "router-serve", + help="Start the streaming router (multi-replica load balancer)", + usage="fastvideo router-serve --config ROUTER_CONFIG", + ) + parser.add_argument( + "--config", + type=str, + default="", + required=False, + help="Path to a YAML/JSON router config. Required.", + ) + return cast(FlexibleArgumentParser, parser) + + +def _load_router_config(path: str) -> RouterConfig: + raw = load_raw_config(path) + router_raw = raw.get("router") if isinstance(raw, dict) else None + if not isinstance(router_raw, dict): + raise ValueError(f"Router config {path!r} must have a top-level `router:` block") + + replicas_raw = router_raw.get("replicas", []) + if not isinstance(replicas_raw, list): + raise ValueError(f"router.replicas must be a list, got {type(replicas_raw).__name__}") + replicas = [] + for i, r in enumerate(replicas_raw): + if not isinstance(r, dict): + raise ValueError(f"router.replicas[{i}] must be a mapping, got {type(r).__name__}") + url = r.get("url") + if not url: + raise ValueError(f"router.replicas[{i}] is missing required key 'url'") + replicas.append( + ReplicaEndpoint( + url=url, + name=r.get("name"), + primary=bool(r.get("primary", False)), + weight=float(r.get("weight", 1.0)), + )) + if not replicas: + raise ValueError("Router config must list at least one replica under `router.replicas`") + + health_check = router_raw.get("health_check") or {} + return RouterConfig( + host=str(router_raw.get("host", "0.0.0.0")), + port=int(router_raw.get("port", 9000)), + replicas=replicas, + health_check_path=str(health_check.get("path", "/health")), + health_check_interval_seconds=float(health_check.get("interval_seconds", 5.0)), + health_check_timeout_seconds=float(health_check.get("timeout_seconds", 2.0)), + failure_threshold=int(health_check.get("failure_threshold", 3)), + recovery_threshold=int(health_check.get("recovery_threshold", 2)), + ) + + +def cmd_init() -> list[CLISubcommand]: + return [RouterServeSubcommand()] + + +__all__ = ["RouterServeSubcommand", "cmd_init"] diff --git a/fastvideo/entrypoints/streaming/router/__init__.py b/fastvideo/entrypoints/streaming/router/__init__.py new file mode 100644 index 000000000..25394caa6 --- /dev/null +++ b/fastvideo/entrypoints/streaming/router/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Multi-replica load balancer + WebSocket proxy for the streaming server. + +Sits in front of one-or-more streaming-server replicas and forwards +WebSocket sessions to a healthy primary, with failover to secondaries. +Kept in-repo under ``fastvideo/entrypoints/streaming/router/`` per the +PR plan's default; the alternative (separate package) is an open +question deferred to review. +""" +from fastvideo.entrypoints.streaming.router.registry import ( + Replica, + ReplicaHealth, + ReplicaRegistry, + ReplicaStatus, +) +from fastvideo.entrypoints.streaming.router.config import RouterConfig +from fastvideo.entrypoints.streaming.router.main import build_router_app, run_router + +__all__ = [ + "Replica", + "ReplicaHealth", + "ReplicaRegistry", + "ReplicaStatus", + "RouterConfig", + "build_router_app", + "run_router", +] diff --git a/fastvideo/entrypoints/streaming/router/config.py b/fastvideo/entrypoints/streaming/router/config.py new file mode 100644 index 000000000..496e2e489 --- /dev/null +++ b/fastvideo/entrypoints/streaming/router/config.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Typed router configuration.""" +from __future__ import annotations + +from dataclasses import dataclass, field +from urllib.parse import urlparse + + +@dataclass +class ReplicaEndpoint: + """One backend replica the router can route to.""" + + url: str + """HTTP base URL, e.g. ``http://host:8000``. WebSocket URL is + derived automatically by replacing the scheme.""" + name: str | None = None + primary: bool = False + """``True`` = prefer this replica over others in steady state.""" + weight: float = 1.0 + + +@dataclass +class RouterConfig: + """Typed router config loaded from a YAML file. + + Example:: + + router: + host: 0.0.0.0 + port: 9000 + replicas: + - url: http://streamer-a:8000 + primary: true + - url: http://streamer-b:8000 + health_check: + path: /health + interval_seconds: 5 + failure_threshold: 3 + + Validation runs in ``__post_init__``: empty replicas, non-positive + intervals/timeouts, thresholds < 1, non-http(s) URLs, and more than + one primary all raise ``ValueError`` so misconfigurations surface at + load time rather than as confusing runtime failures. + """ + + host: str = "0.0.0.0" + port: int = 9000 + replicas: list[ReplicaEndpoint] = field(default_factory=list) + health_check_path: str = "/health" + health_check_interval_seconds: float = 5.0 + health_check_timeout_seconds: float = 2.0 + failure_threshold: int = 3 + recovery_threshold: int = 2 + + def __post_init__(self) -> None: + if not self.replicas: + raise ValueError("RouterConfig.replicas must list at least one replica") + if self.health_check_interval_seconds <= 0: + raise ValueError(f"health_check_interval_seconds must be > 0, got {self.health_check_interval_seconds}") + if self.health_check_timeout_seconds <= 0: + raise ValueError(f"health_check_timeout_seconds must be > 0, got {self.health_check_timeout_seconds}") + if self.failure_threshold < 1: + raise ValueError(f"failure_threshold must be >= 1, got {self.failure_threshold}") + if self.recovery_threshold < 1: + raise ValueError(f"recovery_threshold must be >= 1, got {self.recovery_threshold}") + seen_urls: set[str] = set() + for replica in self.replicas: + if not replica.url.startswith(("http://", "https://")): + raise ValueError(f"ReplicaEndpoint.url must start with http:// or https://, got {replica.url!r}") + parsed = urlparse(replica.url) + if parsed.path not in ("", "/"): + raise ValueError(f"ReplicaEndpoint.url must be a base host[:port] URL without a path; " + f"got {replica.url!r} with path {parsed.path!r}. The router appends " + "`/health` and `/v1/stream` itself.") + if parsed.query or parsed.fragment: + raise ValueError(f"ReplicaEndpoint.url must not include query/fragment; got {replica.url!r}") + if replica.url in seen_urls: + raise ValueError(f"Duplicate ReplicaEndpoint.url {replica.url!r}; " + "router selection keys by URL so duplicates would silently collapse") + seen_urls.add(replica.url) + primaries = sum(1 for r in self.replicas if r.primary) + if primaries > 1: + raise ValueError(f"RouterConfig allows at most one primary replica; got {primaries}. " + "Multi-primary load distribution is deferred — promote one replica to " + "primary and treat the rest as secondaries.") + + +__all__ = ["ReplicaEndpoint", "RouterConfig"] diff --git a/fastvideo/entrypoints/streaming/router/main.py b/fastvideo/entrypoints/streaming/router/main.py new file mode 100644 index 000000000..3c6d65b34 --- /dev/null +++ b/fastvideo/entrypoints/streaming/router/main.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Router FastAPI entry point. + +Exposes the same ``/v1/stream`` WebSocket path the backend servers do, +accepts a client, picks a healthy replica from the registry, and +proxies frames bidirectionally. + +PR 7.9 ships the minimum-viable shape: explicit replica list, single +primary, JSON + binary passthrough in both directions, and a +``/status`` endpoint for operators. Sticky-session routing (so a +reconnect lands on the same backend) is left for a follow-up. +""" +from __future__ import annotations + +import asyncio +import contextlib +from dataclasses import dataclass + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from fastapi.responses import JSONResponse + +from fastvideo.entrypoints.streaming.router.config import RouterConfig +from fastvideo.entrypoints.streaming.router.registry import ( + ReplicaRegistry, + run_health_check_loop, +) +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass +class _RouterState: + config: RouterConfig + registry: ReplicaRegistry + stop_event: asyncio.Event + health_task: asyncio.Task | None = None + + +def build_router_app( + config: RouterConfig, + *, + registry: ReplicaRegistry | None = None, +) -> FastAPI: + """Build the router FastAPI app. + + ``registry`` can be injected for tests; defaults to one built from + ``config.replicas``. + """ + registry = registry or ReplicaRegistry(config.replicas) + state = _RouterState( + config=config, + registry=registry, + stop_event=asyncio.Event(), + ) + + @contextlib.asynccontextmanager + async def _lifespan(_app: FastAPI): + state.health_task = asyncio.create_task( + run_health_check_loop( + registry=state.registry, + config=state.config, + stop_event=state.stop_event, + )) + try: + yield + finally: + state.stop_event.set() + if state.health_task is not None: + with contextlib.suppress(asyncio.CancelledError): + await state.health_task + + app = FastAPI(title="FastVideo Streaming Router", lifespan=_lifespan) + + @app.get("/status") + async def _status() -> JSONResponse: + return JSONResponse({ + "replicas": [{ + "url": r.url, + "primary": r.primary, + "status": r.health.status.value, + "last_ok_at": r.health.last_ok_at, + "last_latency_ms": r.health.last_latency_ms, + "consecutive_failures": r.health.consecutive_failures, + } for r in state.registry.all()], + }) + + @app.websocket("/v1/stream") + async def _proxy(websocket: WebSocket) -> None: + await websocket.accept() + replica = state.registry.select() + if replica is None: + await websocket.send_json({ + "type": "error", + "code": "gpu_unavailable", + "message": "router: no healthy replica available", + "retryable": True, + }) + await websocket.close(code=1013, reason="no_healthy_replica") + return + + ws_url = _websocket_url_for(replica.url) + try: + await _bridge_session(websocket, ws_url) + except WebSocketDisconnect: + logger.info("router: client disconnected") + except Exception as exc: + logger.exception("router: bridge failed: %s", exc) + with contextlib.suppress(RuntimeError): + await websocket.send_json({ + "type": "error", + "code": "worker_failed", + "message": f"router bridge failed: {exc}", + "retryable": True, + }) + with contextlib.suppress(RuntimeError): + await websocket.close(code=1011) + + app.state.router_state = state + return app + + +def run_router(config: RouterConfig) -> None: # pragma: no cover - CLI + import uvicorn + + app = build_router_app(config) + uvicorn.run(app, host=config.host, port=config.port) + + +async def _bridge_session( + client_ws: WebSocket, + backend_ws_url: str, +) -> None: + """Connect to backend and shuttle messages in both directions. + + Uses ``websockets`` for the backend side; imported lazily to keep + the router's import graph small for users who only want the server. + + Cancellation: when either direction completes (client disconnect, + backend close, exception), the other is cancelled explicitly and + both are drained before returning. Unexpected exceptions from the + direction that completed first are re-raised; normal disconnect + paths (``WebSocketDisconnect``, ``ConnectionClosed``, + ``CancelledError``) are swallowed. + """ + try: + import websockets + except ImportError as exc: # pragma: no cover - optional extra + raise RuntimeError("router requires the `websockets` package for backend proxying") from exc + + async with websockets.connect(backend_ws_url + "/v1/stream") as backend_ws: + c2b = asyncio.create_task(_forward_client_to_backend(client_ws, backend_ws)) + b2c = asyncio.create_task(_forward_backend_to_client(backend_ws, client_ws)) + try: + done, _pending = await asyncio.wait( + {c2b, b2c}, + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + for task in (c2b, b2c): + if not task.done(): + task.cancel() + await asyncio.gather(c2b, b2c, return_exceptions=True) + for task in done: + task_exc = task.exception() + if task_exc is not None and not _is_normal_disconnect(task_exc): + raise task_exc + + +def _is_normal_disconnect(exc: BaseException) -> bool: + """Whether ``exc`` is a routine WebSocket teardown vs a real bridge fault.""" + if isinstance(exc, asyncio.CancelledError | WebSocketDisconnect): + return True + name = type(exc).__name__ + # websockets.exceptions.ConnectionClosed{,OK,Error} all subclass + # WebSocketException; check by name to avoid the lazy-import dance. + return name.startswith("ConnectionClosed") + + +async def _forward_client_to_backend(client_ws: WebSocket, backend_ws) -> None: + try: + while True: + msg = await client_ws.receive() + if msg.get("type") == "websocket.disconnect": + break + if "text" in msg and msg["text"] is not None: + await backend_ws.send(msg["text"]) + elif "bytes" in msg and msg["bytes"] is not None: + await backend_ws.send(msg["bytes"]) + finally: + with contextlib.suppress(Exception): + await backend_ws.close() + + +async def _forward_backend_to_client(backend_ws, client_ws: WebSocket) -> None: + try: + async for frame in backend_ws: + if isinstance(frame, bytes): + await client_ws.send_bytes(frame) + else: + await client_ws.send_text(frame) + finally: + with contextlib.suppress(Exception): + await client_ws.close() + + +def _websocket_url_for(http_url: str) -> str: + if http_url.startswith("https://"): + return "wss://" + http_url[len("https://"):] + if http_url.startswith("http://"): + return "ws://" + http_url[len("http://"):] + return http_url + + +__all__ = [ + "build_router_app", + "run_router", +] diff --git a/fastvideo/entrypoints/streaming/router/registry.py b/fastvideo/entrypoints/streaming/router/registry.py new file mode 100644 index 000000000..bbd3fc6d1 --- /dev/null +++ b/fastvideo/entrypoints/streaming/router/registry.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Replica registry + health-check loop. + +The registry tracks the set of known backend replicas and their live +health. The router consults it for "pick a backend for this session" +decisions and a background task updates it from periodic HTTP probes. + +State machine per replica:: + + HEALTHY ──(N consecutive failures)──▶ UNHEALTHY + ▲ │ + └──────(M consecutive successes)──────┘ + +Where N = :attr:`RouterConfig.failure_threshold` and +M = :attr:`RouterConfig.recovery_threshold`. +""" +from __future__ import annotations + +import asyncio +import contextlib +import enum +import time +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.entrypoints.streaming.router.config import ( + ReplicaEndpoint, + RouterConfig, +) +from fastvideo.logger import init_logger + +HttpProbe = Any +"""Structural alias for health-probe callables. Concrete signature is +``async def __call__(url: str, *, timeout: float) -> tuple[float, +str | None]``; typing.Callable cannot express keyword-only parameters, +so duck-typing is the pragmatic compromise.""" + +logger = init_logger(__name__) + + +class ReplicaStatus(enum.Enum): + UNKNOWN = "unknown" + HEALTHY = "healthy" + UNHEALTHY = "unhealthy" + + +@dataclass +class ReplicaHealth: + status: ReplicaStatus = ReplicaStatus.UNKNOWN + last_ok_at: float | None = None + last_failure_at: float | None = None + consecutive_failures: int = 0 + consecutive_successes: int = 0 + last_latency_ms: float | None = None + + +@dataclass +class Replica: + endpoint: ReplicaEndpoint + health: ReplicaHealth = field(default_factory=ReplicaHealth) + + @property + def url(self) -> str: + return self.endpoint.url + + @property + def primary(self) -> bool: + return self.endpoint.primary + + @property + def is_healthy(self) -> bool: + return self.health.status is ReplicaStatus.HEALTHY + + +class ReplicaRegistry: + """Stateful map of replica URL → :class:`Replica`. + + Selection favors primary replicas when healthy; otherwise the first + healthy non-primary is returned. When none are healthy, the + registry returns ``None`` so the router can reject incoming + sessions with ``gpu_unavailable``. + """ + + def __init__(self, replicas: list[ReplicaEndpoint]) -> None: + if not replicas: + raise ValueError("ReplicaRegistry requires at least one replica") + self._replicas: dict[str, Replica] = {endpoint.url: Replica(endpoint=endpoint) for endpoint in replicas} + self._lock = asyncio.Lock() + + def all(self) -> list[Replica]: + return list(self._replicas.values()) + + def get(self, url: str) -> Replica | None: + return self._replicas.get(url) + + def primaries(self) -> list[Replica]: + return [r for r in self._replicas.values() if r.primary] + + def select(self) -> Replica | None: + """Pick the best healthy replica. + + Priority order: + + 1. The first healthy primary (insertion order). + 2. The first healthy non-primary (insertion order). + 3. ``None`` when nothing is healthy. + + This MVP picks the first match within each tier; it does NOT + load-balance across multiple healthy replicas of the same tier. + Round-robin and weighted distribution are deferred until a real + N-way active deployment exists. + """ + healthy_primaries = [r for r in self._replicas.values() if r.primary and r.is_healthy] + if healthy_primaries: + return healthy_primaries[0] + healthy = [r for r in self._replicas.values() if r.is_healthy] + if healthy: + return healthy[0] + return None + + async def record_success( + self, + replica: Replica, + *, + recovery_threshold: int, + latency_ms: float, + ) -> None: + async with self._lock: + h = replica.health + h.last_ok_at = time.time() + h.last_latency_ms = latency_ms + h.consecutive_failures = 0 + h.consecutive_successes += 1 + # State machine: UNKNOWN -> HEALTHY is immediate; only the + # UNHEALTHY -> HEALTHY transition is gated by recovery_threshold. + if h.status is ReplicaStatus.UNKNOWN: + logger.info("router: replica %s initial probe ok, marking HEALTHY", replica.url) + h.status = ReplicaStatus.HEALTHY + h.consecutive_successes = 0 + elif (h.status is ReplicaStatus.UNHEALTHY and h.consecutive_successes >= recovery_threshold): + logger.info("router: replica %s recovered to HEALTHY after %d successes", replica.url, + h.consecutive_successes) + h.status = ReplicaStatus.HEALTHY + h.consecutive_successes = 0 + + async def record_failure( + self, + replica: Replica, + *, + failure_threshold: int, + reason: str, + ) -> None: + async with self._lock: + h = replica.health + h.last_failure_at = time.time() + h.consecutive_successes = 0 + h.consecutive_failures += 1 + if (h.status is not ReplicaStatus.UNHEALTHY and h.consecutive_failures >= failure_threshold): + logger.warning("router: replica %s marked UNHEALTHY after %d failures: %s", replica.url, + h.consecutive_failures, reason) + h.status = ReplicaStatus.UNHEALTHY + + +async def run_health_check_loop( + registry: ReplicaRegistry, + config: RouterConfig, + *, + stop_event: asyncio.Event, + http_get: HttpProbe | None = None, +) -> None: + """Poll all replicas' health endpoints in parallel on a fixed interval. + + ``http_get`` is pluggable so unit tests can inject a deterministic + probe without hitting the network. The default builds a single + ``httpx.AsyncClient`` shared across the loop's lifetime so the + common case (steady polling against a stable replica set) reuses + TCP/TLS connections instead of paying handshake cost per probe. + + Probes within one polling cycle run concurrently via ``asyncio.gather`` + so a slow replica doesn't push the cycle past + ``health_check_interval_seconds``. + """ + if http_get is not None: + await _run_loop(registry, config, stop_event, http_get) + return + async with _build_default_probe(config) as probe: + await _run_loop(registry, config, stop_event, probe) + + +async def _run_loop( + registry: ReplicaRegistry, + config: RouterConfig, + stop_event: asyncio.Event, + http_get: Callable[..., Awaitable[tuple[float, str | None]]], +) -> None: + while not stop_event.is_set(): + replicas = registry.all() + results = await asyncio.gather( + *[ + http_get(replica.url + config.health_check_path, timeout=config.health_check_timeout_seconds) + for replica in replicas + ], + return_exceptions=True, + ) + for replica, result in zip(replicas, results, strict=True): + if isinstance(result, BaseException): + await registry.record_failure( + replica, + failure_threshold=config.failure_threshold, + reason=f"{type(result).__name__}: {result}", + ) + continue + status_ms, error = result + if error is None: + await registry.record_success( + replica, + recovery_threshold=config.recovery_threshold, + latency_ms=status_ms, + ) + else: + await registry.record_failure( + replica, + failure_threshold=config.failure_threshold, + reason=error, + ) + try: + await asyncio.wait_for( + stop_event.wait(), + timeout=config.health_check_interval_seconds, + ) + except asyncio.TimeoutError: + continue + + +@contextlib.asynccontextmanager +async def _build_default_probe( + config: RouterConfig, ) -> AsyncIterator[Callable[..., Awaitable[tuple[float, str | None]]]]: + try: + import httpx + except ImportError as exc: # pragma: no cover - optional extra + raise RuntimeError("router health checks require httpx; install with " + "`pip install fastvideo[streaming]` or `pip install httpx`") from exc + + async with httpx.AsyncClient(timeout=config.health_check_timeout_seconds) as client: + + async def probe(url: str, *, timeout: float) -> tuple[float, str | None]: + start = time.perf_counter() + try: + response = await client.get(url, timeout=timeout) + except Exception as exc: + return 0.0, f"{type(exc).__name__}: {exc}" + latency_ms = (time.perf_counter() - start) * 1000.0 + if response.status_code >= 400: + return latency_ms, f"HTTP {response.status_code}" + return latency_ms, None + + yield probe + + +__all__ = [ + "HttpProbe", + "Replica", + "ReplicaHealth", + "ReplicaRegistry", + "ReplicaStatus", + "run_health_check_loop", +] diff --git a/fastvideo/tests/entrypoints/streaming/test_router.py b/fastvideo/tests/entrypoints/streaming/test_router.py new file mode 100644 index 000000000..a1e92789d --- /dev/null +++ b/fastvideo/tests/entrypoints/streaming/test_router.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Router tests — registry semantics + health loop behavior. + +Avoids real WebSocket proxying (the bridge requires the ``websockets`` +package and two running uvicorn processes); test_server.py already +covers the end-to-end WS protocol against a direct backend. +""" +from __future__ import annotations + +import asyncio + +import pytest + +from fastvideo.entrypoints.streaming.router.config import ( + ReplicaEndpoint, + RouterConfig, +) +from fastvideo.entrypoints.streaming.router.registry import ( + ReplicaRegistry, + ReplicaStatus, + run_health_check_loop, +) + + +def _registry( + *, + num_primary: int = 1, + num_secondary: int = 1, +) -> ReplicaRegistry: + replicas = [ + ReplicaEndpoint( + url=f"http://primary-{i}:8000", + primary=True, + ) + for i in range(num_primary) + ] + [ + ReplicaEndpoint( + url=f"http://secondary-{i}:8000", + primary=False, + ) + for i in range(num_secondary) + ] + return ReplicaRegistry(replicas) + + +class TestReplicaRegistry: + + def test_requires_replicas(self): + with pytest.raises(ValueError): + ReplicaRegistry([]) + + def test_select_none_when_all_unknown(self): + assert _registry().select() is None + + def test_select_prefers_healthy_primary(self): + reg = _registry() + + async def promote(): + primary = reg.primaries()[0] + await reg.record_success( + primary, recovery_threshold=1, latency_ms=1.0) + # Also mark secondary healthy — primary should still win. + sec = next(r for r in reg.all() if not r.primary) + await reg.record_success( + sec, recovery_threshold=1, latency_ms=1.0) + return reg.select() + + pick = asyncio.run(promote()) + assert pick is not None + assert pick.primary + + def test_falls_back_to_secondary_when_primary_unhealthy(self): + reg = _registry() + + async def run(): + primary = reg.primaries()[0] + sec = next(r for r in reg.all() if not r.primary) + # Fail primary past threshold, succeed secondary. + for _ in range(3): + await reg.record_failure( + primary, failure_threshold=3, reason="mock") + await reg.record_success( + sec, recovery_threshold=1, latency_ms=1.0) + return reg.select() + + pick = asyncio.run(run()) + assert pick is not None + assert not pick.primary + + def test_failure_threshold_transitions_to_unhealthy(self): + reg = _registry() + primary = reg.primaries()[0] + + async def run(): + for _ in range(2): + await reg.record_failure( + primary, failure_threshold=3, reason="x") + assert primary.health.status is not ReplicaStatus.UNHEALTHY + await reg.record_failure( + primary, failure_threshold=3, reason="x") + return primary + + result = asyncio.run(run()) + assert result.health.status is ReplicaStatus.UNHEALTHY + assert result.health.consecutive_failures == 3 + + def test_recovery_threshold_returns_to_healthy(self): + reg = _registry() + primary = reg.primaries()[0] + + async def run(): + for _ in range(3): + await reg.record_failure( + primary, failure_threshold=3, reason="x") + assert primary.health.status is ReplicaStatus.UNHEALTHY + for _ in range(2): + await reg.record_success( + primary, recovery_threshold=2, latency_ms=5.0) + return primary + + result = asyncio.run(run()) + assert result.health.status is ReplicaStatus.HEALTHY + assert result.health.last_latency_ms == 5.0 + + def test_record_success_resets_failure_counter(self): + reg = _registry() + primary = reg.primaries()[0] + + async def run(): + await reg.record_failure( + primary, failure_threshold=10, reason="x") + await reg.record_success( + primary, recovery_threshold=1, latency_ms=1.0) + return primary + + result = asyncio.run(run()) + assert result.health.consecutive_failures == 0 + + +class TestHealthCheckLoop: + + def test_loop_transitions_replicas_on_probe_results(self): + config = RouterConfig( + replicas=[ + ReplicaEndpoint(url="http://a", primary=True), + ReplicaEndpoint(url="http://b"), + ], + health_check_interval_seconds=0.01, + failure_threshold=1, + recovery_threshold=1, + ) + reg = ReplicaRegistry(config.replicas) + stop_event = asyncio.Event() + + calls: list[str] = [] + + async def probe(url, *, timeout): + calls.append(url) + if "http://a" in url: + return 1.0, None + return 0.0, "mock failure" + + async def run() -> None: + task = asyncio.create_task(run_health_check_loop( + registry=reg, config=config, stop_event=stop_event, + http_get=probe, + )) + await asyncio.sleep(0.05) + stop_event.set() + await task + + asyncio.run(run()) + a = reg.get("http://a") + b = reg.get("http://b") + assert a is not None + assert b is not None + assert a.health.status is ReplicaStatus.HEALTHY + assert b.health.status is ReplicaStatus.UNHEALTHY + assert any("/health" in c for c in calls) + + +class TestRouterApp: + + def test_status_endpoint_lists_replicas(self): + from starlette.testclient import TestClient + + from fastvideo.entrypoints.streaming.router.main import build_router_app + + config = RouterConfig( + replicas=[ + ReplicaEndpoint(url="http://a", primary=True), + ReplicaEndpoint(url="http://b"), + ], + health_check_interval_seconds=60, # don't actually poll + ) + reg = ReplicaRegistry(config.replicas) + app = build_router_app(config, registry=reg) + client = TestClient(app) + response = client.get("/status") + body = response.json() + urls = {r["url"] for r in body["replicas"]} + assert urls == {"http://a", "http://b"} + # Initial status is UNKNOWN. + assert all(r["status"] == "unknown" for r in body["replicas"]) + + def test_ws_rejects_when_no_healthy_replica(self): + from starlette.testclient import TestClient + + from fastvideo.entrypoints.streaming.router.main import build_router_app + + config = RouterConfig( + replicas=[ReplicaEndpoint(url="http://a", primary=True)], + health_check_interval_seconds=60, + ) + reg = ReplicaRegistry(config.replicas) + app = build_router_app(config, registry=reg) + client = TestClient(app) + with client.websocket_connect("/v1/stream") as ws: + err = ws.receive_json() + assert err["type"] == "error" + assert err["code"] == "gpu_unavailable" + + +class TestUnknownToHealthyImmediate: + """Initial probe must promote UNKNOWN -> HEALTHY without waiting for recovery_threshold.""" + + def test_first_success_promotes_unknown(self): + reg = _registry(num_primary=1, num_secondary=0) + primary = reg.primaries()[0] + assert primary.health.status is ReplicaStatus.UNKNOWN + + async def run(): + await reg.record_success(primary, recovery_threshold=10, latency_ms=1.0) + + asyncio.run(run()) + assert primary.health.status is ReplicaStatus.HEALTHY + + def test_unhealthy_recovery_still_gated_by_threshold(self): + reg = _registry(num_primary=1, num_secondary=0) + primary = reg.primaries()[0] + + async def run(): + for _ in range(3): + await reg.record_failure(primary, failure_threshold=3, reason="x") + assert primary.health.status is ReplicaStatus.UNHEALTHY + await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0) + assert primary.health.status is ReplicaStatus.UNHEALTHY # 1/2 successes + await reg.record_success(primary, recovery_threshold=2, latency_ms=1.0) + assert primary.health.status is ReplicaStatus.HEALTHY # 2/2 successes + + asyncio.run(run()) + + +class TestConfigValidation: + """RouterConfig.__post_init__ rejects malformed configs.""" + + def test_rejects_path_in_url(self): + with pytest.raises(ValueError, match="without a path"): + RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000/api")]) + + def test_rejects_query_in_url(self): + with pytest.raises(ValueError, match="query/fragment"): + RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000?x=1")]) + + def test_rejects_fragment_in_url(self): + with pytest.raises(ValueError, match="query/fragment"): + RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000#frag")]) + + def test_rejects_duplicate_urls(self): + with pytest.raises(ValueError, match="Duplicate"): + RouterConfig(replicas=[ + ReplicaEndpoint(url="http://host:8000"), + ReplicaEndpoint(url="http://host:8000"), + ]) + + def test_accepts_trailing_slash(self): + # parsed.path == "/" should be allowed + cfg = RouterConfig(replicas=[ReplicaEndpoint(url="http://host:8000/")]) + assert cfg.replicas[0].url == "http://host:8000/" diff --git a/pyproject.toml b/pyproject.toml index 2d3c6b865..b47b7eae4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,6 +145,7 @@ prompt-enhancer = [ streaming = [ "fastvideo[prompt-enhancer]", "fastvideo[prompt-safety]", + "websockets", ] rocm = [