[Model] Overlap MiniMax-H3 reference audio extraction with video decode

tests/model_executor/models/test_minimax_h3_audio_overlap.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+import shutil
+import subprocess
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+
+import pytest
+import torch
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
+from vllm_omni.model_executor.models.minimax_h3 import reference_video
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def reference(tmp_path):
+    if shutil.which("ffmpeg") is None:
+        pytest.skip("ffmpeg is required for real reference-audio tests")
+    path = tmp_path / "reference.mp4"
+    subprocess.run(
+        [
+            "ffmpeg",
+            "-y",
+            "-loglevel",
+            "error",
+            "-f",
+            "lavfi",
+            "-i",
+            "color=c=blue:s=64x64:r=24",
+            "-f",
+            "lavfi",
+            "-i",
+            "sine=frequency=440:sample_rate=44100",
+            "-t",
+            "3",
+            "-c:v",
+            "mpeg4",
+            "-c:a",
+            "aac",
+            str(path),
+        ],
+        check=True,
+        timeout=30,
+    )
+    return {
+        "original_path": str(path),
+        "prepared_path": str(path),
+        "input_has_audio": True,
+        "audio_duration_seconds": 2.5,
+    }
+
+
+def _prepare(references):
+    return processing.prepare_encoder_inputs(
+        {"prompt": "reference", "multi_modal_data": {"video": [item["original_path"] for item in references]}},
+        OmniDiffusionSamplingParams(height=64, width=64, extra_args={"task": "ref2va", "duration": 4.4}),
+        prepared_reference_videos=references,
+    )
+
+
+def test_audio_overlap_preserves_reference_order_and_waveforms(reference):
+    silent = {**reference, "input_has_audio": False}
+    prepared = _prepare([reference, silent, reference])
+    expected, rate = reference_video.load_video_audio(reference["original_path"], duration_seconds=2.5)
+    assert prepared.media.video_audios[1] is None
+    for index in (0, 2):
+        waveform, actual_rate = prepared.media.video_audios[index]
+        assert actual_rate == rate
+        torch.testing.assert_close(waveform, expected, rtol=0, atol=0)
+    assert prepared.condition_labels == [("audio", 1), ("video", 1), ("video", 2), ("audio", 2), ("video", 3)]
+
+
+def test_frame_failure_cancels_queued_audio_and_joins_worker(reference, monkeypatch):
+    # Hold the real executor's worker so both audio jobs remain queued.
+    pool = ThreadPoolExecutor(max_workers=1)
+    release = Event()
+    started = Event()
+    futures = []
+    submit = pool.submit
+    shutdown = pool.shutdown
+
+    def hold_worker():
+        started.set()
+        assert release.wait(10)
+
+    blocker = submit(hold_worker)
+    assert started.wait(10)
+
+    def record_submit(*args, **kwargs):
+        future = submit(*args, **kwargs)
+        futures.append(future)
+        return future
+
+    def release_and_shutdown(*, wait=True, cancel_futures=False):
+        shutdown(wait=False, cancel_futures=cancel_futures)
+        release.set()
+        shutdown(wait=wait)
+
+    def fail_frames(path):
+        raise RuntimeError("frame decoding failed")
+
+    monkeypatch.setattr(pool, "submit", record_submit)
+    monkeypatch.setattr(pool, "shutdown", release_and_shutdown)
+    monkeypatch.setattr(processing, "ThreadPoolExecutor", lambda **kwargs: pool)
+    monkeypatch.setattr(processing, "load_video_frames", fail_frames)
+    try:
+        with pytest.raises(RuntimeError, match="frame decoding failed"):
+            _prepare([reference, reference])
+        assert len(futures) == 2
+        assert all(future.cancelled() for future in futures)
+        assert blocker.done()
+    finally:
+        release.set()
+        shutdown(wait=True, cancel_futures=True)
+
+
+@pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="requires a POSIX FIFO")
+def test_audio_extraction_timeout_propagates_and_cleans_tempdir(tmp_path, monkeypatch):
+    if shutil.which("ffmpeg") is None:
+        pytest.skip("ffmpeg is required for the subprocess timeout test")
+    fifo = tmp_path / "blocked_audio"
+    getattr(os, "mkfifo")(fifo)
+    monkeypatch.setattr(reference_video, "_AUDIO_EXTRACTION_TIMEOUT_SECONDS", 0.2)
+    with pytest.raises(subprocess.TimeoutExpired) as error:
+        reference_video.load_video_audio(str(fifo))
+    assert error.value.timeout == pytest.approx(0.2, abs=0.02)
+    assert not os.path.exists(os.path.dirname(error.value.cmd[-1]))

