From cb082d8bbb2ca03902f660d6495ef52c026efb1b Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 5 May 2026 13:20:08 -0700 Subject: [PATCH] [feat] streaming: promote /healthz /readyz /status to public build_app Promote the Dreamverse-compatible health, readiness, and status probes into the public streaming package so integrators can compose them through build_app without product-local shims. Co-authored-by: Junda Su Co-authored-by: Matthew Noto Co-authored-by: XOR-op Co-authored-by: Zhang Peiyuan --- fastvideo/entrypoints/streaming/__init__.py | 6 + fastvideo/entrypoints/streaming/health.py | 187 ++++++++++++++ fastvideo/entrypoints/streaming/server.py | 21 +- .../streaming/test_health_routes.py | 237 ++++++++++++++++++ 4 files changed, 447 insertions(+), 4 deletions(-) create mode 100644 fastvideo/entrypoints/streaming/health.py create mode 100644 fastvideo/tests/entrypoints/streaming/test_health_routes.py diff --git a/fastvideo/entrypoints/streaming/__init__.py b/fastvideo/entrypoints/streaming/__init__.py index 207bb775c..e4e1730a4 100644 --- a/fastvideo/entrypoints/streaming/__init__.py +++ b/fastvideo/entrypoints/streaming/__init__.py @@ -17,6 +17,10 @@ from fastvideo.entrypoints.streaming.gpu_pool import ( PoolAcquireTimeout, SubprocessGpuPool, ) +from fastvideo.entrypoints.streaming.health import ( + build_health_router, + get_pool_status, +) from fastvideo.entrypoints.streaming.mock_server import ( MockGenerator, build_mock_app, @@ -52,6 +56,8 @@ __all__ = [ "PromptEnhancer", "PromptSafetyFilter", "SafetyDecision", + "build_health_router", + "get_pool_status", "SessionLogEvent", "SessionLogger", "build_mock_app", diff --git a/fastvideo/entrypoints/streaming/health.py b/fastvideo/entrypoints/streaming/health.py new file mode 100644 index 000000000..eb421b4a1 --- /dev/null +++ b/fastvideo/entrypoints/streaming/health.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Health, readiness, and status routes for streaming servers.""" +from __future__ import annotations + +from collections.abc import Callable, Mapping +from datetime import datetime, timezone +from typing import Any, Protocol, TypeGuard, TypeAlias + +from fastapi import APIRouter + +from fastvideo.entrypoints.streaming.gpu_pool import PoolHealth + +SERVICE_NAME = "ltx2-streaming-backend" + + +class _StatusPool(Protocol): + + def get_status(self) -> Mapping[str, Any]: + ... + + +class _HealthPool(Protocol): + + def health(self) -> PoolHealth: + ... + + +_PoolLike: TypeAlias = _StatusPool | _HealthPool +PoolRef: TypeAlias = _PoolLike | Callable[[], _PoolLike | None] | None + + +def build_health_router(pool: PoolRef = None) -> APIRouter: + """Build a router exposing streaming liveness/readiness endpoints. + + ``pool`` may be either a concrete pool or a zero-argument callable that + returns the current pool. The callable form lets product servers keep their + own lifespan-managed runtime singleton without adding public global state. + """ + router = APIRouter() + + @router.get("/health") + @router.get("/healthz") + async def get_healthz() -> dict[str, Any]: + """Liveness probe for process-level health.""" + return { + "status": "ok", + "service": SERVICE_NAME, + "ts": _utc_now_iso(), + } + + @router.get("/readyz") + async def get_readyz() -> dict[str, Any]: + """Readiness probe for router/load-balancer health checks.""" + status_payload = await get_pool_status(pool) + ready_workers = _ready_worker_count(status_payload) + return { + "status": "ready" if ready_workers > 0 else "warming", + "service": SERVICE_NAME, + "ready_gpu_workers": ready_workers, + "total_gpus": _as_int(status_payload.get("total_gpus")), + "available_gpus": _as_int(status_payload.get("available_gpus")), + "warmup_successful_gpus": _as_int(status_payload.get("warmup_successful_gpus")), + "warmup_failed_gpus": _as_int(status_payload.get("warmup_failed_gpus")), + "queue_size": _as_int(status_payload.get("queue_size")), + "ts": _utc_now_iso(), + } + + @router.get("/status") + async def get_status() -> dict[str, Any]: + """Get the current status of the GPU pool.""" + return await get_pool_status(pool) + + return router + + +async def get_pool_status(pool: PoolRef = None) -> dict[str, Any]: + """Return the generic GPU pool status payload used by ``/status``.""" + resolved = _resolve_pool(pool) + if resolved is None: + return _zero_pool_status() + if _has_get_status(resolved): + return dict(resolved.get_status()) + if _has_health(resolved): + return _status_from_health(resolved.health()) + return _zero_pool_status() + + +def _resolve_pool(pool: PoolRef) -> _PoolLike | None: + if callable(pool): + return pool() + return pool + + +def _has_get_status(pool: _PoolLike) -> TypeGuard[_StatusPool]: + return callable(getattr(pool, "get_status", None)) + + +def _has_health(pool: _PoolLike) -> TypeGuard[_HealthPool]: + return callable(getattr(pool, "health", None)) + + +def _zero_pool_status() -> dict[str, Any]: + return { + "total_gpus": 0, + "available_gpus": 0, + "queue_size": 0, + "warmup_enabled": False, + "warmup_successful_gpus": 0, + "warmup_failed_gpus": 0, + "gpu_status": {}, + } + + +def _status_from_health(health: PoolHealth) -> dict[str, Any]: + total_workers = max(0, int(health.total_workers)) + available_workers = max(0, min(int(health.available_workers), total_workers)) + active_sessions = max(0, int(health.active_sessions)) + gpu_status: dict[int, dict[str, Any]] = {} + remaining_active_sessions = active_sessions + for gpu_id in range(total_workers): + is_available = gpu_id < available_workers + client_count = 0 if is_available or remaining_active_sessions <= 0 else 1 + remaining_active_sessions -= client_count + gpu_status[gpu_id] = { + "ready": True, + "available": is_available, + "client_count": client_count, + "current_model_id": None, + "process_alive": True, + "warmup_enabled": False, + "warmup_success": True, + "warmup_error": None, + "warmup_timings": {}, + } + return { + "total_gpus": total_workers, + "available_gpus": available_workers, + "queue_size": max(0, int(health.queued_sessions)), + "warmup_enabled": False, + "warmup_successful_gpus": total_workers, + "warmup_failed_gpus": 0, + "gpu_status": gpu_status, + } + + +def _ready_worker_count(status_payload: Mapping[str, Any]) -> int: + gpu_status = status_payload.get("gpu_status") + if isinstance(gpu_status, Mapping): + ready_workers = 0 + for value in gpu_status.values(): + if not isinstance(value, Mapping): + continue + ready = bool(value.get("ready")) + process_alive = bool(value.get("process_alive", True)) + if ready and process_alive: + ready_workers += 1 + if ready_workers > 0: + return ready_workers + successful = _as_int(status_payload.get("warmup_successful_gpus")) + if successful > 0: + return successful + total = _as_int(status_payload.get("total_gpus")) + failed = _as_int(status_payload.get("warmup_failed_gpus")) + if total > 0 and failed < total: + return total - failed + return 0 + + +def _as_int(value: Any) -> int: + if isinstance(value, bool): + return int(value) + if isinstance(value, int): + return value + if isinstance(value, float): + return int(value) + return 0 + + +def _utc_now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +__all__ = [ + "SERVICE_NAME", + "build_health_router", + "get_pool_status", +] diff --git a/fastvideo/entrypoints/streaming/server.py b/fastvideo/entrypoints/streaming/server.py index 0d618b92d..cabecf614 100644 --- a/fastvideo/entrypoints/streaming/server.py +++ b/fastvideo/entrypoints/streaming/server.py @@ -7,7 +7,7 @@ import contextlib import os import time from dataclasses import dataclass -from typing import Any, Protocol +from typing import Any, Literal, Protocol from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse @@ -56,6 +56,7 @@ from fastvideo.entrypoints.streaming.gpu_pool import ( InProcessGpuPool, PoolAcquireTimeout, ) +from fastvideo.entrypoints.streaming.health import build_health_router from fastvideo.entrypoints.streaming.session_store import ( InMemorySessionStore, SessionStore, @@ -68,6 +69,15 @@ logger = init_logger(__name__) # RFC 6455 WebSocket close codes used by the server. _WS_CLOSE_UNSUPPORTED_DATA = 1003 _WS_CLOSE_TRY_AGAIN_LATER = 1013 +_ErrorCode = Literal[ + "session_rejected", + "invalid_message", + "preset_mismatch", + "gpu_unavailable", + "worker_failed", + "upstream_timeout", + "internal_error", +] class _GeneratorProto(Protocol): @@ -104,6 +114,7 @@ def build_app( if serve_config.streaming is None: raise ValueError("ServeConfig.streaming must be set to launch the streaming " "server; got None. Add a `streaming:` block to your serve config.") + streaming = serve_config.streaming if (generator is None) == (pool is None): raise ValueError("build_app requires exactly one of `generator` or `pool`") @@ -130,9 +141,11 @@ def build_app( return JSONResponse({ "status": "ok", "sessions": len(state.sessions), - "stream_mode": state.serve_config.streaming.stream_mode, + "stream_mode": streaming.stream_mode, }) + app.include_router(build_health_router(pool)) + @app.websocket("/v1/stream") async def _stream(websocket: WebSocket) -> None: await websocket.accept() @@ -507,7 +520,7 @@ def _apply_toggle(session: Session, message: Any) -> None: session.curated_prompts = list(message.seed_prompts) -def _extract_frames(result: Any) -> list: +def _extract_frames(result: Any) -> list[Any]: if hasattr(result, "frames"): return list(result.frames or []) if isinstance(result, dict): @@ -533,7 +546,7 @@ async def _send_json(websocket: WebSocket, message: Any) -> None: async def _send_error( websocket: WebSocket, - code: str, + code: _ErrorCode, message: str, *, retryable: bool, diff --git a/fastvideo/tests/entrypoints/streaming/test_health_routes.py b/fastvideo/tests/entrypoints/streaming/test_health_routes.py new file mode 100644 index 000000000..c385bd394 --- /dev/null +++ b/fastvideo/tests/entrypoints/streaming/test_health_routes.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Health/readiness/status route contracts for streaming servers.""" +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +pytest.importorskip("starlette") +from fastapi import FastAPI # noqa: E402 +from starlette.testclient import TestClient # noqa: E402 + +from fastvideo.api.schema import ( # noqa: E402 + GeneratorConfig, + GenerationRequest, + SamplingConfig, + ServeConfig, + StreamingConfig, +) +from fastvideo.entrypoints.streaming.gpu_pool import ( # noqa: E402 + GpuPool, + PoolAssignment, + PoolHealth, +) +from fastvideo.entrypoints.streaming.health import ( # noqa: E402 + build_health_router, + get_pool_status, +) +from fastvideo.entrypoints.streaming.server import build_app # noqa: E402 + + +_STATUS_KEYS = { + "total_gpus", + "available_gpus", + "queue_size", + "warmup_enabled", + "warmup_successful_gpus", + "warmup_failed_gpus", + "gpu_status", +} + + +class _FakePool(GpuPool): + + def __init__(self, status: dict[str, Any] | None = None) -> None: + self._status: dict[str, Any] = status or _fake_status() + + def get_status(self) -> dict[str, Any]: + return self._status + + def health(self) -> PoolHealth: + return PoolHealth( + total_workers=self._status["total_gpus"], + available_workers=self._status["available_gpus"], + active_sessions=0, + queued_sessions=self._status["queue_size"], + ) + + async def acquire(self, session_id: str, *, timeout: float | None = None) -> PoolAssignment: + raise AssertionError("health tests should not acquire a worker") + + async def run(self, session_id: str, request: GenerationRequest) -> Any: + raise AssertionError("health tests should not run generation") + + async def release(self, session_id: str) -> None: + pass + + async def shutdown(self) -> None: + pass + + +def _fake_status() -> dict[str, Any]: + return { + "total_gpus": 2, + "available_gpus": 1, + "queue_size": 3, + "warmup_enabled": True, + "warmup_successful_gpus": 1, + "warmup_failed_gpus": 1, + "gpu_status": { + 0: { + "ready": True, + "available": True, + "client_count": 0, + "current_model_id": "ltx2", + "process_alive": True, + "warmup_enabled": True, + "warmup_success": True, + "warmup_error": None, + "warmup_timings": {"load": 1.2}, + }, + 1: { + "ready": False, + "available": False, + "client_count": 0, + "current_model_id": None, + "process_alive": False, + "warmup_enabled": True, + "warmup_success": False, + "warmup_error": "boot failed", + "warmup_timings": {}, + }, + }, + } + + +def _build_health_client(pool=None) -> TestClient: + app = FastAPI() + app.include_router(build_health_router(pool)) + return TestClient(app) + + +def _build_serve_config() -> ServeConfig: + return ServeConfig( + generator=GeneratorConfig(model_path="/models/fake"), + default_request=GenerationRequest( + sampling=SamplingConfig( + num_frames=12, + height=64, + width=64, + fps=12, + num_inference_steps=1, + ), + ), + streaming=StreamingConfig( + session_timeout_seconds=60, + generation_segment_cap=2, + ), + ) + + +class TestHealthRouterWithoutPool: + + def test_healthz_stays_live_without_pool(self): + client = _build_health_client(pool=None) + response = client.get("/healthz") + assert response.status_code == 200 + body = response.json() + assert body["status"] == "ok" + assert body["service"] == "ltx2-streaming-backend" + assert isinstance(body["ts"], str) + + def test_readyz_reports_warming_zeroed_pool(self): + client = _build_health_client(pool=None) + response = client.get("/readyz") + assert response.status_code == 200 + body = response.json() + assert body == { + "status": "warming", + "service": "ltx2-streaming-backend", + "ready_gpu_workers": 0, + "total_gpus": 0, + "available_gpus": 0, + "warmup_successful_gpus": 0, + "warmup_failed_gpus": 0, + "queue_size": 0, + "ts": body["ts"], + } + assert isinstance(body["ts"], str) + + def test_status_returns_zeroed_pool_shape(self): + client = _build_health_client(pool=None) + response = client.get("/status") + assert response.status_code == 200 + body = response.json() + assert set(body) == _STATUS_KEYS + assert body["total_gpus"] == 0 + assert body["available_gpus"] == 0 + assert body["queue_size"] == 0 + assert body["gpu_status"] == {} + + +class TestHealthRouterWithPool: + + def test_all_routes_return_expected_shapes(self): + client = _build_health_client(pool=_FakePool()) + + health = client.get("/health").json() + assert health["status"] == "ok" + assert health["service"] == "ltx2-streaming-backend" + + ready = client.get("/readyz").json() + assert ready["status"] == "ready" + assert ready["ready_gpu_workers"] == 1 + assert ready["total_gpus"] == 2 + assert ready["available_gpus"] == 1 + assert ready["warmup_successful_gpus"] == 1 + assert ready["warmup_failed_gpus"] == 1 + assert ready["queue_size"] == 3 + + status = client.get("/status").json() + assert set(status) == _STATUS_KEYS + assert status["total_gpus"] == 2 + assert status["available_gpus"] == 1 + assert status["queue_size"] == 3 + assert "0" in status["gpu_status"] + assert status["gpu_status"]["0"]["ready"] is True + + def test_pool_can_be_resolved_lazily(self): + pool = _FakePool() + client = _build_health_client(pool=lambda: pool) + response = client.get("/status") + assert response.status_code == 200 + assert response.json()["total_gpus"] == 2 + + def test_helper_returns_same_pool_status_shape(self): + status = asyncio.run(get_pool_status(_FakePool())) + assert set(status) == _STATUS_KEYS + assert status["gpu_status"][0]["current_model_id"] == "ltx2" + + +class TestBuildAppHealthMount: + + def test_build_app_keeps_legacy_health_and_mounts_public_routes(self): + app = build_app(_build_serve_config(), pool=_FakePool()) + client = TestClient(app) + + legacy_health = client.get("/health") + assert legacy_health.status_code == 200 + assert legacy_health.json() == { + "status": "ok", + "sessions": 0, + "stream_mode": "av_fmp4", + } + + healthz = client.get("/healthz") + assert healthz.status_code == 200 + assert healthz.json()["service"] == "ltx2-streaming-backend" + + readyz = client.get("/readyz") + assert readyz.status_code == 200 + assert readyz.json()["status"] == "ready" + + status = client.get("/status") + assert status.status_code == 200 + assert status.json()["total_gpus"] == 2