hao-ai-lab-FastVideo/fastvideo/entrypoints/streaming/router/main.py

# 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",
]