[Bugfix] Stop caching synthetic videos without requested audio

tests/helpers/media.py

@@ -308,16 +308,12 @@ def _mux_mp4_bytes_with_synthetic_audio(
     duration_sec = num_frames / fps if fps > 0 else 0.0
     duration_int = max(1, int(math.ceil(duration_sec)))
 
-    try:
-        audio_result = generate_synthetic_audio(
-            duration=duration_int,
-            num_channels=1,
-            sample_rate=sample_rate,
-        )
-        audio_pcm = audio_result["np_array"]
-    except Exception as e:
-        logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e)
-        return video_mp4_bytes
+    audio_result = generate_synthetic_audio(
+        duration=duration_int,
+        num_channels=1,
+        sample_rate=sample_rate,
+    )
+    audio_pcm = audio_result["np_array"]
 
     try:
         import imageio_ffmpeg
@@ -326,47 +322,38 @@ def _mux_mp4_bytes_with_synthetic_audio(
     except Exception:
         ffmpeg_exe = "ffmpeg"
 
-    try:
-        with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
-            vid_path = os.path.join(tmp, "video.mp4")
-            wav_path = os.path.join(tmp, "audio.wav")
-            out_path = os.path.join(tmp, "out.mp4")
-            with open(vid_path, "wb") as f:
-                f.write(video_mp4_bytes)
-            sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
-            cmd = [
-                ffmpeg_exe,
-                "-y",
-                "-nostdin",
-                "-hide_banner",
-                "-loglevel",
-                "error",
-                "-i",
-                vid_path,
-                "-i",
-                wav_path,
-                "-c:v",
-                "copy",
-                "-c:a",
-                "aac",
-                "-b:a",
-                "128k",
-                "-shortest",
-                "-movflags",
-                "+faststart",
-                out_path,
-            ]
-            subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300)
-            with open(out_path, "rb") as f:
-                return f.read()
-    except (
-        FileNotFoundError,
-        subprocess.CalledProcessError,
-        subprocess.TimeoutExpired,
-        OSError,
-    ) as e:
-        logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e)
-        return video_mp4_bytes
+    with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
+        vid_path = os.path.join(tmp, "video.mp4")
+        wav_path = os.path.join(tmp, "audio.wav")
+        out_path = os.path.join(tmp, "out.mp4")
+        with open(vid_path, "wb") as f:
+            f.write(video_mp4_bytes)
+        sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
+        cmd = [
+            ffmpeg_exe,
+            "-y",
+            "-nostdin",
+            "-hide_banner",
+            "-loglevel",
+            "error",
+            "-i",
+            vid_path,
+            "-i",
+            wav_path,
+            "-c:v",
+            "copy",
+            "-c:a",
+            "aac",
+            "-b:a",
+            "128k",
+            "-shortest",
+            "-movflags",
+            "+faststart",
+            out_path,
+        ]
+        subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300)
+        with open(out_path, "rb") as f:
+            return f.read()
 
 
 def generate_synthetic_video(
@@ -383,18 +370,49 @@ def generate_synthetic_video(
     ``width`` / ``height`` / ``num_frames`` / ``embed_audio`` unless
     ``force_regenerate`` is true. Cache root: ``cache_dir`` if given, else the
     default temp subdirectory.
+
+    ``embed_audio=True`` requires successful TTS generation and FFmpeg muxing.
+    Cached MP4s with corrupt/incomplete media or without the requested audio
+    stream are regenerated. Cache files are published atomically.
     """
     root = _resolve_synthetic_media_cache_dir(cache_dir)
     root.mkdir(parents=True, exist_ok=True)
     cache_path = root / f"synth_video_w{width}_h{height}_nf{num_frames}_ea{int(embed_audio)}.mp4"
 
     if not force_regenerate and cache_path.is_file():
         video_bytes = cache_path.read_bytes()
-        return {
-            "np_array": _np_array_from_mp4_bytes(video_bytes),
-            "base64": base64.b64encode(video_bytes).decode("utf-8"),
-            "file_path": str(cache_path.resolve()),
-        }
+        has_required_audio = True
+        if embed_audio:
+            import av
+
+            try:
+                with av.open(io.BytesIO(video_bytes)) as container:
+                    has_required_audio = False
+                    if container.streams.audio:
+                        audio_stream = container.streams.audio[0]
+                        samples = sum(frame.samples for frame in container.decode(audio=0))
+                        expected_samples = 0
+                        if audio_stream.duration is not None and audio_stream.time_base is not None:
+                            expected_samples = math.ceil(
+                                audio_stream.duration * audio_stream.time_base * audio_stream.codec_context.sample_rate
+                            )
+                        has_required_audio = samples > 0 and samples >= expected_samples
+            except av.error.FFmpegError:
+                has_required_audio = False
+        if has_required_audio:
+            try:
+                cached_frames = _np_array_from_mp4_bytes(video_bytes)
+            except RuntimeError:
+                cached_frames = None
+            if cached_frames is not None and len(cached_frames) == num_frames:
+                return {
+                    "np_array": cached_frames,
+                    "base64": base64.b64encode(video_bytes).decode("utf-8"),
+                    "file_path": str(cache_path.resolve()),
+                }
+        logger.warning(
+            "Cached synthetic video %s is incomplete or lacks the requested audio; regenerating.", cache_path
+        )
 
     import cv2
     import imageio
@@ -475,7 +493,13 @@ class _BouncingBall(TypedDict):
         else video_only_bytes
     )
 
-    cache_path.write_bytes(video_bytes)
+    temp_cache = tempfile.NamedTemporaryFile(dir=root, prefix=f".{cache_path.name}.", suffix=".tmp", delete=False)
+    try:
+        with temp_cache:
+            temp_cache.write(video_bytes)
+        os.replace(temp_cache.name, cache_path)
+    finally:
+        Path(temp_cache.name).unlink(missing_ok=True)
 
     return {
         "np_array": np.array(video_frames),

tests/helpers/tests/test_synthetic_video.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import base64
+import io
+import subprocess
+from pathlib import Path
+
+import av
+import numpy as np
+import pytest
+
+from tests.helpers import media
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def synthetic_audio(mocker):
+    """Keep video tests independent of the host's TTS installation."""
+    sample_rate = 48000
+    samples = np.arange(sample_rate, dtype=np.float32) / sample_rate
+    waveform = (0.1 * np.sin(2 * np.pi * 440 * samples)).reshape(-1, 1)
+    return mocker.patch.object(media, "generate_synthetic_audio", return_value={"np_array": waveform})
+
+
+def test_audio_generation_failure_is_not_cached(mocker, tmp_path):
+    failure = RuntimeError("TTS dependency unavailable")
+    mocker.patch.object(media, "generate_synthetic_audio", side_effect=failure)
+
+    with pytest.raises(RuntimeError, match="TTS dependency unavailable") as caught:
+        media.generate_synthetic_video(64, 64, 6, embed_audio=True, cache_dir=tmp_path)
+
+    assert caught.value is failure
+    assert not list(tmp_path.glob("*.mp4"))
+
+
+@pytest.mark.parametrize(
+    "failure",
+    [
+        FileNotFoundError("FFmpeg unavailable"),
+        subprocess.CalledProcessError(1, ["ffmpeg"]),
+        subprocess.TimeoutExpired(["ffmpeg"], 300),
+    ],
+    ids=["missing-ffmpeg", "ffmpeg-error", "ffmpeg-timeout"],
+)
+def test_audio_mux_failure_is_propagated(mocker, synthetic_audio, failure):
+    mocker.patch("imageio_ffmpeg.get_ffmpeg_exe", return_value="ffmpeg")
+    mocker.patch.object(media.subprocess, "run", side_effect=failure)
+
+    with pytest.raises(type(failure)) as caught:
+        media._mux_mp4_bytes_with_synthetic_audio(b"video", num_frames=6)
+
+    assert caught.value is failure
+
+
+def test_audio_in_video_repairs_video_only_cache(synthetic_audio, tmp_path):
+    video_only = media.generate_synthetic_video(64, 64, 6, cache_dir=tmp_path)
+    # Reproduce a cache entry written by the old silent audio-mux fallback.
+    cache_path = tmp_path / "synth_video_w64_h64_nf6_ea1.mp4"
+    original_bytes = Path(video_only["file_path"]).read_bytes()
+    cache_path.write_bytes(original_bytes)
+
+    result = media.generate_synthetic_video(64, 64, 6, embed_audio=True, cache_dir=tmp_path)
+
+    assert result["file_path"] == str(cache_path)
+    raw = base64.b64decode(result["base64"])
+    assert raw == cache_path.read_bytes()
+    assert raw != original_bytes
+    assert result["np_array"].shape == (6, 64, 64, 3)
+    with av.open(io.BytesIO(raw)) as container:
+        assert len(container.streams.video) == 1
+        assert len(container.streams.audio) == 1
+        audio_frames = list(container.decode(audio=0))
+        assert sum(frame.samples for frame in audio_frames) > 0
+        assert any(np.any(frame.to_ndarray()) for frame in audio_frames)
+
+
+@pytest.mark.parametrize("embed_audio", [False, True], ids=["video-only", "with-audio"])
+def test_valid_video_cache_is_reused(mocker, synthetic_audio, tmp_path, embed_audio):
+    result = media.generate_synthetic_video(64, 64, 6, embed_audio=embed_audio, cache_dir=tmp_path)
+    encode = mocker.patch("imageio.get_writer", side_effect=AssertionError("valid cache was regenerated"))
+
+    cached = media.generate_synthetic_video(64, 64, 6, embed_audio=embed_audio, cache_dir=tmp_path)
+
+    encode.assert_not_called()
+    assert cached["base64"] == result["base64"]
+    assert cached["file_path"] == result["file_path"]
+    assert np.array_equal(cached["np_array"], media._np_array_from_mp4_bytes(base64.b64decode(result["base64"])))
+
+
+@pytest.mark.parametrize("corruption", ["empty", "invalid-header", "truncated-audio"])
+def test_corrupt_video_cache_is_regenerated(synthetic_audio, tmp_path, corruption):
+    result = media.generate_synthetic_video(64, 64, 6, embed_audio=True, cache_dir=tmp_path)
+    cache_path = Path(result["file_path"])
+    raw = cache_path.read_bytes()
+    if corruption == "empty":
+        corrupt = b""
+    elif corruption == "invalid-header":
+        corrupt = b"not an MP4"
+    else:
+        with av.open(io.BytesIO(raw)) as container:
+            packets = [packet for packet in container.demux(audio=0) if packet.pos is not None and packet.size]
+        packet = packets[len(packets) // 2]
+        corrupt = raw[: packet.pos + packet.size // 2]
+        # +faststart leaves the stream metadata readable despite missing data.
+        with av.open(io.BytesIO(corrupt)) as container:
+            assert container.streams.audio
+    cache_path.write_bytes(corrupt)
+
+    repaired = media.generate_synthetic_video(64, 64, 6, embed_audio=True, cache_dir=tmp_path)
+
+    repaired_bytes = base64.b64decode(repaired["base64"])
+    assert repaired_bytes == cache_path.read_bytes()
+    assert repaired_bytes != corrupt
+    assert repaired["np_array"].shape == (6, 64, 64, 3)
+    with av.open(io.BytesIO(repaired_bytes)) as container:
+        assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+
+
+def test_video_cache_is_published_atomically(monkeypatch, synthetic_audio, tmp_path):
+    original = media.generate_synthetic_video(64, 64, 6, embed_audio=True, cache_dir=tmp_path)
+    cache_path = Path(original["file_path"])
+    old_bytes = cache_path.read_bytes()
+    replace = media.os.replace
+    published = []
+
+    def observe_replace(source, destination):
+        source = Path(source)
+        assert Path(destination) == cache_path
+        assert source.parent == cache_path.parent
+        # Until publication, readers must still see the old complete MP4.
+        assert cache_path.read_bytes() == old_bytes
+        completed = source.read_bytes()
+        with av.open(io.BytesIO(completed)) as container:
+            assert len(list(container.decode(video=0))) == 6
+        with av.open(io.BytesIO(completed)) as container:
+            assert sum(frame.samples for frame in container.decode(audio=0)) > 0
+        published.append(completed)
+        replace(source, destination)
+
+    monkeypatch.setattr(media.os, "replace", observe_replace)
+    result = media.generate_synthetic_video(64, 64, 6, embed_audio=True, force_regenerate=True, cache_dir=tmp_path)
+
+    assert published == [cache_path.read_bytes()]
+    assert base64.b64decode(result["base64"]) == published[0]
+    assert list(tmp_path.iterdir()) == [cache_path]
+
+
+def test_failed_cache_publication_preserves_old_file(mocker, synthetic_audio, tmp_path):
+    original = media.generate_synthetic_video(64, 64, 6, embed_audio=True, cache_dir=tmp_path)
+    cache_path = Path(original["file_path"])
+    old_bytes = cache_path.read_bytes()
+    mocker.patch.object(media.os, "replace", side_effect=OSError("publication failed"))
+
+    with pytest.raises(OSError, match="publication failed"):
+        media.generate_synthetic_video(64, 64, 6, embed_audio=True, force_regenerate=True, cache_dir=tmp_path)
+
+    assert cache_path.read_bytes() == old_bytes
+    assert list(tmp_path.iterdir()) == [cache_path]