minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Shiyan Deng
2026-07-24 16:18:45 -07:00
committed by GitHub
co-authored by Lianmin Zheng Lucia Fang
parent 14d6e1d3b1
commit 962c076934
8 changed files with 442 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
+2 -1
View File
@@ -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
+3
View File
@@ -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
+5 -2
View File
@@ -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
+117 -58
View File
@@ -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
python/sglang/srt/multimodal/processors/base_processor.py
+7
View File
@@ -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
+32
View File
@@ -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
+5 -2
View File
@@ -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",