From f8197e120f8361b3e951f60e95b3ea0ac938928e Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Fri, 18 Sep 2026 09:27:56 +0800 Subject: [PATCH] [5/N] Schedule chunk transfers as their own stage (#7406) Signed-off-by: specture724 --- .../minimax_h3/test_minimax_h3_contract.py | 8 +- .../models/wan2_2/test_wan22_chunked_mp4.py | 4 +- .../models/wan2_2/test_wan22_s2v_pipeline.py | 7 +- .../diffusion/utils/test_chunked_transfer.py | 261 ++++++++++++ vllm_omni/diffusion/utils/chunked_video.py | 394 ++++++++++++++++-- vllm_omni/diffusion/utils/media_utils.py | 83 +++- 6 files changed, 702 insertions(+), 55 deletions(-) create mode 100644 tests/diffusion/utils/test_chunked_transfer.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 7d2d8715f..4e261b3e2 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -2859,8 +2859,10 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): self.kwargs = kwargs self.__class__.instances.append(self) - def push(self, frames): + def push(self, frames, *, on_consumed=None): self.pushes.append(np.array(frames, copy=True)) + if on_consumed is not None: + on_consumed() def finish(self): return b"mp4" @@ -2986,8 +2988,10 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc self.kwargs = kwargs self.__class__.instances.append(self) - def push(self, frames): + def push(self, frames, *, on_consumed=None): del frames + if on_consumed is not None: + on_consumed() def finish(self): return b"mp4" diff --git a/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py b/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py index a919d1727..2855bf6fb 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py +++ b/tests/diffusion/models/wan2_2/test_wan22_chunked_mp4.py @@ -138,13 +138,13 @@ def test_request_batch_threshold_controls_transfers_without_losing_frames(monkey from vllm_omni.diffusion.utils import chunked_video transfers = [] - quantize = chunked_video.chunk_to_uint8_frames + quantize = chunked_video.quantize_chunk def capture(chunk, value_range): transfers.append(chunk.shape[2]) return quantize(chunk, value_range) - monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture) + monkeypatch.setattr(chunked_video, "quantize_chunk", capture) params = SimpleNamespace(extra_args={"preencode_mp4": True, "preencode_batch_frames": value}) vae = _FakeWanVAE(batch=1, chunks=4, frames_per_chunk=3, height=16, width=32) videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=resolve_wan_preencode_batch_frames(params)) diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index ab8cd132c..a36f3a6a5 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock, patch import numpy as np @@ -43,7 +44,7 @@ def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | def _make_s2v_sampling(**overrides): - values: dict[str, object] = { + values: dict[str, Any] = { "height": 16, "width": 16, "num_frames": 8, @@ -749,13 +750,13 @@ def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch from vllm_omni.diffusion.utils import chunked_video transfers = [] - quantize = chunked_video.chunk_to_uint8_frames + quantize = chunked_video.quantize_chunk def capture(chunk, value_range): transfers.append(chunk.shape[2]) return quantize(chunk, value_range) - monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture) + monkeypatch.setattr(chunked_video, "quantize_chunk", capture) with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform: platform.is_available.return_value = False diff --git a/tests/diffusion/utils/test_chunked_transfer.py b/tests/diffusion/utils/test_chunked_transfer.py new file mode 100644 index 000000000..1a8c52eb5 --- /dev/null +++ b/tests/diffusion/utils/test_chunked_transfer.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Ownership and backpressure contract for chunk transfers to the host. + +The ring hands a chunk one reusable host slot for its whole trip. Getting the +release point wrong is silent: a slot reused while its copy is still landing, +or while an encoder still holds a view of it, corrupts frames rather than +raising. These cover the release point, the backpressure it creates, and the +escape path out of a wait. +""" + +from __future__ import annotations + +import threading + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.utils.chunked_video import ( + ChunkedVideoMP4Session, + ChunkLease, + PinnedChunkRing, + quantize_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _frames(value: int = 7, *, batch: int = 1) -> torch.Tensor: + """One quantized ``BTHWC`` chunk, the shape the ring transfers.""" + return torch.full((batch, 2, 4, 6, 3), value, dtype=torch.uint8) + + +def _cpu_ring(depth: int) -> PinnedChunkRing: + return PinnedChunkRing(depth=depth, device=torch.device("cpu")) + + +def test_slot_returns_only_after_its_last_reader(): + """Each batch entry reads the same buffer, so one release must not free it.""" + ring = _cpu_ring(2) + lease = ring.transfer(_frames(batch=3), readers=3) + + assert ring.slots_in_use == 1 + view = lease.wait() + assert view.shape == (3, 2, 4, 6, 3) + # Waiting is not consuming: an encoder still holds a view of this buffer. + assert ring.slots_in_use == 1 + + lease.release() + lease.release() + assert ring.slots_in_use == 1 + + lease.release() + assert ring.slots_in_use == 0 + # A late extra release must not free the slot a second time. + lease.release() + assert ring.slots_in_use == 0 + + +def test_transfer_waits_for_a_slot_instead_of_growing_the_pool(): + """Depth is the bound on in-flight host buffers, so a full ring blocks.""" + ring = _cpu_ring(1) + held = ring.transfer(_frames(1), readers=1) + second: list[object] = [] + + def take_a_slot() -> None: + second.append(ring.transfer(_frames(2), readers=1)) + + waiter = threading.Thread(target=take_a_slot) + waiter.start() + waiter.join(timeout=0.3) + assert waiter.is_alive(), "a full ring must not hand out a second slot" + + held.wait() + held.release() + waiter.join(timeout=5) + assert not waiter.is_alive() + assert ring.slots_in_use == 1 + assert ring.stats.slot_wait_seconds > 0 + + +def test_abort_wakes_a_transfer_waiting_for_a_slot(): + """A dead consumer must surface as an error, not as a stuck producer.""" + ring = _cpu_ring(1) + ring.transfer(_frames(), readers=1) + failures: list[BaseException] = [] + + def take_a_slot() -> None: + try: + ring.transfer(_frames(), readers=1) + except BaseException as exc: # noqa: BLE001 + failures.append(exc) + + waiter = threading.Thread(target=take_a_slot) + waiter.start() + waiter.join(timeout=0.3) + assert waiter.is_alive() + + ring.abort(RuntimeError("encoder died")) + waiter.join(timeout=5) + assert [str(exc) for exc in failures] == ["encoder died"] + # Later transfers stay failed rather than silently resuming. + with pytest.raises(RuntimeError, match="encoder died"): + ring.transfer(_frames(), readers=1) + + +def test_encoder_releases_the_chunk_it_was_reading_when_muxing_fails(monkeypatch): + """A failed mux abandons the frame generator, which still owns a lent slot.""" + from vllm_omni.diffusion.utils import media_utils + + released: list[int] = [] + + def exploding_mux(frames, **kwargs): + # Read into the chunk, so the generator is suspended inside it. + next(iter(frames)) + raise RuntimeError("mux failed") + + monkeypatch.setattr(media_utils, "mux_av_video_audio_bytes", exploding_mux) + encoder = media_utils.ChunkedMP4Encoder(width=6, height=4, fps=24) + encoder.push(np.zeros((2, 4, 6, 3), dtype=np.uint8), on_consumed=lambda: released.append(1)) + + with pytest.raises(RuntimeError, match="mux failed"): + encoder.finish() + + assert released == [1], "the abandoned generator must not keep holding its chunk" + + +class _RecordingEncoder: + """Stands in for one batch entry's encoder, honouring the release contract.""" + + instances: list[_RecordingEncoder] = [] + + def __init__(self, **kwargs): + self.pushes: list[int] = [] + self.released = 0 + self.fail_on_push = False + self.__class__.instances.append(self) + + def push(self, frames, *, on_consumed=None): + if self.fail_on_push: + if on_consumed is not None: + on_consumed() + self.released += 1 + raise RuntimeError("encoder rejected the chunk") + self.pushes.append(int(frames[0, 0, 0, 0])) + if on_consumed is not None: + on_consumed() + self.released += 1 + + def finish(self) -> bytes: + return b"mp4" + + def abort(self) -> None: + return None + + +@pytest.fixture +def recording_encoders(monkeypatch): + _RecordingEncoder.instances = [] + monkeypatch.setattr( + "vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", + _RecordingEncoder, + ) + return _RecordingEncoder.instances + + +def _push_ramp(session: ChunkedVideoMP4Session, values: list[int]) -> None: + for value in values: + session.push(torch.full((1, 3, 1, 4, 6), value / 255.0)) + + +def test_session_feeds_chunks_in_producer_order(recording_encoders): + """Copies complete in issue order, and the encoder must see that order.""" + session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24) + _push_ramp(session, [10, 20, 30, 40]) + + assert session.finish() == [b"mp4"] + assert recording_encoders[0].pushes == [10, 20, 30, 40] + + +def test_session_returns_the_slot_when_an_encoder_rejects_a_chunk(recording_encoders): + """A rejected chunk must not strand the buffer it was lent.""" + session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, transfer_slots=1) + session.push(torch.zeros(1, 3, 1, 4, 6)) + recording_encoders[0].fail_on_push = True + + with pytest.raises(RuntimeError, match="rejected"): + session.push(torch.zeros(1, 3, 1, 4, 6)) + + assert session.ring is not None + assert session.ring.slots_in_use == 0 + + +def test_session_abort_releases_every_slot_still_in_flight(recording_encoders): + """Teardown returns the pool even when nothing consumed the chunks.""" + session = ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, transfer_slots=2) + session.push(torch.zeros(1, 3, 1, 4, 6)) + + session.abort() + + assert session.ring is not None + assert session.ring.slots_in_use == 0 + + +def test_session_gets_its_slot_back_after_a_real_encoder_failure(): + """A real libx264 error must not strand the producer on the slot it lent.""" + session = ChunkedVideoMP4Session( + value_range=(0.0, 1.0), + fps=24, + transfer_slots=1, + video_codec_options={"preset": "invalid-preset"}, + ) + failures: list[BaseException] = [] + + def produce() -> None: + try: + for _ in range(4): + session.push(torch.zeros(1, 3, 1, 16, 16)) + session.finish() + except BaseException as exc: # noqa: BLE001 + failures.append(exc) + + producer = threading.Thread(target=produce, daemon=True) + producer.start() + producer.join(timeout=30) + # With one slot, a lease the failed encoder never released leaves the + # producer waiting on the ring forever instead of seeing the codec error. + assert not producer.is_alive(), "the producer stalled on a slot the failed encoder kept" + assert len(failures) == 1 + + session.abort() + assert session.ring is not None + assert session.ring.slots_in_use == 0 + + +@hardware_test(res={"cuda": ["H100", "B200"]}) +def test_device_chunks_land_intact_through_the_pinned_ring(): + """On an accelerator the copy is deferred, so its event must order the read.""" + if not torch.accelerator.is_available(): + pytest.skip("no accelerator available") + + device = torch.device(torch.accelerator.current_accelerator().type, 0) + ring = PinnedChunkRing(depth=2, device=device) + pending: list[tuple[int, ChunkLease]] = [] + for step in range(6): + value = 10 * step + 1 + # Quantize on the compute stream immediately before handing the chunk + # over: without the source-ready event the copy would race this work. + chunk = torch.full((1, 3, 2, 4, 6), value / 255.0, device=device) + pending.append((value, ring.transfer(quantize_chunk(chunk, (0.0, 1.0)), readers=1))) + if len(pending) == 2: + for expected, lease in pending: + assert int(lease.wait()[0, 0, 0, 0, 0]) == expected + lease.release() + pending.clear() + + assert ring.stats.transfers == 6 + assert ring.stats.peak_slots_in_use == 2 + assert ring.stats.peak_pending_bytes == 2 * 2 * 4 * 6 * 3 diff --git a/vllm_omni/diffusion/utils/chunked_video.py b/vllm_omni/diffusion/utils/chunked_video.py index 79ceeab8c..57b6d2f82 100644 --- a/vllm_omni/diffusion/utils/chunked_video.py +++ b/vllm_omni/diffusion/utils/chunked_video.py @@ -10,41 +10,293 @@ and lands in a bounded encoder. That consumer lives here so a model only has to publish finished chunks -- :class:`SupportsChunkedVAEDecode` for producers that can be driven, :class:`ChunkedVideoMP4Session` for producers that drive themselves. + +The move to the host runs as its own scheduled stage. A chunk is quantized on +the accelerator, copied without blocking into one of a fixed number of reusable +pinned host slots on a dedicated stream, and only read once its copy-completion +event has fired. The producer therefore returns to decoding while the copy is +still in flight, and the slot count bounds how many chunks can be in flight at +once. """ from __future__ import annotations +import threading +import time +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass from typing import Any import numpy as np import torch +from vllm.logger import init_logger from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder +logger = init_logger(__name__) -def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray: - """Quantize a ``BCTHW`` chunk to ``BTHWC`` uint8, transferring once. +# Re-checked while waiting for a slot so an abort recorded without a notify +# still wakes the producer. +_SLOT_POLL_SECONDS = 0.05 + +_COPY_STREAMS: dict[tuple[str, int], torch.Stream] = {} +_COPY_STREAM_LOCK = threading.Lock() + + +def _copy_stream(device: torch.device) -> torch.Stream: + """Return the process-wide chunk-copy stream for ``device``. + + One stream per device rather than one per request: copies are issued in + order on it anyway, and a stream per request would churn accelerator + resources for no ordering benefit. + """ + with _COPY_STREAM_LOCK: + key = (device.type, device.index or 0) + stream = _COPY_STREAMS.get(key) + if stream is None: + stream = torch.Stream(device=device) + _COPY_STREAMS[key] = stream + return stream + + +def quantize_chunk(chunk: torch.Tensor, value_range: tuple[float, float]) -> torch.Tensor: + """Quantize a ``BCTHW`` chunk to contiguous ``BTHWC`` uint8 on its own device. ``value_range`` is the interval the producer publishes, which differs per checkpoint (see ``SupportsChunkedVAEDecode.chunk_value_range``). Quantizing - on the accelerator means one transfer moves the final bytes rather than - float frames. - - That transfer is synchronous: ``.cpu()`` blocks inside the producer's chunk - callback, so what overlaps the remaining decode today is the CPU H.264 - encode of earlier chunks, not this copy. Splitting device preparation, - device-to-host copy, and encode into three independently scheduled stages -- - a pool of reusable pinned host slots with non-blocking copies on a dedicated - stream, leases released by the encoder -- fits behind this same callback - contract, but is not implemented here. + on the accelerator means the transfer that follows moves the final bytes + rather than float frames, and making the result contiguous there keeps the + copy a single flat move instead of a strided host-side gather. """ low, high = float(value_range[0]), float(value_range[1]) if high <= low: raise ValueError(f"value_range must be increasing, got {value_range!r}") scale = 255.0 / (high - low) frames = chunk.clamp(low, high).sub(low).mul(scale).round().to(torch.uint8) - return frames.permute(0, 2, 3, 4, 1).cpu().numpy() + return frames.permute(0, 2, 3, 4, 1).contiguous() + + +def chunk_to_uint8_frames(chunk: torch.Tensor, value_range: tuple[float, float]) -> np.ndarray: + """Quantize a ``BCTHW`` chunk and block until it reaches the host. + + The path a session takes when no pinned ring is available, and the one a + caller outside a session gets. A session on an accelerator uses + :class:`PinnedChunkRing` instead so the copy does not stall the producer. + """ + return quantize_chunk(chunk, value_range).cpu().numpy() + + +@dataclass +class ChunkTransferStats: + """What the host transfer cost, for checking that it actually overlapped. + + ``copy_wait_seconds`` staying near zero while ``transfers`` climbs is the + signal that copies completed behind the producer; it rising with them means + the producer is outrunning the copy stream. ``slot_wait_seconds`` rising + instead means the encoder is the bottleneck and backpressure has reached + the producer. + """ + + transfers: int = 0 + slot_wait_seconds: float = 0.0 + copy_wait_seconds: float = 0.0 + peak_slots_in_use: int = 0 + peak_pending_bytes: int = 0 + + +class _HostSlot: + """One reusable pinned host buffer. + + The buffer is allocated on first use and never shrinks. ``batch_frames`` is + a flush threshold rather than a cap, so a later batch can be larger than + the first; the slot then keeps the larger buffer for the rest of the + request instead of reallocating per chunk. + """ + + def __init__(self, *, pinned: bool) -> None: + self._pinned = pinned + self._buffer: torch.Tensor | None = None + self.held_bytes = 0 + + def reserve(self, frames: torch.Tensor) -> torch.Tensor: + """Return a host view shaped like ``frames``, growing the buffer if needed.""" + needed = frames.numel() + if self._buffer is None or self._buffer.numel() < needed: + try: + self._buffer = torch.empty(needed, dtype=frames.dtype, pin_memory=self._pinned) + except Exception: + # Pinning can be refused by the driver or unsupported on a + # platform. A pageable buffer still copies correctly, only + # synchronously, so degrade rather than fail the request. + logger.warning("Pinned host buffer unavailable; chunk copies fall back to pageable memory") + self._pinned = False + self._buffer = torch.empty(needed, dtype=frames.dtype) + self.held_bytes = needed * frames.element_size() + return self._buffer[:needed].view(frames.shape) + + +class ChunkLease: + """One chunk in flight: a host slot, its copy event, and its readers. + + The slot returns to the ring only after the last reader has finished with + it. Being handed to an encoder is not enough -- the encoder's queue holds a + view of this buffer until its muxing thread reads it -- so release is + driven by the encoder rather than by dequeuing. + """ + + def __init__( + self, + *, + ring: PinnedChunkRing, + slot: _HostSlot, + view: np.ndarray, + source: torch.Tensor, + landed: torch.Event | None, + readers: int, + ) -> None: + self._ring = ring + self._slot = slot + self._view = view + self._source: torch.Tensor | None = source + self._landed = landed + self._readers = readers + self._lock = threading.Lock() + + def ready(self) -> bool: + """Whether the copy has landed, without blocking on it.""" + return self._landed is None or self._landed.query() + + def wait(self) -> np.ndarray: + """Block until the copy has landed, then hand over the host frames.""" + if self._landed is not None: + started = time.perf_counter() + self._landed.synchronize() + self._ring.stats.copy_wait_seconds += time.perf_counter() - started + self._landed = None + # The source stays referenced until here so the allocator cannot hand + # its device memory to the next chunk while the copy is still reading. + self._source = None + return self._view + + def release(self) -> None: + """Report one reader done; the last one returns the slot to the ring.""" + with self._lock: + if self._readers <= 0: + return + self._readers -= 1 + if self._readers > 0: + return + self._ring.reclaim(self._slot) + + def discard(self) -> None: + """Drop an unread chunk during teardown, keeping the copy's memory alive. + + The copy may still be writing into the slot and reading from the source + tensor, so this waits for it rather than letting either be reused + underneath an in-flight transfer. + """ + if self._landed is not None: + self._landed.synchronize() + self._landed = None + self._source = None + with self._lock: + remaining, self._readers = self._readers, 0 + if remaining > 0: + self._ring.reclaim(self._slot) + + +class PinnedChunkRing: + """Fixed-depth pool of reusable host slots for chunk transfers. + + A slot walks ``free -> copying -> ready -> encoding -> free`` and holds + exactly one chunk for that whole trip, so pending transfer bytes are + bounded by the depth times the largest chunk of the request. Waiting for a + free slot is where encoder backpressure reaches the producer. + """ + + def __init__(self, *, depth: int, device: torch.device) -> None: + if depth <= 0: + raise ValueError("transfer ring depth must be positive") + self.depth = depth + self.stats = ChunkTransferStats() + self._on_device = device.type != "cpu" + self._stream = _copy_stream(device) if self._on_device else None + self._lock = threading.Lock() + self._slot_freed = threading.Condition(self._lock) + self._free = [_HostSlot(pinned=self._on_device) for _ in range(depth)] + self._pending_bytes = 0 + self._failure: BaseException | None = None + + @property + def slots_in_use(self) -> int: + """Slots currently holding a chunk, whether copying, ready or encoding.""" + with self._lock: + return self.depth - len(self._free) + + def transfer(self, frames: torch.Tensor, *, readers: int) -> ChunkLease: + """Start one non-blocking copy of ``frames`` into a free slot.""" + if readers <= 0: + raise ValueError("a transferred chunk needs at least one reader") + slot = self._acquire() + try: + host = slot.reserve(frames) + with self._lock: + self._pending_bytes += slot.held_bytes + self.stats.transfers += 1 + self.stats.peak_pending_bytes = max(self.stats.peak_pending_bytes, self._pending_bytes) + landed = self._issue_copy(host, frames) + except BaseException: + self.reclaim(slot) + raise + return ChunkLease(ring=self, slot=slot, view=host.numpy(), source=frames, landed=landed, readers=readers) + + def abort(self, error: BaseException | None = None) -> None: + """Fail every present and future wait for a slot.""" + with self._slot_freed: + if self._failure is None: + self._failure = error or RuntimeError("chunk transfer ring aborted") + self._slot_freed.notify_all() + + def reclaim(self, slot: _HostSlot) -> None: + """Return a slot whose last reader is done.""" + with self._slot_freed: + self._pending_bytes -= slot.held_bytes + slot.held_bytes = 0 + self._free.append(slot) + self._slot_freed.notify() + + def _issue_copy(self, host: torch.Tensor, frames: torch.Tensor) -> torch.Event | None: + if self._stream is None: + host.copy_(frames) + return None + # The chunk was quantized on the compute stream, so the copy stream has + # to wait for that work before reading it. + source_ready = torch.Event() + source_ready.record() + self._stream.wait_event(source_ready) + with self._stream: + host.copy_(frames, non_blocking=True) + landed = torch.Event() + landed.record(self._stream) + return landed + + def _acquire(self) -> _HostSlot: + started = time.perf_counter() + with self._slot_freed: + while not self._free and self._failure is None: + # ponytail: no request deadline yet, only an abort escape. A + # unified deadline across producer, slot, queue and encoder + # waits is RFC #6872 R4. + self._slot_freed.wait(timeout=_SLOT_POLL_SECONDS) + if self._failure is not None: + raise self._failure + slot = self._free.pop() + in_use = self.depth - len(self._free) + self.stats.peak_slots_in_use = max(self.stats.peak_slots_in_use, in_use) + self.stats.slot_wait_seconds += time.perf_counter() - started + return slot class ChunkedVideoMP4Session: @@ -58,7 +310,8 @@ class ChunkedVideoMP4Session: silent entry), so a caller with several outputs per prompt repeats a request's waveform across its entries. ``batch_frames`` coalesces transfers for producers that publish finer than a transfer is worth; ``crop`` trims - the decoder's padding to the requested output size. + the decoder's padding to the requested output size. ``transfer_slots`` sets + how many chunks may be in flight to the host at once. """ def __init__( @@ -72,6 +325,7 @@ class ChunkedVideoMP4Session: max_pending: int = 2, video_codec_options: dict[str, str] | None = None, crop: tuple[int, int] | None = None, + transfer_slots: int = 2, ) -> None: if batch_frames <= 0: raise ValueError("batch_frames must be positive") @@ -83,9 +337,13 @@ class ChunkedVideoMP4Session: self._max_pending = max_pending self._video_codec_options = video_codec_options self._crop = crop + self._transfer_slots = transfer_slots self._encoders: list[ChunkedMP4Encoder] = [] self._pending: list[torch.Tensor] = [] self._pending_frames = 0 + self._ring: PinnedChunkRing | None = None + self._ring_resolved = False + self._inflight: deque[ChunkLease] = deque() def push(self, chunk: torch.Tensor) -> None: """Queue one committed ``BCTHW`` chunk.""" @@ -100,33 +358,96 @@ class ChunkedVideoMP4Session: def finish(self) -> list[bytes]: """Flush what is pending and return one MP4 per batch entry.""" self._flush() + while self._inflight: + self._feed_lease(self._inflight.popleft()) return [encoder.finish() for encoder in self._encoders] def abort(self) -> None: + if self._ring is not None: + self._ring.abort() + while self._inflight: + self._inflight.popleft().discard() for encoder in self._encoders: encoder.abort() + def transfer_stats(self) -> ChunkTransferStats | None: + """Per-request transfer accounting, or ``None`` if no ring was used.""" + return None if self._ring is None else self._ring.stats + + @property + def ring(self) -> PinnedChunkRing | None: + """The transfer ring this session settled on, once a chunk has flushed.""" + return self._ring + def _flush(self) -> None: if not self._pending: return - frames = chunk_to_uint8_frames(torch.cat(self._pending, dim=2), self._value_range) - if not self._encoders: - self._encoders = [ - ChunkedMP4Encoder( - width=frames.shape[3], - height=frames.shape[2], - fps=self._fps, - audio_waveform=self._waveform_for(index, frames.shape[0]), - audio_sample_rate=self._audio_sample_rate, - max_pending=self._max_pending, - video_codec_options=self._video_codec_options, - ) - for index in range(frames.shape[0]) - ] - for index, encoder in enumerate(self._encoders): - encoder.push(np.ascontiguousarray(frames[index])) + frames = quantize_chunk(torch.cat(self._pending, dim=2), self._value_range) self._pending.clear() self._pending_frames = 0 + self._ensure_encoders(frames) + ring = self._resolve_ring(frames.device) + if ring is None: + self._feed(frames.cpu().numpy(), release=None) + return + # Free a slot before asking for one, so the ring's own wait is only + # ever for the encoder to catch up rather than for this loop. + while len(self._inflight) >= ring.depth: + self._feed_lease(self._inflight.popleft()) + self._inflight.append(ring.transfer(frames, readers=len(self._encoders))) + while self._inflight and self._inflight[0].ready(): + self._feed_lease(self._inflight.popleft()) + + def _feed_lease(self, lease: ChunkLease) -> None: + try: + self._feed(lease.wait(), release=lease.release) + except BaseException: + lease.discard() + raise + + def _feed(self, frames: np.ndarray, *, release: Callable[[], None] | None) -> None: + for index, encoder in enumerate(self._encoders): + try: + encoder.push(frames[index], on_consumed=release) + except BaseException: + # push() owns the callback for the entry it was given, so only + # the encoders this chunk never reached still hold a reader. + for _ in range(len(self._encoders) - index - 1): + if release is not None: + release() + raise + + def _ensure_encoders(self, frames: torch.Tensor) -> None: + if self._encoders: + return + self._encoders = [ + ChunkedMP4Encoder( + width=int(frames.shape[3]), + height=int(frames.shape[2]), + fps=self._fps, + audio_waveform=self._waveform_for(index, int(frames.shape[0])), + audio_sample_rate=self._audio_sample_rate, + max_pending=self._max_pending, + video_codec_options=self._video_codec_options, + ) + for index in range(int(frames.shape[0])) + ] + + def _resolve_ring(self, device: torch.device) -> PinnedChunkRing | None: + if self._ring_resolved: + return self._ring + self._ring_resolved = True + if self._transfer_slots <= 0: + return None + try: + self._ring = PinnedChunkRing(depth=self._transfer_slots, device=device) + except Exception: + logger.warning( + "Pinned chunk transfers unavailable on %s; falling back to blocking copies", + device, + exc_info=True, + ) + return self._ring def _waveform_for(self, index: int, batch_size: int) -> np.ndarray | None: if self._audio_waveforms is None: @@ -154,7 +475,18 @@ def decode_to_mp4(vae: Any, z: torch.Tensor, **session_kwargs: Any) -> list[byte session = ChunkedVideoMP4Session(value_range=vae.chunk_value_range, **session_kwargs) try: vae.decode_with_chunks(z, on_chunk=session.push) - return session.finish() + videos = session.finish() except BaseException: session.abort() raise + stats = session.transfer_stats() + if stats is not None: + logger.debug( + "Chunked MP4 transfers=%d slot_wait=%.3fs copy_wait=%.3fs peak_slots=%d peak_bytes=%d", + stats.transfers, + stats.slot_wait_seconds, + stats.copy_wait_seconds, + stats.peak_slots_in_use, + stats.peak_pending_bytes, + ) + return videos diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index b2141400c..a53359cc8 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -7,16 +7,36 @@ from __future__ import annotations import io import queue import threading -from collections.abc import Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping from fractions import Fraction -from typing import Any, cast +from typing import Any, NamedTuple, cast import av import numpy as np +from vllm.logger import init_logger + +logger = init_logger(__name__) _CHUNKED_MP4_DONE = object() +class _QueuedChunk(NamedTuple): + """One queued chunk plus the callback that returns its buffer to its owner.""" + + frames: np.ndarray + on_consumed: Callable[[], None] | None + + +def _release_chunk(on_consumed: Callable[[], None] | None) -> None: + """Hand a chunk's buffer back, never letting the callback wedge the worker.""" + if on_consumed is None: + return + try: + on_consumed() + except BaseException: # noqa: BLE001 + logger.exception("Chunked MP4 encoder failed to release a consumed chunk") + + def normalize_preencode_batch_frames(value: Any) -> int: """Validate the request's MP4 transfer/encoding batch threshold.""" if isinstance(value, bool) or not isinstance(value, int) or value <= 0: @@ -88,9 +108,10 @@ class ChunkedMP4Encoder: self._state_lock = threading.Lock() def run() -> None: + frames = self._frames() try: self._result = mux_av_video_audio_bytes( - self._frames(), + frames, width=self.width, height=self.height, audio_waveform=audio_waveform, @@ -103,6 +124,11 @@ class ChunkedMP4Encoder: ) except BaseException as exc: self._error = exc + # Closing is what runs the generator's cleanup and releases the + # chunk it was mid-way through. Waiting for collection would + # never get there: the traceback just stored on this encoder + # keeps the abandoned generator alive for the process's life. + frames.close() self._drain_until_done() self._thread = threading.Thread(target=run, name="chunked-mp4", daemon=True) @@ -110,19 +136,29 @@ class ChunkedMP4Encoder: def _drain_until_done(self) -> None: # A flush/mux failure can occur after _frames consumed the sentinel. - if not self._input_done: - while self._queue.get() is not _CHUNKED_MP4_DONE: - pass + if self._input_done: + return + while True: + entry = self._queue.get() + if entry is _CHUNKED_MP4_DONE: + return + _release_chunk(cast(_QueuedChunk, entry).on_consumed) def _frames(self): while True: - chunk = self._queue.get() - if chunk is _CHUNKED_MP4_DONE: + entry = self._queue.get() + if entry is _CHUNKED_MP4_DONE: self._input_done = True return - assert isinstance(chunk, np.ndarray) - for frame_data in chunk: - yield av.VideoFrame.from_ndarray(frame_data, format="rgb24") + assert isinstance(entry, _QueuedChunk) + try: + for frame_data in entry.frames: + yield av.VideoFrame.from_ndarray(frame_data, format="rgb24") + finally: + # Also runs when the muxer raises part-way through this chunk + # and abandons the generator, so a lent buffer is never + # stranded in the encoder. + _release_chunk(entry.on_consumed) def _validate_chunk(self, chunk: np.ndarray) -> None: _validate_video_chunk(chunk, width=self.width, height=self.height) @@ -133,11 +169,22 @@ class ChunkedMP4Encoder: if self._closed: raise RuntimeError("ChunkedMP4Encoder is already closed") - def push(self, chunk: np.ndarray) -> None: - """Queue a chunk; asynchronous failures surface here or at finish().""" - self._validate_chunk(chunk) - self._raise_if_failed() - self._queue.put(chunk) + def push(self, chunk: np.ndarray, *, on_consumed: Callable[[], None] | None = None) -> None: + """Queue a chunk; asynchronous failures surface here or at finish(). + + ``on_consumed`` fires exactly once after this encoder has read + ``chunk``, including when the chunk is discarded by an abort or by a + failed encode. Ownership transfers on call: whether this returns or + raises, the callback is this encoder's responsibility, so a caller + lending a pooled host buffer never compensates on the error path. + """ + try: + self._validate_chunk(chunk) + self._raise_if_failed() + except BaseException: + _release_chunk(on_consumed) + raise + self._queue.put(_QueuedChunk(chunk, on_consumed)) self._raise_if_failed() def _send_done(self) -> None: @@ -168,9 +215,11 @@ class ChunkedMP4Encoder: self._aborted = True while True: try: - self._queue.get_nowait() + entry = self._queue.get_nowait() except queue.Empty: break + if entry is not _CHUNKED_MP4_DONE: + _release_chunk(cast(_QueuedChunk, entry).on_consumed) self._send_done() self._thread.join()