131 lines
4.9 KiB
Python
131 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Shared asynchronous execution substrate for OpenAI-compatible routes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any, TypeVar
|
|
|
|
from fastvideo.api.schema import GenerationRequest
|
|
from fastvideo.entrypoints.video_generator import VideoGenerator
|
|
|
|
_T = TypeVar("_T")
|
|
|
|
|
|
class OpenAIServingEngine:
|
|
"""Own generator lifecycle and serialize access to its mutable pipeline.
|
|
|
|
FastVideo pipelines contain request-mutated sampling state and some LoRA
|
|
implementations merge weights in place. Running two Python threads through
|
|
one pipeline is therefore unsafe even if the HTTP layer accepts requests
|
|
concurrently. This engine gives every OpenAI route one model-agnostic async
|
|
entrypoint while preserving that invariant. A future scheduler can replace
|
|
the lock without changing the transport contract.
|
|
"""
|
|
|
|
def __init__(self, generator: VideoGenerator) -> None:
|
|
self._generator = generator
|
|
self._generation_lock = asyncio.Lock()
|
|
self._closed = False
|
|
self._unhealthy_reason: str | None = None
|
|
|
|
@property
|
|
def generator(self) -> VideoGenerator:
|
|
return self._generator
|
|
|
|
@property
|
|
def closed(self) -> bool:
|
|
return self._closed
|
|
|
|
@property
|
|
def healthy(self) -> bool:
|
|
if self._closed or self._unhealthy_reason is not None:
|
|
return False
|
|
executor = getattr(self._generator, "executor", None)
|
|
workers = getattr(executor, "workers", None)
|
|
if workers is None:
|
|
return True
|
|
return bool(workers) and all(worker.proc.is_alive() for worker in workers)
|
|
|
|
@property
|
|
def unhealthy_reason(self) -> str | None:
|
|
if self._unhealthy_reason is not None:
|
|
return self._unhealthy_reason
|
|
if not self.healthy:
|
|
return "one or more generation workers are not alive"
|
|
return None
|
|
|
|
async def generate(
|
|
self,
|
|
request: GenerationRequest,
|
|
*,
|
|
on_start: Callable[[], Awaitable[None]] | None = None,
|
|
) -> Any:
|
|
"""Generate one typed request without blocking the event loop."""
|
|
return await self.run_serialized(self._generator.generate, request, on_start=on_start)
|
|
|
|
async def run_serialized(
|
|
self,
|
|
function: Callable[..., _T],
|
|
*args: Any,
|
|
on_start: Callable[[], Awaitable[None]] | None = None,
|
|
**kwargs: Any,
|
|
) -> _T:
|
|
"""Run a synchronous pipeline operation under the serving lock."""
|
|
if self._closed:
|
|
raise RuntimeError("FastVideo serving engine is shutting down")
|
|
async with self._generation_lock:
|
|
if self._closed:
|
|
raise RuntimeError("FastVideo serving engine is shutting down")
|
|
if on_start is not None:
|
|
await on_start()
|
|
worker = asyncio.create_task(asyncio.to_thread(function, *args, **kwargs))
|
|
try:
|
|
return await asyncio.shield(worker)
|
|
except asyncio.CancelledError:
|
|
# Python cannot stop a running worker thread. Keep the lock
|
|
# until the pipeline call really exits so cancellation cannot
|
|
# expose mutable model state to a second request.
|
|
await self._wait_after_cancellation(worker)
|
|
raise
|
|
except (BrokenPipeError, EOFError) as error:
|
|
self._unhealthy_reason = str(error)
|
|
raise
|
|
|
|
async def run_async_serialized(self, function: Callable[[], Awaitable[_T]]) -> _T:
|
|
"""Run an async operation under the same pipeline lock."""
|
|
if self._closed:
|
|
raise RuntimeError("FastVideo serving engine is shutting down")
|
|
async with self._generation_lock:
|
|
if self._closed:
|
|
raise RuntimeError("FastVideo serving engine is shutting down")
|
|
worker: asyncio.Future[_T] = asyncio.ensure_future(function())
|
|
try:
|
|
return await asyncio.shield(worker)
|
|
except asyncio.CancelledError:
|
|
await self._wait_after_cancellation(worker)
|
|
raise
|
|
|
|
@staticmethod
|
|
async def _wait_after_cancellation(worker: asyncio.Future[Any]) -> None:
|
|
"""Keep waiting for an uninterruptible worker despite repeated cancellation."""
|
|
while not worker.done():
|
|
try:
|
|
await asyncio.shield(worker)
|
|
except asyncio.CancelledError:
|
|
continue
|
|
except Exception:
|
|
break
|
|
if not worker.cancelled():
|
|
worker.exception()
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Stop accepting requests and release the generator after in-flight work."""
|
|
self._closed = True
|
|
async with self._generation_lock:
|
|
await asyncio.to_thread(self._generator.shutdown)
|
|
|
|
|
|
__all__ = ["OpenAIServingEngine"]
|