From 962c0769341bd1509f7af6fd84f6df078c50bb93 Mon Sep 17 00:00:00 2001 From: Shiyan Deng Date: Fri, 24 Jul 2026 16:18:45 -0700 Subject: [PATCH] Decode input_audio media containers with PyAV & Update memory profiler (#31832) Signed-off-by: Shiyan Deng Signed-off-by: Lianmin Zheng Co-authored-by: Lianmin Zheng Co-authored-by: Lucia Fang <116399278+luccafong@users.noreply.github.com> --- python/pyproject.toml | 3 +- python/sglang/srt/environ.py | 3 + .../scheduler_components/profiler_manager.py | 7 +- .../sglang/srt/multimodal/audio_from_video.py | 175 +++++++---- .../multimodal/processors/base_processor.py | 7 + python/sglang/srt/utils/common.py | 32 +++ python/sglang/srt/utils/profile_utils.py | 7 +- .../multimodal/test_audio_container_decode.py | 271 ++++++++++++++++++ 8 files changed, 442 insertions(+), 63 deletions(-) create mode 100644 test/registered/unit/multimodal/test_audio_container_decode.py diff --git a/python/pyproject.toml b/python/pyproject.toml index dfe105b7b..d612751ce 100755 --- a/python/pyproject.toml +++ b/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", diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 506ba4cfe..5aea0593d 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index f1c65a3b1..f60d53675 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/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", diff --git a/python/sglang/srt/multimodal/audio_from_video.py b/python/sglang/srt/multimodal/audio_from_video.py index 4ed725c34..df32c95b4 100644 --- a/python/sglang/srt/multimodal/audio_from_video.py +++ b/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). - -PyAV wraps FFmpeg's C libraries in-process, avoiding subprocess forks which -would crash CUDA-active workers. -""" +"""Decode audio from media containers with PyAV.""" 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. - - 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. - """ + """Extract optional mono audio for callers that accept silent videos.""" try: - import av - except ImportError: - logger.warning( - "PyAV (av) is not installed. Cannot extract audio from video. " - "Install with: pip install av" + return decode_audio_container( + video_bytes, + target_sr=target_sr, + mono=True, ) - 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 diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 04c181aaa..65b0f4151 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/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", diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index f99242d3f..852ae9065 100644 --- a/python/sglang/srt/utils/common.py +++ b/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. diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 693a7418d..b3946eeb7 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/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", diff --git a/test/registered/unit/multimodal/test_audio_container_decode.py b/test/registered/unit/multimodal/test_audio_container_decode.py new file mode 100644 index 000000000..f0c34443b --- /dev/null +++ b/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()