diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 3a37b52d2..331487fac 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -82,6 +82,7 @@ depends on the installed kernels and model path. | --- | --- | --- | --- | --- | | `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable | | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | +| `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental | | `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off) | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental | | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index c6e91df60..1152ebce3 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -109,6 +109,16 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4 The final content is available from `/v1/videos/{video_id}/content` after the job status becomes `completed`. +`queued` means the request is still waiting for diffusion scheduler admission. +The status changes to `in_progress` when the scheduler first selects the +request for execution. + +`DELETE /v1/videos/{video_id}` issues a bounded engine abort +(`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend +task. Cancellation cleanup is also bounded and best-effort: it confirms +the abort was queued, and the current request batch may still drain. +The job is then re-read so a completed save is not orphaned. + ### Synchronous Response `POST /v1/videos/sync` blocks until generation finishes and returns raw video diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index ce608dca3..b73062689 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 27, + EnvironmentVariableCategory.PUBLIC_OMNI: 28, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 64, diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 1ee7424e9..77a6da5db 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -15,7 +15,12 @@ from pytest_mock import MockerFixture import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + DIFFUSION_REQUEST_LIFECYCLE_KEY, + DIFFUSION_REQUEST_STARTED, + DiffusionOutput, + OmniDiffusionConfig, +) from vllm_omni.diffusion.diffusion_engine import ( DiffusionEngine, DiffusionExecutionMode, @@ -183,6 +188,28 @@ def _make_request_mode_sched_output(*request_ids: str) -> RealDiffusionScheduler ) +@pytest.mark.cpu +def test_request_started_output_is_emitted_only_for_opted_in_requests() -> None: + sched_output = _make_request_mode_sched_output("tracked", "untracked") + sched_output.scheduled_new_reqs[0].req.sampling_params.emit_request_lifecycle = True + engine = object.__new__(DiffusionEngine) + emitted = [] + engine._put_output = lambda request_id, output: emitted.append((request_id, output)) + + engine._emit_request_started_outputs(sched_output) + + assert len(emitted) == 1 + request_id, output = emitted[0] + assert request_id == "tracked" + assert output.request_started is True + assert output.finished is False + [formatted] = engine.postprocess_output(sched_output.scheduled_new_reqs[0].req, output) + assert formatted.custom_output == { + DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED, + } + assert formatted.finished is False + + class TestRequestBatchCapability: pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index 4578d8368..a1dfec637 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -443,3 +443,14 @@ def test_format_empty_diffusion_outputs_preserves_empty_response_shape() -> None assert [result.prompt for result in results] == ["prompt-0"] assert [result.images for result in results] == [[]] assert [result.metrics for result in results] == [{}] + + +def test_format_empty_diffusion_outputs_preserves_custom_output() -> None: + [result] = format_empty_diffusion_outputs( + _request("prompt-0"), + finished=False, + custom_output={"lifecycle": "started"}, + ) + + assert result.custom_output == {"lifecycle": "started"} + assert result.finished is False diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index e39c18ff4..3eadc53c1 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import asyncio @@ -7,7 +10,11 @@ from unittest.mock import MagicMock, patch import pytest from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + DIFFUSION_REQUEST_LIFECYCLE_KEY, + DIFFUSION_REQUEST_STARTED, + OmniDiffusionConfig, +) from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -80,6 +87,43 @@ async def test_inline_dispatch_request_success(client, mock_engine): assert output.request_id == "req-1" +@pytest.mark.asyncio +async def test_inline_non_streaming_dispatches_lifecycle_before_final(client, mock_engine): + lifecycle = OmniRequestOutput.from_diffusion( + request_id="req-lifecycle", + images=[], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + finished=False, + ) + intermediate = OmniRequestOutput.from_diffusion( + request_id="req-lifecycle", + images=[], + custom_output={"chunk": 0}, + finished=False, + ) + final = OmniRequestOutput.from_diffusion(request_id="req-lifecycle", images=[MagicMock()]) + + async def _step_streaming(_request): + yield [lifecycle] + yield [intermediate] + yield [final] + + mock_engine.step_streaming = _step_streaming + + await client.add_request_async("req-lifecycle", "A test prompt", OmniDiffusionSamplingParams()) + + outputs = [] + for _ in range(20): + output = client.get_diffusion_output_nowait() + if output is not None: + outputs.append(output) + if output.finished: + break + await asyncio.sleep(0.01) + + assert outputs == [lifecycle, final] + + @pytest.mark.asyncio async def test_inline_requests_clone_shared_sampling_params(client, mock_engine): requests = [] diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index 90e753c9b..c6b5761d9 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -9,6 +9,7 @@ import pytest import vllm_omni.diffusion.stage_diffusion_proc as stage_diffusion_proc import vllm_omni.plugins as omni_plugins +from vllm_omni.diffusion.data import DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -103,6 +104,49 @@ async def test_proc_streaming_request_yields_each_engine_chunk(): assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}} +@pytest.mark.asyncio +async def test_proc_non_streaming_forwards_lifecycle_before_final_output(): + lifecycle = OmniRequestOutput.from_diffusion( + request_id="", + images=[], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + finished=False, + ) + intermediate = OmniRequestOutput.from_diffusion( + request_id="", + images=[], + custom_output={"chunk": 0}, + finished=False, + ) + final = OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True) + + class _LifecycleEngine: + async def step_streaming(self, request): + del request + yield [lifecycle] + yield [intermediate] + yield [final] + + stage_proc = object.__new__(StageDiffusionProc) + stage_proc._engine = _LifecycleEngine() + intermediate_outputs = [] + + async def _capture(output): + intermediate_outputs.append(output) + + result = await stage_proc._process_request( + request_id="req-lifecycle", + prompt="prompt", + sampling_params_dict=asdict(OmniDiffusionSamplingParams()), + on_request_started=_capture, + ) + + assert intermediate_outputs == [lifecycle] + assert lifecycle.request_id == "req-lifecycle" + assert result is final + assert result.request_id == "req-lifecycle" + + @pytest.mark.asyncio async def test_proc_process_request_with_batching_async_output(): stage_proc = object.__new__(StageDiffusionProc) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 57b2daaf8..2de5c4056 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -25,6 +25,7 @@ from PIL import Image from pytest_mock import MockerFixture from vllm import envs +from vllm_omni.diffusion.data import DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes from vllm_omni.entrypoints.openai import api_server, video_api_utils from vllm_omni.entrypoints.openai.api_server import router @@ -58,6 +59,7 @@ class MockVideoResult: multimodal_output=None, stage_durations=None, peak_memory_mb=0.0, + custom_output=None, ): self.multimodal_output = dict(multimodal_output or {"video": videos}) if audios is not None: @@ -66,6 +68,7 @@ class MockVideoResult: self.multimodal_output["audio_sample_rate"] = sample_rate self.stage_durations = stage_durations or {} self.peak_memory_mb = peak_memory_mb + self.custom_output = custom_output or {} class FakeAsyncOmni: @@ -102,9 +105,17 @@ class FakeAsyncOmni: ): self.captured_reference_video_bytes = [Path(item).read_bytes() for item in reference_videos] num_outputs = sampling_params_list[0].num_outputs_per_prompt + if sampling_params_list[0].emit_request_lifecycle: + yield MockVideoResult( + [], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + ) videos = [object() for _ in range(num_outputs)] yield MockVideoResult(videos) + async def abort(self, request_id, *, timeout=None): + del request_id, timeout + def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture): engine = FakeAsyncOmni() @@ -224,9 +235,18 @@ class BlockingVideoHandler: self.stage_configs = stage_configs async def generate_video_bytes( - self, request, reference_id, *, reference_image=None, reference_video=None, reference_audio=None + self, + request, + reference_id, + *, + reference_image=None, + reference_video=None, + reference_audio=None, + on_started=None, ): del request, reference_id, reference_image, reference_video, reference_audio + if on_started is not None: + await on_started() self.started.set() try: await asyncio.Future() @@ -234,6 +254,104 @@ class BlockingVideoHandler: self.cancelled.set() raise + async def abort_request(self, request_id): + del request_id + + +class HangingAbortHandler(BlockingVideoHandler): + """Engine abort never returns; DELETE must time out and still cancel.""" + + async def abort_request(self, request_id): + del request_id + await asyncio.sleep(30) + + +class CompletingDuringAbortHandler(BlockingVideoHandler): + """Finishes and persists output while DELETE is still awaiting abort.""" + + def __init__(self): + super().__init__() + self._finish = asyncio.Event() + + async def generate_video_bytes( + self, + request, + reference_id, + *, + reference_image=None, + reference_video=None, + reference_audio=None, + on_started=None, + ): + del request, reference_image, reference_video, reference_audio + if on_started is not None: + await on_started() + self.started.set() + await self._finish.wait() + return b"late-complete", {}, 0.0, None + + async def abort_request(self, request_id): + self._finish.set() + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline: + job = await api_server.VIDEO_STORE.get(request_id) + if job is not None and job.status is VideoGenerationStatus.COMPLETED and job.file_name is not None: + return + await asyncio.sleep(0.01) + raise RuntimeError(f"video job {request_id} did not complete during abort") + + +class SchedulerQueuedVideoHandler(BlockingVideoHandler): + """Parks after engine submission but before scheduler admission.""" + + def __init__(self): + super().__init__() + self.admit = threading.Event() + self.in_progress = threading.Event() + + async def generate_video_bytes( + self, + request, + reference_id, + *, + reference_image=None, + reference_video=None, + reference_audio=None, + on_started=None, + ): + del request, reference_id, reference_image, reference_video, reference_audio + self.started.set() + try: + while not self.admit.is_set(): + await asyncio.sleep(0.01) + if on_started is not None: + await on_started() + self.in_progress.set() + await asyncio.Future() + except asyncio.CancelledError: + self.cancelled.set() + raise + + +class AbortTrackingOmni(FakeAsyncOmni): + def __init__(self): + super().__init__() + self.aborted: list[str] = [] + self.entered = threading.Event() + + async def generate(self, prompt, request_id, sampling_params_list): + del prompt, request_id, sampling_params_list + self.entered.set() + yield MockVideoResult( + [], + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + ) + await asyncio.Future() + + async def abort(self, request_id, *, timeout=None): + assert timeout is not None + self.aborted.append(request_id) + class FakeServerSocket: def __init__(self): @@ -592,12 +710,12 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): engine = FakeAsyncOmni() - engine.get_diffusion_od_config = lambda: SimpleNamespace( + engine.get_diffusion_od_config = lambda: SimpleNamespace( # type: ignore[method-assign] model="org/model", model_class_name="ExamplePipeline", revision="pinned-revision", ) - captured = {} + captured: dict[str, str | None] = {} def fake_policy(model_class_name, *, model, revision=None): captured.update( @@ -625,6 +743,7 @@ def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch): ) ) + assert engine.captured_prompt is not None input_image = engine.captured_prompt["multi_modal_data"]["image"] assert isinstance(input_image, Image.Image) assert input_image.size == (48, 32) @@ -718,6 +837,7 @@ def test_video_generation_bridges_request_fields(generation_request, expected_nu asyncio.run(handler._run_and_extract(generation_request, "field-bridge")) + assert engine.captured_sampling_params_list is not None sampling = engine.captured_sampling_params_list[0] # Top-level ``seconds`` bridges into extra_args["duration"]; num_frames is # passed through (or derived as seconds x fps when omitted). No private @@ -1503,7 +1623,7 @@ def test_video_generation_response_exposes_action_payload(mocker: MockerFixture) }, ) - engine.generate = _generate + engine.generate = _generate # type: ignore[method-assign] mocker.patch( "vllm_omni.entrypoints.openai.serving_video.encode_video_base64", return_value="encoded-video", @@ -2096,6 +2216,80 @@ def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client): assert retrieve_resp.status_code == 404 +def test_async_video_stays_queued_until_scheduler_admission(test_client): + handler = SchedulerQueuedVideoHandler() + test_client.app.state.openai_serving_video = handler + + create_resp = test_client.post("/v1/videos", data={"prompt": "Queue this video"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert handler.started.wait(timeout=2.0) + + queued = test_client.get(f"/v1/videos/{video_id}") + assert queued.status_code == 200 + assert queued.json()["status"] == VideoGenerationStatus.QUEUED.value + + handler.admit.set() + assert handler.in_progress.wait(timeout=2.0) + in_progress = test_client.get(f"/v1/videos/{video_id}") + assert in_progress.status_code == 200 + assert in_progress.json()["status"] == VideoGenerationStatus.IN_PROGRESS.value + + assert test_client.delete(f"/v1/videos/{video_id}").status_code == 200 + + +def test_delete_times_out_engine_abort_and_still_cancels(test_client, monkeypatch): + monkeypatch.setattr(api_server, "VIDEO_ABORT_TIMEOUT_S", 0.05) + handler = HangingAbortHandler() + test_client.app.state.openai_serving_video = handler + + create_resp = test_client.post("/v1/videos", data={"prompt": "Hang abort"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert handler.started.wait(timeout=2.0) + + started = time.monotonic() + delete_resp = test_client.delete(f"/v1/videos/{video_id}") + assert time.monotonic() - started < 2.0 + assert delete_resp.status_code == 200 + assert handler.cancelled.wait(timeout=2.0) + + +def test_delete_removes_artifact_if_job_completes_during_abort(test_client): + handler = CompletingDuringAbortHandler() + test_client.app.state.openai_serving_video = handler + + create_resp = test_client.post("/v1/videos", data={"prompt": "Complete during delete"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert handler.started.wait(timeout=2.0) + + assert test_client.delete(f"/v1/videos/{video_id}").status_code == 200 + file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id) + assert not os.path.exists(file_path) + assert asyncio.run(api_server.VIDEO_STORE.get(video_id)) is None + + +def test_delete_aborts_engine_request_before_cancelling_task(test_client): + engine = AbortTrackingOmni() + test_client.app.state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion( + engine, + model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + ) + + create_resp = test_client.post("/v1/videos", data={"prompt": "Abort this video"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + assert engine.entered.wait(timeout=2.0) + _wait_for_status(test_client, video_id, VideoGenerationStatus.IN_PROGRESS.value) + + delete_resp = test_client.delete(f"/v1/videos/{video_id}") + assert delete_resp.status_code == 200 + assert delete_resp.json()["deleted"] is True + _wait_until(lambda: engine.aborted == [video_id]) + assert asyncio.run(api_server.VIDEO_STORE.get(video_id)) is None + + def test_video_response_file_extension_is_robust(): response = VideoResponse(model="test-model", prompt="Make something beautiful") assert response.file_extension == "mp4" diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index a69bcce09..2379db78e 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -12,6 +12,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.entrypoints import async_omni as async_omni_mod from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.outputs import OmniRequestOutput @@ -38,7 +39,8 @@ def get_fake_add_request(submitted_request_ids, submitted_lora_requests=None): def get_fake_abort(aborted_request_batches): - async def fake_abort_async(request_ids): + async def fake_abort_async(request_ids, timeout=None): + del timeout aborted_request_batches.append(list(request_ids)) return fake_abort_async @@ -196,8 +198,8 @@ def test_abort_keeps_request_states_until_generate_cleanup(): release = asyncio.Event() seen_during_wait: list[int] = [] - async def slow_abort_async(request_ids): - del request_ids + async def slow_abort_async(request_ids, timeout=None): + del request_ids, timeout seen_during_wait.append(len(omni.request_states)) await release.wait() @@ -224,8 +226,8 @@ def test_abort_keeps_request_states_until_generate_cleanup(): @pytest.mark.cpu def test_abort_propagates_engine_errors_without_popping_state(): async def run(): - async def failing_abort_async(request_ids): - del request_ids + async def failing_abort_async(request_ids, timeout=None): + del request_ids, timeout raise RuntimeError("orchestrator abort failed") omni = get_async_omni_instance(fake_abort_request=failing_abort_async) @@ -253,7 +255,8 @@ def test_abort_enqueues_prefix_tokens_from_engine(): queue: asyncio.Queue = asyncio.Queue() - async def abort_with_prefix(request_ids): + async def abort_with_prefix(request_ids, timeout=None): + del timeout rid = request_ids[0] engine_output = OmniRequestOutput( request_id=rid, @@ -308,8 +311,8 @@ def test_abort_enqueues_synthetic_finished_when_engine_returns_empty(): async def run(): queue: asyncio.Queue = asyncio.Queue() - async def empty_abort_async(request_ids): - del request_ids + async def empty_abort_async(request_ids, timeout=None): + del request_ids, timeout return [] omni = get_async_omni_instance(fake_abort_request=empty_abort_async) @@ -467,6 +470,42 @@ def test_generate_accepts_request_after_repeated_cancellations(): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_cancel_bounds_cleanup_abort(monkeypatch): + """Cancelled generate() must not wait forever on a wedged orchestrator abort.""" + monkeypatch.setattr(async_omni_mod, "ABORT_TIMEOUT_S", 0.05) + + async def run(): + seen_timeouts: list[float | None] = [] + + async def hanging_abort_async(request_ids, timeout=None): + del request_ids + seen_timeouts.append(timeout) + await asyncio.Event().wait() + + omni = get_async_omni_instance(fake_abort_request=hanging_abort_async) + + async def collect(): + async for _ in omni.generate( + prompt={"prompt": "prompt"}, + request_id="cancel-hang-abort", + sampling_params_list=[SimpleNamespace()], + output_modalities=["image"], + ): + pass + + task = asyncio.create_task(collect()) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1.0) + assert seen_timeouts + assert seen_timeouts[-1] == 0.05 + assert omni.request_states == {} + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker): async def run_test(): diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 95059a06e..01264eaea 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -167,7 +167,8 @@ class FakeAsyncOmniEngine: def abort(self, request_ids: list[str]) -> None: self.aborted.append(list(request_ids)) - async def abort_async(self, request_ids: list[str]) -> None: + async def abort_async(self, request_ids: list[str], timeout=None) -> None: + del timeout self.abort(request_ids) async def collective_rpc_async(self, **_: Any) -> list[Any]: diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index b5214720a..97782328d 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -63,6 +63,7 @@ _PUBLIC_OMNI = ( "OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS", "SPEAKER_MAX_UPLOADED", "SPEAKER_SAMPLES_DIR", + "VLLM_OMNI_ABORT_TIMEOUT", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", "VLLM_OMNI_EVENT_DRIVEN_ORCH", "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c333fb11a..10f8642c6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1637,6 +1637,18 @@ class OmniDiffusionConfig: return cls(**cls.normalize_init_kwargs(kwargs)) +DIFFUSION_REQUEST_LIFECYCLE_KEY = "_diffusion_request_lifecycle" +DIFFUSION_REQUEST_STARTED = "started" + + +def is_diffusion_request_started_output(output: Any) -> bool: + custom_output = getattr(output, "custom_output", None) + return ( + isinstance(custom_output, dict) + and custom_output.get(DIFFUSION_REQUEST_LIFECYCLE_KEY) == DIFFUSION_REQUEST_STARTED + ) + + @dataclass class DiffusionOutput: """ @@ -1685,6 +1697,9 @@ class DiffusionOutput: # mode) and the receiving side must not initialise a stray CUDA context. to_cpu: bool = False + # Internal control-plane event emitted on first scheduler admission. + request_started: bool = False + # Typed video-media contract. Declared last so the pre-existing positional # constructor order (output, trajectory_timesteps, ...) that out-of-tree # pipelines rely on is preserved. Mutually exclusive with ``output``. diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 417090503..6e53b5909 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -26,6 +26,8 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.kv_cache_interface import KVCacheConfig from vllm_omni.diffusion.data import ( + DIFFUSION_REQUEST_LIFECYCLE_KEY, + DIFFUSION_REQUEST_STARTED, DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig, @@ -55,7 +57,7 @@ from vllm_omni.diffusion.registry import ( ) from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest from vllm_omni.diffusion.sched import BaseScheduler, RequestScheduler, StepScheduler -from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus +from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus, DiffusionSchedulerOutput from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput, RunnerOutput from vllm_omni.errors import client_error_from_metadata, is_client_error_status from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt @@ -533,7 +535,12 @@ class DiffusionEngine: error_type=output.error_type, ) raise RuntimeError(output.error) - logger.debug("Generation completed successfully.") + if output.request_started: + return format_empty_diffusion_outputs( + request, + finished=False, + custom_output={DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED}, + ) if output.media is not None: if output.output is not None: @@ -599,6 +606,8 @@ class DiffusionEngine: sched_output = self.scheduler.schedule() self._scheduler_num_waiting_reqs = max(int(sched_output.num_waiting_reqs), 0) + self._emit_request_started_outputs(sched_output) + if sched_output.is_empty: self._emit_finished_outputs(sched_output.finished_req_ids, None) continue @@ -629,6 +638,17 @@ class DiffusionEngine: # Engine is stopping: fail any RPCs still queued so callers don't hang. self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down.")) + def _emit_request_started_outputs(self, sched_output: DiffusionSchedulerOutput) -> None: + """Notify opted-in callers when requests leave the scheduler queue.""" + for new_req in sched_output.scheduled_new_reqs: + request = getattr(new_req, "req", None) + sampling_params = getattr(request, "sampling_params", None) + if getattr(sampling_params, "emit_request_lifecycle", False): + self._put_output( + new_req.request_id, + DiffusionOutput(finished=False, request_started=True), + ) + def _wait_for_admission_if_needed_locked(self) -> None: """Apply scheduler admission policy while holding the engine condition. diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 6c8cc236c..f583aba14 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -19,7 +19,10 @@ from typing import TYPE_CHECKING, Any from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.diffusion.data import DiffusionRequestAbortedError +from vllm_omni.diffusion.data import ( + DiffusionRequestAbortedError, + is_diffusion_request_started_output, +) from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.engine.stage_client import StageClientBase @@ -158,7 +161,13 @@ class InlineStageDiffusionClient(StageClientBase): # only publish the final output. result = None async for results in self._engine.step_streaming(request): - result = results[0] + output = results[0] + if is_diffusion_request_started_output(output): + if not output.request_id: + output.request_id = request_id + self._output_queue.put_nowait(output) + continue + result = output if result is None: raise RuntimeError("Diffusion execution finished without output.") if not result.request_id: diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py index 62c2330d5..4074806e1 100644 --- a/vllm_omni/diffusion/output_formatter.py +++ b/vllm_omni/diffusion/output_formatter.py @@ -5,7 +5,7 @@ from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass, field -from typing import TypeGuard, cast +from typing import Any, TypeGuard, cast from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.io_support import get_diffusion_output_type, supports_audio_output @@ -106,6 +106,7 @@ def format_empty_diffusion_outputs( request: OmniDiffusionRequest, *, finished: bool = True, + custom_output: dict[str, Any] | None = None, ) -> list[OmniRequestOutput]: return [ OmniRequestOutput.from_diffusion( @@ -114,6 +115,7 @@ def format_empty_diffusion_outputs( prompt=request.prompt, metrics={}, latents=None, + custom_output=custom_output, finished=finished, ) ] diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index aedb730e3..220158f1f 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -13,7 +13,7 @@ import asyncio import contextlib import multiprocessing.connection import signal -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace from typing import TYPE_CHECKING, Any @@ -33,7 +33,10 @@ from vllm.v1.engine.utils import ( ) from vllm.v1.utils import shutdown -from vllm_omni.diffusion.data import DiffusionRequestAbortedError +from vllm_omni.diffusion.data import ( + DiffusionRequestAbortedError, + is_diffusion_request_started_output, +) from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.utils.serialization import ( @@ -163,6 +166,7 @@ class StageDiffusionProc: prompt: Any, sampling_params_dict: dict, kv_sender_info: dict[str, Any] | None = None, + on_request_started: Callable[[OmniRequestOutput], Awaitable[None]] | None = None, kv_transfer_params: dict[str, Any] | None = None, ) -> OmniRequestOutput: """Build a diffusion request and consume DiffusionEngine.step_streaming() to completion.""" @@ -180,7 +184,13 @@ class StageDiffusionProc: # return the final output. result = None async for results in self._engine.step_streaming(request): - result = results[0] + output = results[0] + if is_diffusion_request_started_output(output) and on_request_started is not None: + if not output.request_id: + output.request_id = request_id + await on_request_started(output) + continue + result = output if result is None: raise RuntimeError("Diffusion execution finished without output.") if not result.request_id: @@ -354,11 +364,16 @@ class StageDiffusionProc: """Process a single diffusion request and send the response.""" try: if not self._od_config.streaming_output: + + async def _send_request_started(output: OmniRequestOutput) -> None: + await response_socket.send(encoder.encode({"type": "result", "output": output})) + result = await self._process_request( request_id, prompt, sampling_params_dict, kv_sender_info=kv_sender_info, + on_request_started=_send_request_started, kv_transfer_params=kv_transfer_params, ) await response_socket.send(encoder.encode({"type": "result", "output": result})) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 498a9dc89..75620bec9 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -35,6 +35,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker @@ -1199,7 +1200,8 @@ class Orchestrator: idle = False continue - pool.record_output_timestamps([diffusion_output]) + if not is_diffusion_request_started_output(diffusion_output): + pool.record_output_timestamps([diffusion_output]) processed = [diffusion_output] else: raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001) @@ -1441,7 +1443,8 @@ class Orchestrator: if self._absorb_diffusion_metrics(stage_id, replica_id, payload): self._orch_monitor.note_loop(idle=False) continue - pool.record_output_timestamps([payload]) + if not is_diffusion_request_started_output(payload): + pool.record_output_timestamps([payload]) processed = [payload] else: processed = await self._process_llm_stage_outputs( diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 0012039e0..993f0ca1e 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -11,6 +11,7 @@ StageEngineCoreClient instances) instead of OmniStage with worker processes. from __future__ import annotations import asyncio +import os import time import uuid from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence @@ -64,6 +65,8 @@ _FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 # the janus queue's condition variable; this timeout only bounds how often the # orchestrator liveness check runs while the pipeline is idle. _FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 +# Shared DELETE / generate() cleanup abort bound. Env is the documented knob. +ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) class AsyncEventResolver: @@ -661,12 +664,12 @@ class AsyncOmni(EngineClient, OmniBase): except (asyncio.CancelledError, GeneratorExit): self._record_request_failure_once(request_id, reason="client_disconnect") - await self._abort_internal_requests(request_id) + await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: self._record_request_failure_once(request_id, reason="stage_error") - await self._abort_internal_requests(request_id) + await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise finally: @@ -1120,14 +1123,14 @@ class AsyncOmni(EngineClient, OmniBase): return all(bool(item) for item in result) return bool(result) - async def abort(self, request_id: str | Iterable[str]) -> None: + async def abort(self, request_id: str | Iterable[str], *, timeout: float | None = None) -> None: """Abort request(s) via the Orchestrator.""" request_ids = [request_id] if isinstance(request_id, str) else list(request_id) # Map the external user request IDs to internal IDs used by the Orchestrator. # NOTE: If the user request_id matches multiple requests, all of them will be # aborted. This is also what happens in this case in vLLM's output processor. internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id in request_ids] - await self._abort(internal_ids) + await self._abort(internal_ids, timeout=timeout) async def submit_interaction_async( self, @@ -1169,16 +1172,32 @@ class AsyncOmni(EngineClient, OmniBase): if self.log_stats: logger.info("[AsyncOmni] Queued interaction for request %s", request_id) - async def _abort_internal_requests(self, request_id: str | Iterable[str]): + async def _abort_internal_requests( + self, + request_id: str | Iterable[str], + *, + timeout: float = ABORT_TIMEOUT_S, + ): """Abort request(s) via the Orchestrator given internal request IDs, which take the format -. """ request_ids = [request_id] if isinstance(request_id, str) else list(request_id) # Request IDs are already internal, so we just need to get the matching states. internal_req_ids = [rid for rid in request_ids if rid in self.request_states] - await self._abort(internal_req_ids) + try: + # Unbind generate() if abort_async blocks in the executor. + await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) + except TimeoutError: + logger.warning( + "[AsyncOmni] Timed out aborting %s after %.1fs; " + "engine abort is best-effort until the current batch drains", + ",".join(internal_req_ids), + timeout, + ) + except Exception: + logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) - async def _abort(self, request_ids: list[str]) -> None: + async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: """Abort request IDs via the engine and enqueue terminal abort outputs. Waits for orchestrator abort acknowledgment, enqueues any AR terminal @@ -1191,7 +1210,7 @@ class AsyncOmni(EngineClient, OmniBase): registered yet, unbound replica, or orchestrator id drop), enqueue a synthetic finished abort so ``generate()`` cannot hang on ``queue.get``. """ - abort_outputs = await self.engine.abort_async(request_ids) or [] + abort_outputs = await self.engine.abort_async(request_ids, timeout=timeout) or [] delivered: set[str] = set() for output_msg in abort_outputs: req_id = getattr(output_msg, "request_id", None) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 8acd1d6fa..36b7544ee 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -86,7 +86,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.endpoint_policy import ( shutdown_unsupported_routes, ) -from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni from vllm_omni.entrypoints.duplex.capability import should_enable_duplex_endpoint from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime @@ -201,6 +201,8 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingName logger = init_logger(__name__) router = APIRouter() +VIDEO_ABORT_TIMEOUT_S = ABORT_TIMEOUT_S + profiler_router = APIRouter() @@ -2518,11 +2520,11 @@ async def retrieve_video(video_id: str) -> VideoResponse | JSONResponse: @router.delete("/v1/videos/{video_id}") -async def delete_video(video_id: str) -> VideoDeleteResponse: +async def delete_video(video_id: str, raw_request: Request) -> VideoDeleteResponse: """Delete a stored video job and any generated output. - If the job is still queued or running, this endpoint first attempts to - cancel the in-flight generation task before removing the stored metadata. + In-flight jobs get a bounded engine abort, then frontend cancel. The job + is re-read afterwards so a completed save is not orphaned. Args: video_id: Identifier of the video job to delete. @@ -2539,19 +2541,37 @@ async def delete_video(video_id: str) -> VideoDeleteResponse: raise HTTPException(status_code=404, detail="Video not found") if job.status in (VideoGenerationStatus.QUEUED, VideoGenerationStatus.IN_PROGRESS): + handler = raw_request.app.state.openai_serving_video + try: + await asyncio.wait_for(handler.abort_request(video_id), timeout=VIDEO_ABORT_TIMEOUT_S) + except asyncio.TimeoutError: + logger.warning( + "Timed out aborting video request %s after %.1fs; " + "engine abort is best-effort until the current batch drains", + video_id, + VIDEO_ABORT_TIMEOUT_S, + ) + except Exception: + logger.exception("Failed to abort in-flight video request %s", video_id) task = await VIDEO_TASKS.get(video_id) if task is not None: task.cancel() try: - await asyncio.wait_for(task, timeout=2.0) + # Cancel cleanup may spend a full abort budget; +2s covers scheduling slack. + await asyncio.wait_for(task, timeout=VIDEO_ABORT_TIMEOUT_S + 2.0) except asyncio.TimeoutError: raise HTTPException(status_code=409, detail="Cancellation in progress. Please try again later.") except asyncio.CancelledError: pass + job = await VIDEO_STORE.get(video_id) + if job is None: + return VideoDeleteResponse(id=video_id, deleted=True) + if job.status in (VideoGenerationStatus.QUEUED, VideoGenerationStatus.IN_PROGRESS): await VIDEO_STORE.pop(video_id) return VideoDeleteResponse(id=job.id, deleted=True) - elif job.status is VideoGenerationStatus.FAILED: + + if job.status is VideoGenerationStatus.FAILED: if job.file_name is not None: try: await STORAGE_MANAGER.delete(video_id) diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 5ff055fa2..9a24a1ef9 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -6,7 +6,7 @@ from __future__ import annotations import copy import math import time -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from functools import cached_property from http import HTTPStatus @@ -19,9 +19,10 @@ from PIL import Image from vllm.engine.protocol import EngineClient from vllm.logger import init_logger +from vllm_omni.diffusion.data import is_diffusion_request_started_output from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.media_utils import count_mp4_frames, normalize_preencode_batch_frames -from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni import ABORT_TIMEOUT_S, AsyncOmni from vllm_omni.entrypoints.openai.protocol.videos import ( VideoAction, VideoData, @@ -256,6 +257,9 @@ class OmniOpenAIServingVideo: def shutdown(self) -> None: self._video_frame_converter.shutdown() + async def abort_request(self, request_id: str) -> None: + await self._engine_client.abort(request_id, timeout=ABORT_TIMEOUT_S) + async def _run_and_extract( self, request: VideoGenerationRequest, @@ -264,6 +268,7 @@ class OmniOpenAIServingVideo: reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, + on_started: Callable[[], Awaitable[None]] | None = None, ) -> VideoGenerationArtifacts: """Run the generation pipeline and extract video/audio/profiler outputs.""" prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt, modalities=["video"]) @@ -444,7 +449,12 @@ class OmniOpenAIServingVideo: gen_params.seed, ) - result = await self._run_generation(prompt, gen_params, reference_id) + result = await self._run_generation( + prompt, + gen_params, + reference_id, + on_started=on_started, + ) multimodal_output = self._extract_multimodal_output(result) metadata = multimodal_output.get("metadata") if isinstance(multimodal_output, dict) else {} common_metadata = metadata.get("common") if isinstance(metadata, dict) else {} @@ -528,6 +538,7 @@ class OmniOpenAIServingVideo: reference_image: ReferenceImage | None = None, reference_video: ReferenceVideo | None = None, reference_audio: ReferenceAudio | None = None, + on_started: Callable[[], Awaitable[None]] | None = None, ) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: """Generate a video and return raw MP4 bytes, bypassing base64 encoding.""" artifacts = await self._run_and_extract( @@ -536,6 +547,7 @@ class OmniOpenAIServingVideo: reference_image=reference_image, reference_video=reference_video, reference_audio=reference_audio, + on_started=on_started, ) if len(artifacts.videos) > 1: logger.warning( @@ -641,6 +653,8 @@ class OmniOpenAIServingVideo: prompt: OmniTextPrompt, gen_params: OmniDiffusionSamplingParams, request_id: str, + *, + on_started: Callable[[], Awaitable[None]] | None = None, ) -> object: stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None) @@ -660,6 +674,7 @@ class OmniOpenAIServingVideo: # Common generation logic for both paths engine_client = cast(AsyncOmni, self._engine_client) + gen_params.emit_request_lifecycle = on_started is not None sampling_params_list = build_stage_sampling_params_list( list(stage_configs), get_default_sampling_params_list(engine_client), @@ -668,11 +683,17 @@ class OmniOpenAIServingVideo: ) result = None + started_notified = False async for output in engine_client.generate( prompt=prompt, request_id=request_id, sampling_params_list=sampling_params_list, ): + if is_diffusion_request_started_output(output): + if on_started is not None and not started_notified: + await on_started() + started_notified = True + continue result = output if result is None: diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py index e23d876e5..4473fce07 100644 --- a/vllm_omni/entrypoints/openai/video/generation/helpers.py +++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py @@ -330,9 +330,12 @@ async def _run_video_generation_job( _cleanup_video_references(reference_video, reference_audio, control_path) return - await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) started_at = time.perf_counter() try: + + async def _mark_started() -> None: + await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) + video_bytes, stage_durations, peak_memory_mb, action, video_metadata = _unpack_video_generation_result( await handler.generate_video_bytes( request, @@ -340,6 +343,7 @@ async def _run_video_generation_job( reference_image=reference_image, reference_video=reference_video, reference_audio=reference_audio, + on_started=_mark_started, ) ) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index f9f0a9a9a..287ddb3fb 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -349,6 +349,11 @@ class OmniDiffusionSamplingParams: # results output: torch.Tensor | None = None + # Internal control-plane event used by asynchronous APIs. When enabled, + # the diffusion engine emits a non-terminal output when the scheduler + # first admits this request for execution. + emit_request_lifecycle: bool = False + def __post_init__(self) -> None: if self.quality is not None and self.quality not in DIFFUSION_QUALITY_LEVELS: raise ValueError(f"quality must be one of {list(DIFFUSION_QUALITY_LEVELS)}, got {self.quality!r}")