Decode input_audio media containers with PyAV & Update memory profiler (#31832)

Signed-off-by: Shiyan Deng <dsy842974287@meta.com>
Signed-off-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
Co-authored-by: Lucia Fang <116399278+luccafong@users.noreply.github.com>

python/pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
   "aiohttp",
   "anthropic>=0.20.0",
   "apache-tvm-ffi==0.1.11",
-  "av ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
+  "av==16.1.0 ; sys_platform == 'linux' and (platform_machine == 'aarch64' or platform_machine == 'arm64' or platform_machine == 'armv7l')",
   "blobfile==3.0.0",
   "build",
   "compressed-tensors",
@@ -144,6 +144,7 @@ test = [
   "accelerate",
   "addict",
   "auto-round>=0.13.1",
+  "av==16.1.0",
   "bitsandbytes",
   "pymupdf",
   "diff-cover",

python/sglang/srt/environ.py

@@ -331,6 +331,9 @@ class Envs:
     SGLANG_SIMULATE_UNIFORM_EXPERTS = EnvBool(False)
     SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS = EnvBool(False)
     SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp")
+    # Allocator-history buffer for /start_profile activities=["MEM"]; the
+    # default truncates long windows (each entry is one alloc/free event).
+    SGLANG_MEM_PROFILE_MAX_ENTRIES = EnvInt(100000)
     SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500)
     SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64)
     SGLANG_NATIVE_MOVE_KV_CACHE = EnvBool(False)

python/sglang/srt/managers/scheduler_components/profiler_manager.py

@@ -251,7 +251,9 @@ class SchedulerProfilerManager:
             self.profile_in_progress = True
 
         if "MEM" in activities:
-            torch.cuda.memory._record_memory_history(max_entries=100000)
+            torch.cuda.memory._record_memory_history(
+                max_entries=envs.SGLANG_MEM_PROFILE_MAX_ENTRIES.get()
+            )
             self.profile_in_progress = True
 
         if "CUDA_PROFILER" in activities:
