[feat] Improve API: streaming router (multi-replica load balancer + ws proxy) (#1286)
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
eb3a394224
commit
2aaeee2ab8
@@ -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
|
||||
|
||||
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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/"
|
||||
@@ -145,6 +145,7 @@ prompt-enhancer = [
|
||||
streaming = [
|
||||
"fastvideo[prompt-enhancer]",
|
||||
"fastvideo[prompt-safety]",
|
||||
"websockets",
|
||||
]
|
||||
|
||||
rocm = [
|
||||
|
||||
Reference in New Issue
Block a user