vllm_omni/model_executor/models/minimax_h3/encoder_processing.py

@@ -9,6 +9,7 @@
 import subprocess
 import tempfile
 from collections.abc import Callable, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass
 from typing import Any
@@ -582,41 +583,52 @@ def prepare_encoder_inputs(
                         workdir=workdir,
                         start_time_seconds=extra_args.get("start_time_seconds"),
                     )
-                for item in prepared_videos:
-                    full_frames = load_video_frames(item["prepared_path"])
-                    encoded_video_inputs.append(_frames_to_tensor(full_frames))
-                    sampled = sample_reference_video_frames(
-                        item["prepared_path"],
-                        decoded_frames=full_frames,
-                    )
-                    video_timestamps.append(sampled["block_timestamps"])
-                    frames = np.stack(sampled["frames"])
-                    frame_count = int(frames.shape[0])
-                    qwen_video_inputs.append(
-                        (
-                            frames,
-                            {
-                                "total_num_frames": frame_count,
-                                "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
-                                "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
-                                "video_backend": "minimax_h3",
-                                "frames_indices": list(range(frame_count)),
-                                "do_sample_frames": False,
-                            },
-                        )
-                    )
-                    if item["input_has_audio"]:
-                        waveform, sample_rate = load_video_audio(
+                audio_pool = ThreadPoolExecutor(max_workers=1)
+                try:
+                    audio_futures = [
+                        audio_pool.submit(
+                            load_video_audio,
                             item["original_path"],
                             start_time_seconds=float(item.get("start_time_seconds", 0.0)),
                             duration_seconds=item.get(
                                 "audio_duration_seconds",
                                 item.get("duration_seconds"),
                             ),
                         )
-                        video_audio_inputs.append((waveform.float().contiguous(), int(sample_rate)))
-                    else:
-                        video_audio_inputs.append(None)
+                        if item["input_has_audio"]
+                        else None
+                        for item in prepared_videos
+                    ]
+                    for item, audio_future in zip(prepared_videos, audio_futures):
+                        full_frames = load_video_frames(item["prepared_path"])
+                        encoded_video_inputs.append(_frames_to_tensor(full_frames))
+                        sampled = sample_reference_video_frames(
+                            item["prepared_path"],
+                            decoded_frames=full_frames,
+                        )
+                        video_timestamps.append(sampled["block_timestamps"])
+                        frames = np.stack(sampled["frames"])
+                        frame_count = int(frames.shape[0])
+                        qwen_video_inputs.append(
+                            (
+                                frames,
+                                {
+                                    "total_num_frames": frame_count,
+                                    "fps": MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
+                                    "duration": frame_count / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS,
+                                    "video_backend": "minimax_h3",
+                                    "frames_indices": list(range(frame_count)),
+                                    "do_sample_frames": False,
+                                },
+                            )
+                        )
+                        if audio_future is not None:
+                            waveform, sample_rate = audio_future.result()
+                            video_audio_inputs.append((waveform.float().contiguous(), int(sample_rate)))
+                        else:
+                            video_audio_inputs.append(None)
+                finally:
+                    audio_pool.shutdown(wait=True, cancel_futures=True)
         audio_index = 0
         for video_index, item in enumerate(prepared_videos, start=1):
             if item["input_has_audio"]:

vllm_omni/model_executor/models/minimax_h3/reference_video.py

@@ -688,6 +688,9 @@ def load_audio_file(path: str) -> tuple[torch.Tensor, int]:
             return _soundfile_to_waveform(wav)
 
 
+_AUDIO_EXTRACTION_TIMEOUT_SECONDS = 120
+
+
 def load_video_audio(
     path: str,
     *,
@@ -718,7 +721,7 @@ def load_video_audio(
         if duration_seconds is not None:
             command.extend(["-t", f"{float(duration_seconds):.6f}"])
         command.append(output)
-        subprocess.run(command, check=True)
+        subprocess.run(command, check=True, timeout=_AUDIO_EXTRACTION_TIMEOUT_SECONDS)
         return load_audio_file(output)