From d9736d5e5d93e4db27f9d9491dcbabb4f178c0b0 Mon Sep 17 00:00:00 2001 From: Mac Lee Date: Sat, 30 May 2026 01:35:00 +0000 Subject: [PATCH] [feat]: add OpenAI video batching scheduler --- .../multimodal-gen-batching-port.md | 64 ++++++- fastvideo/entrypoints/openai/api_server.py | 21 ++- fastvideo/entrypoints/openai/batching.py | 170 ++++++++++++++++++ fastvideo/entrypoints/openai/state.py | 14 +- fastvideo/entrypoints/openai/video_api.py | 13 +- .../tests/entrypoints/test_openai_api.py | 110 ++++++++++++ 6 files changed, 377 insertions(+), 15 deletions(-) create mode 100644 fastvideo/entrypoints/openai/batching.py diff --git a/.agents/exploration/multimodal-gen-batching-port.md b/.agents/exploration/multimodal-gen-batching-port.md index d7998b045..275c29223 100644 --- a/.agents/exploration/multimodal-gen-batching-port.md +++ b/.agents/exploration/multimodal-gen-batching-port.md @@ -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. diff --git a/fastvideo/entrypoints/openai/api_server.py b/fastvideo/entrypoints/openai/api_server.py index 2b64087af..27faf67d4 100644 --- a/fastvideo/entrypoints/openai/api_server.py +++ b/fastvideo/entrypoints/openai/api_server.py @@ -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.") diff --git a/fastvideo/entrypoints/openai/batching.py b/fastvideo/entrypoints/openai/batching.py new file mode 100644 index 000000000..68597b865 --- /dev/null +++ b/fastvideo/entrypoints/openai/batching.py @@ -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 diff --git a/fastvideo/entrypoints/openai/state.py b/fastvideo/entrypoints/openai/state.py index c95f84400..5630dcc37 100644 --- a/fastvideo/entrypoints/openai/state.py +++ b/fastvideo/entrypoints/openai/state.py @@ -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 diff --git a/fastvideo/entrypoints/openai/video_api.py b/fastvideo/entrypoints/openai/video_api.py index 4134175f3..b02a748d2 100644 --- a/fastvideo/entrypoints/openai/video_api.py +++ b/fastvideo/entrypoints/openai/video_api.py @@ -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] = { diff --git a/fastvideo/tests/entrypoints/test_openai_api.py b/fastvideo/tests/entrypoints/test_openai_api.py index e4f7179a6..edb50357d 100644 --- a/fastvideo/tests/entrypoints/test_openai_api.py +++ b/fastvideo/tests/entrypoints/test_openai_api.py @@ -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 # ---------------------------------------------------------------------------