diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index ab51abf93..3eee51ed1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -16,6 +16,7 @@ import shutil import subprocess import tempfile import threading +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from copy import copy from dataclasses import dataclass, field @@ -48,6 +49,8 @@ from sglang.srt.observability.trace import TraceReqContext logger = init_logger(__name__) _MAX_CACHED_CUDA_VIDEO_BUFFER_BYTES = 1024 * 1024 * 1024 +_MAX_CUDA_VIDEO_CONVERSION_CHUNK_BYTES = 128 * 1024 * 1024 +_MAX_PARALLEL_CUDA_VIDEO_SAVES = 2 _cuda_video_buffer_cache_lock = threading.Lock() _cached_cuda_video_buffer: "_CudaMemfdVideoBuffer | None" = None @@ -464,6 +467,25 @@ def _x264_auto_thread_count(height: int) -> int: return min(cpu_limit, row_limit, 128) +def _cuda_video_conversion_chunk_frames(video: torch.Tensor) -> int: + _, num_frames, height, width = video.shape + temporary_bytes_per_frame = 3 * height * width * (video.element_size() + 2) + return min( + num_frames, + max(1, _MAX_CUDA_VIDEO_CONVERSION_CHUNK_BYTES // temporary_bytes_per_frame), + ) + + +def _sendfile_all(output_fd: int, input_fd: int, count: int) -> None: + offset = 0 + while count: + sent = os.sendfile(output_fd, input_fd, offset, count) + if sent <= 0: + raise RuntimeError("sendfile made no progress") + offset += sent + count -= sent + + def _try_save_cuda_video_direct( *, save_file_path: str, @@ -472,8 +494,8 @@ def _try_save_cuda_video_direct( audio_sample_rate: Optional[int], output_compression: Optional[int], ) -> bool: - """Save a CUDA RGB video through a registered memfd instead of a Python pipe.""" - if not hasattr(os, "memfd_create") or not os.path.isdir("/proc/self/fd"): + """Stream CUDA RGB chunks to ffmpeg through a registered memfd.""" + if not hasattr(os, "memfd_create") or not hasattr(os, "sendfile"): return False sample_without_audio, audio = _split_sample_audio(sample) @@ -492,9 +514,8 @@ def _try_save_cuda_video_direct( if video.shape[0] != 3: return False - frames = (video * 255).clamp_(0, 255).to(torch.uint8) - frames = frames.permute(1, 2, 3, 0).contiguous() - num_frames, height, width, _ = frames.shape + _, num_frames, height, width = video.shape + chunk_frames = _cuda_video_conversion_chunk_frames(video) quality = output_compression / 10 if output_compression is not None else 5 if not 1 <= quality <= 10: @@ -518,76 +539,105 @@ def _try_save_cuda_video_direct( scipy_wavfile.write(tmp_wav_path, selected_sr, audio_np) ffmpeg_exe = _resolve_ffmpeg_exe() - shape = tuple(frames.shape) - with _acquire_cuda_video_buffer(shape) as buffer: - assert buffer.tensor is not None - buffer.tensor.copy_(frames, non_blocking=True) - torch.cuda.current_stream(frames.device).synchronize() - del frames - os.lseek(buffer.fd, 0, os.SEEK_SET) + command = [ + ffmpeg_exe, + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{width}x{height}", + "-pix_fmt", + "rgb24", + "-r", + f"{fps:.02f}", + "-i", + "pipe:0", + ] + if tmp_wav_path is None: + command += ["-an"] + else: + command += ["-i", tmp_wav_path] + command += [ + "-vcodec", + "libx264", + "-pix_fmt", + "yuv420p", + "-crf", + str(crf), + ] - command = [ - ffmpeg_exe, - "-y", - "-f", - "rawvideo", - "-vcodec", - "rawvideo", - "-s", - f"{width}x{height}", - "-pix_fmt", - "rgb24", - "-r", - f"{fps:.02f}", - "-i", - f"/proc/self/fd/{buffer.fd}", - ] - if tmp_wav_path is None: - command += ["-an"] - else: - command += ["-i", tmp_wav_path] - command += [ - "-vcodec", - "libx264", - "-pix_fmt", - "yuv420p", - "-crf", - str(crf), - ] - - macro_block_size = 16 - if width % macro_block_size or height % macro_block_size: - output_width = ( - width - if width % macro_block_size == 0 - else width + macro_block_size - width % macro_block_size - ) - output_height = ( - height - if height % macro_block_size == 0 - else height + macro_block_size - height % macro_block_size - ) - command += ["-vf", f"scale={output_width}:{output_height}"] - - command += ["-threads", str(_x264_auto_thread_count(height))] - if tmp_wav_path is not None: - command += [ - "-acodec", - "aac", - "-map", - "0:v:0", - "-map", - "1:a:0", - ] - command += ["-v", "warning", save_file_path] - - subprocess.run( - command, - check=True, - pass_fds=(buffer.fd,), - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, + macro_block_size = 16 + if width % macro_block_size or height % macro_block_size: + output_width = ( + width + if width % macro_block_size == 0 + else width + macro_block_size - width % macro_block_size ) + output_height = ( + height + if height % macro_block_size == 0 + else height + macro_block_size - height % macro_block_size + ) + command += ["-vf", f"scale={output_width}:{output_height}"] + + command += ["-threads", str(_x264_auto_thread_count(height))] + if tmp_wav_path is not None: + command += [ + "-acodec", + "aac", + "-map", + "0:v:0", + "-map", + "1:a:0", + ] + command += ["-v", "warning", save_file_path] + + with tempfile.TemporaryFile() as stderr_file: + process = subprocess.Popen( + command, + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=stderr_file, + ) + try: + if process.stdin is None: + raise RuntimeError("ffmpeg stdin pipe was not created") + with _acquire_cuda_video_buffer( + (chunk_frames, height, width, 3) + ) as buffer: + assert buffer.tensor is not None + for start in range(0, num_frames, chunk_frames): + end = min(start + chunk_frames, num_frames) + frames = ( + (video[:, start:end] * 255).clamp_(0, 255).to(torch.uint8) + ) + frames = frames.permute(1, 2, 3, 0).contiguous() + buffer.tensor[: end - start].copy_(frames, non_blocking=True) + torch.cuda.current_stream(video.device).synchronize() + del frames + _sendfile_all( + process.stdin.fileno(), + buffer.fd, + (end - start) * height * width * 3, + ) + process.stdin.close() + process.stdin = None + returncode = process.wait() + finally: + if process.stdin is not None: + process.stdin.close() + if process.poll() is None: + process.kill() + process.wait() + if returncode: + stderr_file.seek(0) + raise subprocess.CalledProcessError( + returncode, + command, + stderr=stderr_file.read(), + ) return True except Exception as e: logger.warning( @@ -603,6 +653,86 @@ def _try_save_cuda_video_direct( pass +def _try_save_cuda_videos_direct( + samples: Sequence[Any], + save_file_paths: Sequence[str], + *, + fps: int, + audio_sample_rate: Optional[int], + output_compression: Optional[int], +) -> list[bool] | None: + """Save independent CUDA videos concurrently when memory permits.""" + if len(samples) < 2 or len(samples) != len(save_file_paths): + return None + + videos = [] + for sample, save_file_path in zip(samples, save_file_paths): + video, _ = _split_sample_audio(sample) + if not ( + isinstance(video, torch.Tensor) + and video.device.type == "cuda" + and video.dim() in (3, 4) + and int(video.shape[0]) == 3 + and os.path.splitext(save_file_path)[1].lower() == ".mp4" + ): + return None + if video.dim() == 3: + video = video.unsqueeze(1) + if videos and video.device != videos[0].device: + return None + videos.append(video) + + # Each direct save creates one multiply result and two uint8 layouts for a + # temporal chunk. Keep two such chunks below 25% of currently free device + # memory so postprocessing cannot turn a tight inference into an OOM. + temporary_bytes = sorted( + ( + _cuda_video_conversion_chunk_frames(video) + * 3 + * int(video.shape[-2]) + * int(video.shape[-1]) + * (int(video.element_size()) + 2) + for video in videos + ), + reverse=True, + )[:_MAX_PARALLEL_CUDA_VIDEO_SAVES] + try: + free_bytes, _ = torch.cuda.mem_get_info(videos[0].device) + except (RuntimeError, TypeError): + return None + if sum(temporary_bytes) > int(free_bytes) // 4: + return None + try: + available_cpus = len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + available_cpus = os.cpu_count() or 1 + encoder_threads = sorted( + (_x264_auto_thread_count(int(video.shape[-2])) for video in videos), + reverse=True, + )[:_MAX_PARALLEL_CUDA_VIDEO_SAVES] + if sum(encoder_threads) > available_cpus: + return None + + def save_one(idx: int) -> bool: + return _try_save_cuda_video_direct( + save_file_path=save_file_paths[idx], + sample=samples[idx], + fps=fps, + audio_sample_rate=audio_sample_rate, + output_compression=output_compression, + ) + + try: + with ThreadPoolExecutor(max_workers=_MAX_PARALLEL_CUDA_VIDEO_SAVES) as pool: + return list(pool.map(save_one, range(len(samples)))) + except Exception as exc: + logger.warning( + "Parallel CUDA video save failed; falling back to serial output: %s", + str(exc), + ) + return None + + def _mux_audio_np_into_mp4( *, save_file_path: str, @@ -1013,8 +1143,35 @@ def save_outputs( upscaling_scale: int = 4, ) -> list[str]: output_paths: list[str] = [] - for idx, sample in enumerate(outputs): - save_file_path = build_output_path(idx) + samples = ( + [ + attach_audio_to_video_sample(sample, audio, idx) + for idx, sample in enumerate(outputs) + ] + if data_type == DataType.VIDEO + else outputs + ) + save_file_paths = [build_output_path(idx) for idx in range(len(outputs))] + parallel_results = None + if ( + data_type == DataType.VIDEO + and len(outputs) > 1 + and save_output + and frames_out is None + and not enable_frame_interpolation + and not enable_upscaling + ): + for path in save_file_paths: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + parallel_results = _try_save_cuda_videos_direct( + samples, + save_file_paths, + fps=fps, + audio_sample_rate=audio_sample_rate, + output_compression=output_compression, + ) + + for idx, (sample, save_file_path) in enumerate(zip(samples, save_file_paths)): if data_type == DataType.ACTION: if samples_out is not None: samples_out.append(sample) @@ -1031,7 +1188,6 @@ def save_outputs( continue if data_type == DataType.VIDEO: - sample = attach_audio_to_video_sample(sample, audio, idx) if ( save_output and save_file_path @@ -1040,13 +1196,16 @@ def save_outputs( and not enable_upscaling ): os.makedirs(os.path.dirname(save_file_path) or ".", exist_ok=True) - if _try_save_cuda_video_direct( - save_file_path=save_file_path, - sample=sample, - fps=fps, - audio_sample_rate=audio_sample_rate, - output_compression=output_compression, - ): + direct_saved = parallel_results is not None and parallel_results[idx] + if not direct_saved: + direct_saved = _try_save_cuda_video_direct( + save_file_path=save_file_path, + sample=sample, + fps=fps, + audio_sample_rate=audio_sample_rate, + output_compression=output_compression, + ) + if direct_saved: if samples_out is not None: samples_out.append(sample) if audios_out is not None: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py index a51c69099..32dfd26d1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py @@ -6,6 +6,7 @@ from __future__ import annotations import json import math import subprocess +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( @@ -329,16 +330,22 @@ class MiniMaxH3VideoModelAdapter: if shape.get("width") is not None and shape.get("height") is not None: expected_size = (int(shape["width"]), int(shape["height"])) - final_media_fields: dict[str, str] = {} - for output_index, output_path in enumerate(output_paths): - media_fields = _probe_minimax_h3_output_fields( + def probe_output(output_path: str) -> dict[str, str]: + return _probe_minimax_h3_output_fields( output_path, expected_frame_count=expected_frame_count, expected_size=expected_size, ) - if output_index == 0: - final_media_fields = media_fields - elif media_fields != final_media_fields: + + if len(output_paths) > 1: + with ThreadPoolExecutor(max_workers=min(4, len(output_paths))) as pool: + media_fields_by_output = list(pool.map(probe_output, output_paths)) + else: + media_fields_by_output = [probe_output(output_paths[0])] + + final_media_fields = media_fields_by_output[0] + for output_index, media_fields in enumerate(media_fields_by_output[1:], 1): + if media_fields != final_media_fields: raise RuntimeError( "generated MiniMax H3 outputs have inconsistent media metadata: " f"output 0={final_media_fields}, output " diff --git a/python/sglang/multimodal_gen/test/unit/test_output_saving.py b/python/sglang/multimodal_gen/test/unit/test_output_saving.py index 5157816a6..c4de2acaf 100644 --- a/python/sglang/multimodal_gen/test/unit/test_output_saving.py +++ b/python/sglang/multimodal_gen/test/unit/test_output_saving.py @@ -199,3 +199,47 @@ def test_video_direct_save_short_circuits_materialization(tmp_path, monkeypatch) assert paths == [str(output_path)] assert len(direct_calls) == 1 + + +def test_multiple_videos_use_parallel_direct_save_with_serial_fallback( + tmp_path, monkeypatch +): + outputs = [torch.zeros((3, 1, 2, 3)), torch.ones((3, 1, 2, 3))] + direct_calls = [] + + def parallel_save(samples, paths, **kwargs): + direct_calls.append((samples, paths, kwargs)) + return [True, False] + + serial_calls = [] + + monkeypatch.setattr(output_utils, "_try_save_cuda_videos_direct", parallel_save) + monkeypatch.setattr( + output_utils, + "_try_save_cuda_video_direct", + lambda **kwargs: serial_calls.append(kwargs) or True, + ) + monkeypatch.setattr( + output_utils, + "post_process_sample", + lambda *_args, **_kwargs: pytest.fail( + "successful parallel direct saves should skip frame materialization" + ), + ) + + paths = output_utils.save_outputs( + outputs, + DataType.VIDEO, + fps=24, + save_output=True, + build_output_path=lambda idx: str(tmp_path / f"sample_{idx}.mp4"), + ) + + assert paths == [str(tmp_path / "sample_0.mp4"), str(tmp_path / "sample_1.mp4")] + assert len(direct_calls) == 1 + samples, save_paths, kwargs = direct_calls[0] + assert all(actual is expected for actual, expected in zip(samples, outputs)) + assert save_paths == paths + assert kwargs["fps"] == 24 + assert len(serial_calls) == 1 + assert serial_calls[0]["save_file_path"] == paths[1]