@@ -356,7 +358,8 @@ class SchedulerProfilerManager:
         if self.profiler_activities is not None and "MEM" in self.profiler_activities:
             memory_profile_path = os.path.join(
                 self.torch_profiler_output_dir,
-                str(time.time())
+                stage_prefix
+                + str(time.time())
                 + f"-TP-{self.ps.tp_rank}-memory"
                 + stage_suffix
                 + ".pickle",

python/sglang/srt/multimodal/audio_from_video.py

@@ -11,11 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Extract audio from video bytes using PyAV (in-process, CUDA-safe).
+"""Decode audio from media containers with PyAV."""
-
-PyAV wraps FFmpeg's C libraries in-process, avoiding subprocess forks which
-would crash CUDA-active workers.
-"""
 
 import io
 import logging
@@ -24,66 +20,129 @@ import numpy as np
 
 logger = logging.getLogger(__name__)
 
+_INVALID_AUDIO_CONTAINER_MESSAGE = (
+    "Invalid input_audio: no decodable audio stream was found in the media container."
+)
+_AUDIO_CONTAINER_SIGNATURES = (
+    ((4, b"ftyp"),),
+    ((0, b"RIFF"), (8, b"AVI ")),
+    ((0, b"#!AMR\n"),),
+    ((0, b"#!AMR-WB\n"),),
+    # EBML magic: WebM / Matroska (e.g. browser MediaRecorder output)
+    ((0, b"\x1a\x45\xdf\xa3"),),
+)
+
+
+class _AudioContainerDecodeError(ValueError):
+    pass
+
+
+def is_audio_container(data: bytes) -> bool:
+    """Return whether the header identifies a supported media container."""
+    return any(
+        all(data[offset : offset + len(magic)] == magic for offset, magic in signature)
+        for signature in _AUDIO_CONTAINER_SIGNATURES
+    )
+
+
+def _append_resampled_frames(
+    chunks: list[np.ndarray],
+    frames,
+    *,
+    mono: bool,
+) -> None:
+    for frame in frames:
+        array = frame.to_ndarray()
+        chunks.append(array.reshape(-1) if mono else array.T)
+
+
+def decode_audio_container(
+    source: bytes | str,
+    *,
+    target_sr: int,
+    mono: bool,
+) -> np.ndarray:
+    """Strictly decode the first audio stream from a supported container."""
+    if not isinstance(target_sr, int) or target_sr <= 0:
+        raise ValueError(_INVALID_AUDIO_CONTAINER_MESSAGE)
+    if isinstance(source, bytes) and not source:
+        raise ValueError(_INVALID_AUDIO_CONTAINER_MESSAGE)
+
+    try:
+        import av
+
+        input_source = io.BytesIO(source) if isinstance(source, bytes) else source
+        with av.open(input_source) as container:
+            if not container.streams.audio:
+                raise _AudioContainerDecodeError(_INVALID_AUDIO_CONTAINER_MESSAGE)
+
+            audio_stream = container.streams.audio[0]
+            if mono:
+                resampler = av.audio.resampler.AudioResampler(
+                    format="fltp", layout="mono", rate=target_sr
+                )
+            else:
+                resampler = av.audio.resampler.AudioResampler(
+                    format="fltp", rate=target_sr
+                )
+
+            chunks: list[np.ndarray] = []
+            skipped_packets = 0
+            first_decode_error = None
+            for packet in container.demux(audio_stream):
+                try:
+                    for frame in packet.decode():
+                        _append_resampled_frames(
+                            chunks,
+                            resampler.resample(frame),
+                            mono=mono,
+                        )
+                except av.error.FFmpegError as error:
+                    skipped_packets += 1
+                    if first_decode_error is None:
+                        first_decode_error = error
+
+            _append_resampled_frames(
+                chunks,
+                resampler.resample(None),
+                mono=mono,
+            )
+            if skipped_packets:
+                logger.warning(
+                    "Skipped %d undecodable audio packet(s); kept %d decoded "
+                    "chunk(s). First decode error: %s",
+                    skipped_packets,
+                    len(chunks),
+                    first_decode_error,
+                )
+    except _AudioContainerDecodeError:
+        raise
+    except Exception as error:
+        raise ValueError(_INVALID_AUDIO_CONTAINER_MESSAGE) from error
+
+    if not chunks:
+        if first_decode_error is not None:
+            raise ValueError(_INVALID_AUDIO_CONTAINER_MESSAGE) from first_decode_error
+        raise ValueError(_INVALID_AUDIO_CONTAINER_MESSAGE)
+
+    waveform = np.concatenate(chunks, axis=0)
+    expected_ndim = 1 if mono else 2
+    if waveform.ndim != expected_ndim or waveform.size == 0:
+        raise ValueError(_INVALID_AUDIO_CONTAINER_MESSAGE)
+    return np.ascontiguousarray(waveform, dtype=np.float32)
+
 
 def extract_audio_from_video_bytes(
     video_bytes: bytes,
     target_sr: int = 16000,
 ) -> np.ndarray | None:
-    """Extract mono audio from video bytes at the target sample rate.
+    """Extract optional mono audio for callers that accept silent videos."""
-
-    Args:
-        video_bytes: Raw video file bytes (e.g. MP4).
-        target_sr: Target sample rate for the output waveform.
-
-    Returns:
-        1-D float32 numpy array of audio samples, or None if the video
-        has no audio track.
-    """
     try:
-        import av
+        return decode_audio_container(
-    except ImportError:
+            video_bytes,
-        logger.warning(
+            target_sr=target_sr,
-            "PyAV (av) is not installed. Cannot extract audio from video. "
+            mono=True,
-            "Install with: pip install av"
         )
-        return None
-
-    try:
-        container = av.open(io.BytesIO(video_bytes))
-    except Exception:
-        logger.warning("Failed to open video bytes for audio extraction")
-        return None
-
-    if not container.streams.audio:
-        container.close()
-        return None
-
-    try:
-        audio_stream = container.streams.audio[0]
-        native_sr = audio_stream.rate or target_sr
-
-        resampler = av.audio.resampler.AudioResampler(
-            format="flt",
-            layout="mono",
-            rate=target_sr,
-        )
-
-        chunks = []
-        for frame in container.decode(audio=0):
-            resampled = resampler.resample(frame)
-            for rf in resampled:
-                arr = rf.to_ndarray().flatten()
-                chunks.append(arr)
-
-        container.close()
-
-        if not chunks:
-            return None
-
-        waveform = np.concatenate(chunks).astype(np.float32)
-        return waveform
-
     except Exception:
         logger.warning("Error extracting audio from video", exc_info=True)
-        container.close()
         return None

python/sglang/srt/multimodal/processors/base_processor.py

@@ -1013,6 +1013,13 @@ class BaseMultimodalProcessor(ABC):
         for modality, idx, future in futures:
             try:
                 result = await asyncio.wrap_future(future)
+            except ValueError:
+                logger.exception(
+                    "[load_mm_data(simple)] error loading %s data at index=%d",
+                    modality.name,
+                    idx,
+                )
+                raise
             except Exception as e:
                 logger.exception(
                     "[load_mm_data(simple)] error loading %s data at index=%d",

python/sglang/srt/utils/common.py

@@ -1539,6 +1539,24 @@ def load_audio(
     else:
         raise ValueError(f"Invalid audio format: {audio_file}")
 
+    from sglang.srt.multimodal.audio_from_video import (
+        decode_audio_container,
+        is_audio_container,
+    )
+
+    if isinstance(source, bytes):
+        header = source[:16]
+    else:
+        with open(source, "rb") as audio_stream:
+            header = audio_stream.read(16)
+
+    if is_audio_container(header):
+        return decode_audio_container(
+            source,
+            target_sr=sr,
+            mono=mono,
+        )
+
     if _BACKEND == "torchcodec":
         from torchcodec.decoders import AudioDecoder
 
@@ -1740,6 +1758,20 @@ def _normalize_video_input(
         return None
 
 
+def get_video_bytes(video_file: Union[str, bytes, VideoData]) -> bytes:
+    """Normalize a video input and return its encoded bytes."""
+    if isinstance(video_file, VideoData):
+        video_file = video_file.url
+
+    source = _normalize_video_input(video_file)
+    if isinstance(source, bytes):
+        return source
+    if isinstance(source, str):
+        with open(source, "rb") as f:
+            return f.read()
+    raise ValueError(f"Unsupported video input type: {type(video_file)}")
+
+
 def load_video(video_file: Union[str, bytes, VideoData], use_gpu: bool = True):
     if isinstance(video_file, VideoData):
         # preprocess_kwargs is consumed by the multimodal processor, not here.

python/sglang/srt/utils/profile_utils.py

@@ -359,14 +359,17 @@ class _ProfilerTorch(_ProfilerConcreteBase):
 
 class _ProfilerMemory(_ProfilerConcreteBase):
     def start(self):
-        torch.cuda.memory._record_memory_history(max_entries=100000)
+        torch.cuda.memory._record_memory_history(
+            max_entries=envs.SGLANG_MEM_PROFILE_MAX_ENTRIES.get()
+        )
 
     def stop(self):
         Path(self.output_dir).mkdir(parents=True, exist_ok=True)
 
         memory_profile_path = os.path.join(
             self.output_dir,
-            str(time.time())
+            (self.output_prefix + "-" if self.output_prefix else "")
+            + str(time.time())
             + f"-TP-{self.ps.tp_rank}-memory"
             + self.output_suffix
             + ".pickle",

test/registered/unit/multimodal/test_audio_container_decode.py

@@ -0,0 +1,271 @@
+"""Regression tests for explicit ``input_audio`` media containers."""
+
+import asyncio
+import base64
+import concurrent.futures
+import io
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import av
+import numpy as np
+import soundfile
+
+from sglang.srt.managers.schedule_batch import Modality
+from sglang.srt.multimodal.audio_from_video import (
+    decode_audio_container,
+    extract_audio_from_video_bytes,
+    is_audio_container,
+)
+from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
+from sglang.srt.utils.common import load_audio
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=15, suite="base-a-test-cpu")
+
+
+def _tone(*, sample_rate: int, channels: int = 1) -> np.ndarray:
+    samples = np.arange(sample_rate // 10)
+    mono = 0.2 * np.sin(2 * np.pi * 440 * samples / sample_rate)
+    return np.repeat(mono[np.newaxis, :], channels, axis=0).astype(np.float32)
+
+
+def _encode_audio_container(
+    *,
+    codec: str,
+    container_format: str,
+    sample_rate: int,
+    channels: int = 1,
+) -> bytes:
+    output = io.BytesIO()
+    layout = "mono" if channels == 1 else "stereo"
+    sample_format = "s16" if codec == "libopencore_amrnb" else "fltp"
+    samples = _tone(sample_rate=sample_rate, channels=channels)
+    if sample_format == "s16":
+        samples = (samples * np.iinfo(np.int16).max).astype(np.int16)
+
+    with av.open(output, mode="w", format=container_format) as container:
+        stream = container.add_stream(codec, rate=sample_rate)
+        stream.layout = layout
+        frame = av.AudioFrame.from_ndarray(
+            samples,
+            format=sample_format,
+            layout=layout,
+        )
+        frame.sample_rate = sample_rate
+        for packet in stream.encode(frame):
+            container.mux(packet)
+        for packet in stream.encode(None):
+            container.mux(packet)
+    return output.getvalue()
+
+
+def _encode_video_only_mp4() -> bytes:
+    output = io.BytesIO()
+    with av.open(output, mode="w", format="mp4") as container:
+        stream = container.add_stream("mpeg4", rate=1)
+        stream.width = 16
+        stream.height = 16
+        stream.pix_fmt = "yuv420p"
+        frame = av.VideoFrame.from_ndarray(
+            np.zeros((16, 16, 3), dtype=np.uint8),
+            format="rgb24",
+        )
+        for packet in stream.encode(frame):
+            container.mux(packet)
+        for packet in stream.encode(None):
+            container.mux(packet)
+    return output.getvalue()
+
+
+def _encode_video_only_webm() -> bytes:
+    output = io.BytesIO()
+    with av.open(output, mode="w", format="webm") as container:
+        stream = container.add_stream("libvpx", rate=1)
+        stream.width = 16
+        stream.height = 16
+        stream.pix_fmt = "yuv420p"
+        frame = av.VideoFrame.from_ndarray(
+            np.zeros((16, 16, 3), dtype=np.uint8),
+            format="rgb24",
+        )
+        for packet in stream.encode(frame):
+            container.mux(packet)
+        for packet in stream.encode(None):
+            container.mux(packet)
+    return output.getvalue()
+
+
+def _encode_wav() -> bytes:
+    output = io.BytesIO()
+    soundfile.write(output, _tone(sample_rate=16000)[0], 16000, format="WAV")
+    return output.getvalue()
+
+
+class _StubProcessor(BaseMultimodalProcessor):
+    async def process_mm_data_async(
+        self,
+        image_data,
+        audio_data,
+        input_text,
+        request_obj,
+        **kwargs,
+    ):
+        raise NotImplementedError
+
+
+class TestAudioContainerDetection(CustomTestCase):
+    def test_supported_container_signatures(self):
+        """Guards the external container signatures used for decoder routing."""
+        supported_headers = (
+            b"\x00\x00\x00\x18ftypisom",
+            b"\x00\x00\x00\x18ftypM4A ",
+            b"RIFF\x00\x00\x00\x00AVI ",
+            b"#!AMR\n",
+            b"#!AMR-WB\n",
+            b"\x1a\x45\xdf\xa3\x01\x00\x00\x00",
+        )
+        for header in supported_headers:
+            with self.subTest(header=header):
+                self.assertTrue(is_audio_container(header))
+
+    def test_non_container_audio_is_not_routed(self):
+        """Keeps WAV and unknown short inputs on their existing decoder path."""
+        self.assertFalse(is_audio_container(b"RIFF\x00\x00\x00\x00WAVE"))
+        self.assertFalse(is_audio_container(b"random"))
+
+
+class TestAudioContainerDecode(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.mp4 = _encode_audio_container(
+            codec="aac",
+            container_format="mp4",
+            sample_rate=16000,
+            channels=2,
+        )
+        cls.amr = _encode_audio_container(
+            codec="libopencore_amrnb",
+            container_format="amr",
+            sample_rate=8000,
+        )
+        cls.webm = _encode_audio_container(
+            codec="libopus",
+            container_format="webm",
+            sample_rate=48000,
+        )
+
+    def test_mp4_aac_decodes_and_resamples_to_mono(self):
+        """Reproduces the AAC-in-MP4 input that libsndfile rejected."""
+        waveform = decode_audio_container(
+            self.mp4,
+            target_sr=8000,
+            mono=True,
+        )
+        self.assertEqual(waveform.ndim, 1)
+        self.assertEqual(waveform.dtype, np.float32)
+        self.assertTrue(waveform.flags.c_contiguous)
+        self.assertGreater(len(waveform), 700)
+
+    def test_mp4_bypasses_torchcodec_backend(self):
+        """Recognized containers must use PyAV regardless of video backend."""
+        with patch("sglang.srt.utils.common._BACKEND", "torchcodec"):
+            waveform = load_audio(self.mp4, sr=16000)
+        self.assertGreater(len(waveform), 0)
+
+    def test_wav_keeps_existing_decoder_path(self):
+        """Ordinary WAV must not be redirected to the container decoder."""
+        with (
+            patch("sglang.srt.utils.common._BACKEND", "decord"),
+            patch(
+                "sglang.srt.multimodal.audio_from_video.decode_audio_container"
+            ) as decoder,
+        ):
+            waveform = load_audio(_encode_wav(), sr=16000)
+        decoder.assert_not_called()
+        self.assertEqual(waveform.shape, (1600,))
+
+    def test_amr_decodes_despite_wav_data_url_mime(self):
+        """Reproduces AMR bytes mislabeled as WAV by production clients."""
+        data_url = "data:audio/wav;base64," + base64.b64encode(self.amr).decode()
+        waveform = load_audio(data_url, sr=16000)
+        self.assertEqual(waveform.ndim, 1)
+        self.assertGreater(len(waveform), 1500)
+
+    def test_webm_opus_decodes_despite_mp4_data_url_mime(self):
+        """Reproduces WebM bytes mislabeled as audio/mp4 by production clients."""
+        data_url = "data:audio/mp4;base64," + base64.b64encode(self.webm).decode()
+        waveform = load_audio(data_url, sr=16000)
+        self.assertEqual(waveform.ndim, 1)
+        self.assertGreater(len(waveform), 1500)
+
+    def test_video_only_webm_is_a_value_error(self):
+        """A video-only WebM sent as input_audio must be a clear client error."""
+        with self.assertRaisesRegex(ValueError, "Invalid input_audio"):
+            load_audio(_encode_video_only_webm(), sr=16000)
+
+    def test_recognized_path_is_passed_directly_to_pyav(self):
+        """Prevents path inputs from regressing to a full in-memory read."""
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "audio.m4a"
+            path.write_bytes(self.mp4)
+            for source in (str(path), path.as_uri()):
+                with (
+                    self.subTest(source=source),
+                    patch(
+                        "sglang.srt.multimodal.audio_from_video.decode_audio_container",
+                        wraps=decode_audio_container,
+                    ) as decoder,
+                ):
+                    waveform = load_audio(source, sr=16000)
+                    self.assertGreater(len(waveform), 0)
+                    self.assertEqual(decoder.call_args.args[0], str(path))
+
+    def test_invalid_container_remains_a_value_error(self):
+        """Malformed recognized media must remain a client input error."""
+        with self.assertRaisesRegex(ValueError, "Invalid input_audio"):
+            load_audio(b"\x00\x00\x00\x18ftypisom-invalid", sr=16000)
+
+    def test_container_without_audio_is_a_value_error(self):
+        """A valid video-only MP4 must not be accepted as empty audio."""
+        with self.assertRaisesRegex(ValueError, "Invalid input_audio"):
+            decode_audio_container(
+                _encode_video_only_mp4(),
+                target_sr=16000,
+                mono=True,
+            )
+
+    def test_invalid_container_is_classified_as_bad_input(self):
+        """Invalid input_audio must not become an internal server error."""
+        with self.assertRaisesRegex(ValueError, "Invalid input_audio"):
+            _StubProcessor._load_single_item(
+                b"\x00\x00\x00\x18ftypisom-invalid",
+                Modality.AUDIO,
+                audio_sample_rate=16000,
+            )
+
+    def test_fast_loader_preserves_bad_input_error(self):
+        """Fast multimodal loading must preserve client input errors."""
+        processor = _StubProcessor.__new__(_StubProcessor)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+            processor.io_executor = executor
+            with self.assertRaisesRegex(ValueError, "Invalid input_audio"):
+                asyncio.run(
+                    processor.fast_load_mm_data(
+                        prompt="",
+                        multimodal_tokens=None,
+                        audio_data=[b"\x00\x00\x00\x18ftypisom-invalid"],
+                        audio_sample_rate=16000,
+                    )
+                )
+
+    def test_lenient_video_wrapper_returns_none(self):
+        """Silent or corrupt video remains optional for video-specific callers."""
+        self.assertIsNone(extract_audio_from_video_bytes(b"invalid"))
+
+
+if __name__ == "__main__":
+    unittest.main()