nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat] streaming: promote /healthz /readyz /status to public build_app

Browse Source 
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 <junda@berkeley.edu>
Co-authored-by: Matthew Noto <matthewnoto@gmail.com>
Co-authored-by: XOR-op <xor-op@users.noreply.github.com>
Co-authored-by: Zhang Peiyuan <peiyuan_zhang@berkeley.edu>
This commit is contained in:
SolitaryThinker
2026-05-07 13:37:25 -07:00
co-authored by Junda Su Matthew Noto XOR-op Zhang Peiyuan
parent 02ab3e45cd
commit cb082d8bbb
4 changed files with 447 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/entrypoints/streaming/__init__.py
+6
View File
@@ -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",
fastvideo/entrypoints/streaming/health.py
+187
View File
@@ -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",
]
fastvideo/entrypoints/streaming/server.py
+17 -4
View File
@@ -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,
fastvideo/tests/entrypoints/streaming/test_health_routes.py
+237
View File
@@ -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