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

[feat]: add OpenAI video batching scheduler

Browse Source
This commit is contained in:
Mac Lee
2026-07-05 13:39:16 -07:00
committed by SolitaryThinker
parent 7396024f95
commit d9736d5e5d
6 changed files with 377 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/exploration/multimodal-gen-batching-port.md
+56 -8
View File
@@ -56,6 +56,18 @@ FastVideo files inspected:
- [x] Stage 2 remote validation passed on Modal L40S.
- [x] Stage 2 commit: `fe5572f3`
(`[feat]: add generator dynamic batching path`).
- [x] Stage 2 state commit: `39cb1d43`
(`[misc]: record batching stage 2 state`).
- [x] Stage 5 local implementation:
- added `VideoBatchScheduler` for OpenAI video requests,
- stores the scheduler in OpenAI server state,
- starts/stops the scheduler in API server lifespan,
- routes `video_api._run_generation()` through the scheduler when
dynamic batching is enabled,
- added async scheduler tests for compatible grouping and incompatible
fallback.
- [x] Stage 5 remote validation passed on Modal L40S.
- [ ] Stage 5 commit.
- [ ] After approval, implement remaining stages with commits and remote Modal validation.
- [ ] Produce final Markdown write-up with test and benchmark results and commit it.
@@ -193,6 +205,46 @@ Validation clean rerun:
- Pre-commit passed: yapf, ruff, codespell, mypy, filename spaces, and
suggestion hooks.
### Stage 5: OpenAI Server Queue Integration
Files changed/added:
- `fastvideo/entrypoints/openai/batching.py`
- Added `VideoBatchScheduler`, an async FIFO queue with batching delay,
compatibility checks, background dispatch, and per-request futures.
- `fastvideo/entrypoints/openai/state.py`
- Added global scheduler storage and accessor.
- `fastvideo/entrypoints/openai/api_server.py`
- Starts the scheduler during lifespan when `batching_mode=dynamic` and
`batching_max_size>1`; stops it before generator shutdown.
- `fastvideo/entrypoints/openai/video_api.py`
- `_run_generation()` submits to the scheduler when enabled; otherwise keeps
the prior direct executor-thread path.
- `fastvideo/tests/entrypoints/test_openai_api.py`
- Added scheduler grouping and incompatible fallback tests using a fake
generator.
Pending for Stage 5:
- Commit Stage 5 if validation passes.
Validation attempt:
- Modal app: `ap-2PJ4b5eKPnxR9HTEb9x3UM`
- Command:
`pytest fastvideo/tests/entrypoints/test_openai_api.py -q && pre-commit run --files ...`
- Result:
- Tests passed: `61 passed, 14 warnings`.
- Pre-commit failed only on mypy for assigning to an `exc` variable outside
an `except` block in `openai/batching.py`.
- Fix applied locally:
- Renamed that local variable to `error`.
Validation clean rerun:
- Modal app: `ap-sgV5gRsGJeHE4g9a1Dswk3`
- Command:
`pytest fastvideo/tests/entrypoints/test_openai_api.py -q && pre-commit run --files ...`
- Result:
- Tests passed: `61 passed, 14 warnings`.
- Pre-commit passed: yapf, ruff, codespell, mypy, filename spaces, and
suggestion hooks.
Validation attempt:
- Modal L40S command:
`pytest fastvideo/tests/batching fastvideo/tests/api/test_compat_translation.py -q && pre-commit run --files ...`
@@ -242,8 +294,9 @@ Current branch: `multimodal-gen-batching`.
Current local implementation state:
- Stage 1 code is committed as `3d3157fb`; state commit is `ce758820`.
- Stage 2 code is committed as `fe5572f3`.
- The state-file update recording that commit is the only local change before
Stage 5 begins.
- Stage 2 state commit is `39cb1d43`; branch is pushed to origin through
Stage 2.
- Stage 5 code is local, remote validated, and ready to commit.
Important constraints:
- Do not edit unrelated untracked files already present in the worktree.
@@ -253,9 +306,4 @@ Important constraints:
test environment.
Next step:
Begin Stage 5 OpenAI server queue integration:
- add a queue/scheduler object that groups compatible queued API requests,
- start/stop it from API server lifespan,
- route `video_api.create_video` through the scheduler when dynamic batching is
enabled,
- keep image/reference requests sequential via compatibility fallback.
Commit the Stage 5 slice, then begin GPU parity and benchmark validation.
fastvideo/entrypoints/openai/api_server.py
+20 -1
View File
@@ -10,6 +10,7 @@ from fastapi.middleware.cors import CORSMiddleware
from fastvideo.api.presets import validate_preset_selection
from fastvideo.api.schema import GenerationRequest
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
from fastvideo.entrypoints.openai.state import (
DEFAULT_OUTPUT_DIR,
clear_state,
@@ -59,11 +60,29 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
generator = VideoGenerator.from_fastvideo_args(args)
logger.info("Model loaded successfully.")
set_state(generator, args, output_dir, default_request=default_request)
video_batch_scheduler: VideoBatchScheduler | None = None
if args.batching_mode == "dynamic" and args.batching_max_size > 1:
video_batch_scheduler = VideoBatchScheduler(generator, args)
await video_batch_scheduler.start()
logger.info(
"Started dynamic video batch scheduler: max_size=%d delay_ms=%.2f",
args.batching_max_size,
args.batching_delay_ms,
)
set_state(
generator,
args,
output_dir,
default_request=default_request,
video_batch_scheduler=video_batch_scheduler,
)
yield # server is running
logger.info("Shutting down — releasing model resources ...")
if video_batch_scheduler is not None:
await video_batch_scheduler.stop()
generator.shutdown()
clear_state()
logger.info("Shutdown complete.")
fastvideo/entrypoints/openai/batching.py
+170
View File
@@ -0,0 +1,170 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import asyncio
import time
from collections import deque
from dataclasses import dataclass
from typing import Any
from fastvideo.api.sampling_param import SamplingParam
from fastvideo.batching.signature import can_dynamic_batch
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.logger import init_logger
logger = init_logger(__name__)
@dataclass
class _VideoBatchJob:
request_id: str
kwargs: dict[str, Any]
future: asyncio.Future
enqueue_time: float
class VideoBatchScheduler:
"""Async FIFO scheduler for OpenAI-compatible video generation."""
def __init__(self, generator: Any, fastvideo_args: FastVideoArgs) -> None:
self._generator = generator
self._fastvideo_args = fastvideo_args
self._queue: asyncio.Queue[_VideoBatchJob | None] = asyncio.Queue()
self._pending: deque[_VideoBatchJob] = deque()
self._task: asyncio.Task | None = None
self._stopped = False
@property
def enabled(self) -> bool:
return self._fastvideo_args.batching_mode == "dynamic" and self._fastvideo_args.batching_max_size > 1
async def start(self) -> None:
if self._task is not None:
return
self._task = asyncio.create_task(self._run(), name="fastvideo-video-batch-scheduler")
async def stop(self) -> None:
self._stopped = True
await self._queue.put(None)
if self._task is not None:
await self._task
self._task = None
async def submit(self, request_id: str, kwargs: dict[str, Any]) -> Any:
loop = asyncio.get_running_loop()
future = loop.create_future()
await self._queue.put(
_VideoBatchJob(
request_id=request_id,
kwargs=dict(kwargs),
future=future,
enqueue_time=time.perf_counter(),
))
return await future
async def _run(self) -> None:
while not self._stopped:
job = await self._get_next_job()
if job is None:
break
batch = await self._collect_batch(job)
await self._dispatch(batch)
while self._pending:
pending = self._pending.popleft()
if not pending.future.done():
pending.future.set_exception(RuntimeError("Video batch scheduler stopped before dispatch"))
async def _get_next_job(self) -> _VideoBatchJob | None:
if self._pending:
return self._pending.popleft()
return await self._queue.get()
async def _collect_batch(self, first: _VideoBatchJob) -> list[_VideoBatchJob]:
batch = [first]
max_size = self._fastvideo_args.batching_max_size
delay_s = max(0.0, self._fastvideo_args.batching_delay_ms / 1000.0)
deadline = first.enqueue_time + delay_s
while len(batch) < max_size:
timeout = deadline - time.perf_counter()
if timeout <= 0 and delay_s > 0:
break
try:
candidate = await asyncio.wait_for(self._get_next_job(), timeout=max(0.0, timeout))
except TimeoutError:
break
if candidate is None:
await self._queue.put(None)
break
if self._jobs_are_compatible(batch[0], candidate):
batch.append(candidate)
continue
self._pending.append(candidate)
break
return batch
async def _dispatch(self, batch: list[_VideoBatchJob]) -> None:
loop = asyncio.get_running_loop()
request_ids = [job.request_id for job in batch]
queue_wait_ms = (time.perf_counter() - min(job.enqueue_time for job in batch)) * 1000.0
if self._fastvideo_args.enable_batching_metrics:
logger.info(
"Dispatching video batch: request_ids=%s size=%d queue_wait_ms=%.2f",
request_ids,
len(batch),
queue_wait_ms,
)
try:
results = await loop.run_in_executor(
None,
lambda: self._generator.generate_video_batch([job.kwargs for job in batch]),
)
except Exception as exc:
for job in batch:
if not job.future.done():
job.future.set_exception(exc)
return
if len(results) != len(batch):
error = RuntimeError(f"Video batch returned {len(results)} results for {len(batch)} requests")
for job in batch:
if not job.future.done():
job.future.set_exception(error)
return
for job, result in zip(batch, results, strict=True):
if not job.future.done():
job.future.set_result(result)
def _jobs_are_compatible(self, base: _VideoBatchJob, candidate: _VideoBatchJob) -> bool:
try:
base_sampling, base_extra = self._sampling_param_from_kwargs(base.kwargs)
candidate_sampling, candidate_extra = self._sampling_param_from_kwargs(candidate.kwargs)
except Exception:
return False
return can_dynamic_batch(
base_sampling,
candidate_sampling,
base_extra=base_extra,
candidate_extra=candidate_extra,
).can_batch
def _sampling_param_from_kwargs(self, kwargs: dict[str, Any]) -> tuple[SamplingParam, dict[str, Any]]:
sampling_param = SamplingParam.from_pretrained(self._fastvideo_args.model_path)
updates = dict(kwargs)
prompt = updates.pop("prompt", None)
extra: dict[str, Any] = {}
for key in (
"ltx2_audio_latents",
"ltx2_audio_clean_latent",
"ltx2_audio_denoise_mask",
"audio_num_frames",
"video_position_offset_sec",
):
if key in updates:
extra[key] = updates.pop(key)
sampling_param.update(updates)
sampling_param.prompt = prompt
return sampling_param, extra
fastvideo/entrypoints/openai/state.py
+12 -2
View File
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING
if TYPE_CHECKING:
from fastvideo.api.schema import GenerationRequest
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
from fastvideo.entrypoints.video_generator import VideoGenerator
from fastvideo.fastvideo_args import FastVideoArgs
@@ -20,6 +21,7 @@ _generator: VideoGenerator | None = None
_fastvideo_args: FastVideoArgs | None = None
_output_dir: str = DEFAULT_OUTPUT_DIR
_default_request: GenerationRequest | None = None
_video_batch_scheduler: VideoBatchScheduler | None = None
def get_generator() -> VideoGenerator:
@@ -44,23 +46,31 @@ def get_default_request() -> GenerationRequest | None:
return _default_request
def get_video_batch_scheduler() -> VideoBatchScheduler | None:
"""Return the video batch scheduler when dynamic batching is enabled."""
return _video_batch_scheduler
def set_state(
generator: VideoGenerator,
fastvideo_args: FastVideoArgs,
output_dir: str,
default_request: GenerationRequest | None = None,
video_batch_scheduler: VideoBatchScheduler | None = None,
) -> None:
"""Set all server state at once (called from lifespan)."""
global _generator, _fastvideo_args, _output_dir, _default_request
global _generator, _fastvideo_args, _output_dir, _default_request, _video_batch_scheduler
_generator = generator
_fastvideo_args = fastvideo_args
_output_dir = output_dir
_default_request = default_request
_video_batch_scheduler = video_batch_scheduler
def clear_state() -> None:
"""Clear server state on shutdown."""
global _generator, _fastvideo_args, _default_request
global _generator, _fastvideo_args, _default_request, _video_batch_scheduler
_generator = None
_fastvideo_args = None
_default_request = None
_video_batch_scheduler = None
fastvideo/entrypoints/openai/video_api.py
+9 -4
View File
@@ -26,6 +26,7 @@ from fastvideo.entrypoints.openai.state import (
get_generator,
get_output_dir,
get_server_args,
get_video_batch_scheduler,
)
from fastvideo.entrypoints.openai.protocol import (
VideoGenerationsRequest,
@@ -151,15 +152,19 @@ async def _run_generation(request_id: str, kwargs: dict[str, Any]) -> None:
is synchronous) and update the store on completion or failure.
"""
generator = get_generator()
scheduler = get_video_batch_scheduler()
loop = asyncio.get_running_loop()
try:
start = time.perf_counter()
result = await loop.run_in_executor(
None,
lambda: generator.generate_video(**kwargs),
)
if scheduler is not None and scheduler.enabled:
result = await scheduler.submit(request_id, kwargs)
else:
result = await loop.run_in_executor(
None,
lambda: generator.generate_video(**kwargs),
)
elapsed = time.perf_counter() - start
update: dict[str, Any] = {
fastvideo/tests/entrypoints/test_openai_api.py
+110
View File
@@ -1,10 +1,14 @@
"""Unit tests for the OpenAI-compatible API server helpers (no GPU needed)."""
import asyncio
import os
from types import SimpleNamespace
from unittest.mock import patch
import pytest
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.entrypoints.openai.batching import VideoBatchScheduler
from fastvideo.api.parser import parse_config
from fastvideo.api.schema import GenerationRequest
from fastvideo.entrypoints.openai.protocol import (
@@ -21,6 +25,112 @@ from fastvideo.entrypoints.openai.utils import (
parse_size,
)
class _FakeBatchGenerator:
def __init__(self):
self.calls = []
def generate_video_batch(self, request_kwargs):
self.calls.append([dict(item) for item in request_kwargs])
return [{"prompts": item["prompt"], "video_path": item["output_path"]} for item in request_kwargs]
def _batch_scheduler_args(**overrides):
defaults = dict(
model_path="test-model",
batching_mode="dynamic",
batching_max_size=2,
batching_delay_ms=25.0,
enable_batching_metrics=False,
pipeline_config=PipelineConfig(),
)
defaults.update(overrides)
return SimpleNamespace(**defaults)
def test_video_batch_scheduler_groups_compatible_requests(tmp_path):
async def run():
generator = _FakeBatchGenerator()
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
await scheduler.start()
try:
first = {
"prompt": "first",
"height": 256,
"width": 256,
"num_frames": 1,
"num_inference_steps": 2,
"seed": 1,
"output_path": str(tmp_path / "first.mp4"),
"save_video": False,
}
second = {
"prompt": "second",
"height": 256,
"width": 256,
"num_frames": 1,
"num_inference_steps": 2,
"seed": 2,
"output_path": str(tmp_path / "second.mp4"),
"save_video": False,
}
results = await asyncio.gather(
scheduler.submit("req-1", first),
scheduler.submit("req-2", second),
)
finally:
await scheduler.stop()
return generator.calls, results
calls, results = asyncio.run(run())
assert len(calls) == 1
assert [item["prompt"] for item in calls[0]] == ["first", "second"]
assert [result["prompts"] for result in results] == ["first", "second"]
def test_video_batch_scheduler_keeps_incompatible_requests_separate(tmp_path):
async def run():
generator = _FakeBatchGenerator()
scheduler = VideoBatchScheduler(generator, _batch_scheduler_args())
await scheduler.start()
try:
text_only = {
"prompt": "first",
"height": 256,
"width": 256,
"num_frames": 1,
"num_inference_steps": 2,
"seed": 1,
"output_path": str(tmp_path / "first.mp4"),
"save_video": False,
}
image_conditioned = {
"prompt": "second",
"height": 256,
"width": 256,
"num_frames": 1,
"num_inference_steps": 2,
"seed": 2,
"image_path": str(tmp_path / "input.png"),
"output_path": str(tmp_path / "second.mp4"),
"save_video": False,
}
results = await asyncio.gather(
scheduler.submit("req-1", text_only),
scheduler.submit("req-2", image_conditioned),
)
finally:
await scheduler.stop()
return generator.calls, results
calls, results = asyncio.run(run())
assert len(calls) == 2
assert [[item["prompt"] for item in call] for call in calls] == [["first"], ["second"]]
assert [result["prompts"] for result in results] == ["first", "second"]
# ---------------------------------------------------------------------------
# parse_size
# ---------------------------------------------------------------------------