diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 053d2590b..cf7b9136a 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -335,6 +335,13 @@ window and output frame count stay unchanged. Wan S2V keeps its existing per-clip behavior by default. Zero, negative, fractional, boolean, string, and null values are rejected when pre-encoding is enabled. +The transfer ring applies backpressure before another D2H copy when both of +its two slots are occupied or the combined pending uint8 payload would exceed +256 MiB. This byte bound keeps large resolutions and uneven final chunks from +turning a fixed item count into unexpectedly large memory growth. A native VAE +chunk larger than 256 MiB is admitted only when the ring is otherwise empty, +so supported chunk shapes cannot deadlock. + `preencode_mp4` applies to the complete-MP4 response paths only. The `/v1/realtime/video` WebSocket endpoint rejects it, because that path already overlaps encoding through its own incremental fragmented-MP4 encoder. Wan also diff --git a/tests/diffusion/utils/test_chunked_transfer.py b/tests/diffusion/utils/test_chunked_transfer.py index 1a8c52eb5..a62ca8612 100644 --- a/tests/diffusion/utils/test_chunked_transfer.py +++ b/tests/diffusion/utils/test_chunked_transfer.py @@ -33,8 +33,8 @@ def _frames(value: int = 7, *, batch: int = 1) -> torch.Tensor: 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 _cpu_ring(depth: int, *, max_pending_bytes: int | None = None) -> PinnedChunkRing: + return PinnedChunkRing(depth=depth, device=torch.device("cpu"), max_pending_bytes=max_pending_bytes) def test_slot_returns_only_after_its_last_reader(): @@ -81,6 +81,44 @@ def test_transfer_waits_for_a_slot_instead_of_growing_the_pool(): assert ring.stats.slot_wait_seconds > 0 +def test_transfer_waits_for_byte_budget_even_when_a_slot_is_free(): + """Variable-sized chunks cannot bypass the memory bound via free slots.""" + small = _frames() + large = _frames(batch=2) + large_bytes = large.numel() * large.element_size() + ring = _cpu_ring(3, max_pending_bytes=large_bytes) + held = ring.transfer(small, readers=1) + second: list[ChunkLease] = [] + + waiter = threading.Thread(target=lambda: second.append(ring.transfer(large, readers=1))) + waiter.start() + waiter.join(timeout=0.3) + assert waiter.is_alive(), "the byte budget must backpressure a producer even with free slots" + + held.wait() + held.release() + waiter.join(timeout=5) + assert not waiter.is_alive() + assert ring.slots_in_use == 1 + assert ring.pending_bytes == large_bytes + second[0].wait() + second[0].release() + + +def test_one_oversized_chunk_is_admitted_without_deadlock(): + """A native VAE chunk larger than the budget must make progress alone.""" + chunk = _frames() + chunk_bytes = chunk.numel() * chunk.element_size() + ring = _cpu_ring(2, max_pending_bytes=chunk_bytes - 1) + + lease = ring.transfer(chunk, readers=1) + + assert ring.pending_bytes == chunk_bytes + assert ring.stats.peak_pending_bytes == chunk_bytes + lease.wait() + lease.release() + + 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) @@ -180,6 +218,26 @@ def test_session_feeds_chunks_in_producer_order(recording_encoders): assert recording_encoders[0].pushes == [10, 20, 30, 40] +def test_session_drains_oldest_chunk_when_byte_budget_fills(recording_encoders, monkeypatch): + """Byte pressure must drain queued D2H work instead of deadlocking it.""" + monkeypatch.setattr(ChunkLease, "ready", lambda self: False) + chunk_bytes = 1 * 1 * 4 * 6 * 3 # one quantized BTHWC chunk from _push_ramp + session = ChunkedVideoMP4Session( + value_range=(0.0, 1.0), + fps=24, + transfer_slots=2, + max_pending_bytes=chunk_bytes, + ) + + _push_ramp(session, [10, 20, 30]) + + assert session.finish() == [b"mp4"] + assert recording_encoders[0].pushes == [10, 20, 30] + assert session.ring is not None + assert session.ring.stats.peak_slots_in_use == 1 + assert session.ring.stats.peak_pending_bytes == chunk_bytes + + 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) diff --git a/tests/diffusion/utils/test_chunked_video.py b/tests/diffusion/utils/test_chunked_video.py index 849b49290..80b1c49c5 100644 --- a/tests/diffusion/utils/test_chunked_video.py +++ b/tests/diffusion/utils/test_chunked_video.py @@ -116,3 +116,9 @@ def test_session_rejects_a_waveform_count_that_does_not_match_the_batch(): def test_session_rejects_a_non_positive_batch(): with pytest.raises(ValueError, match="batch_frames"): ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, batch_frames=0) + + +@pytest.mark.parametrize("value", [0, -1, 1.5, True]) +def test_session_rejects_an_invalid_pending_byte_budget(value): + with pytest.raises(ValueError, match="max_pending_bytes"): + ChunkedVideoMP4Session(value_range=(0.0, 1.0), fps=24, max_pending_bytes=value) diff --git a/vllm_omni/diffusion/utils/chunked_video.py b/vllm_omni/diffusion/utils/chunked_video.py index 57b6d2f82..bf88552cf 100644 --- a/vllm_omni/diffusion/utils/chunked_video.py +++ b/vllm_omni/diffusion/utils/chunked_video.py @@ -41,6 +41,11 @@ logger = init_logger(__name__) # still wakes the producer. _SLOT_POLL_SECONDS = 0.05 +# Bound queued/in-flight uint8 video independently of resolution and chunk +# length. A single larger chunk is still admitted when the ring is empty so a +# legal native VAE chunk can never deadlock the producer. +_DEFAULT_MAX_PENDING_BYTES = 256 * 1024**2 + _COPY_STREAMS: dict[tuple[str, int], torch.Stream] = {} _COPY_STREAM_LOCK = threading.Lock() @@ -211,15 +216,24 @@ 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. + exactly one chunk for that whole trip. Admission is bounded by both slot + count and pending bytes; reaching either limit is where encoder + backpressure reaches the producer. """ - def __init__(self, *, depth: int, device: torch.device) -> None: + def __init__( + self, + *, + depth: int, + device: torch.device, + max_pending_bytes: int | None = _DEFAULT_MAX_PENDING_BYTES, + ) -> None: if depth <= 0: raise ValueError("transfer ring depth must be positive") + if max_pending_bytes is not None and (type(max_pending_bytes) is not int or max_pending_bytes <= 0): + raise ValueError("max_pending_bytes must be a positive integer or None") self.depth = depth + self.max_pending_bytes = max_pending_bytes self.stats = ChunkTransferStats() self._on_device = device.type != "cpu" self._stream = _copy_stream(device) if self._on_device else None @@ -235,15 +249,26 @@ class PinnedChunkRing: with self._lock: return self.depth - len(self._free) + @property + def pending_bytes(self) -> int: + """Bytes held by chunks currently copying, ready, or encoding.""" + with self._lock: + return self._pending_bytes + + def would_block(self, chunk_bytes: int) -> bool: + """Whether admitting ``chunk_bytes`` now would exceed either bound.""" + with self._lock: + return not self._free or not self._bytes_available(chunk_bytes) + 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() + chunk_bytes = frames.numel() * frames.element_size() + slot = self._acquire(chunk_bytes) 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) @@ -282,10 +307,17 @@ class PinnedChunkRing: landed.record(self._stream) return landed - def _acquire(self) -> _HostSlot: + def _bytes_available(self, chunk_bytes: int) -> bool: + return ( + self.max_pending_bytes is None + or self._pending_bytes == 0 + or self._pending_bytes + chunk_bytes <= self.max_pending_bytes + ) + + def _acquire(self, chunk_bytes: int) -> _HostSlot: started = time.perf_counter() with self._slot_freed: - while not self._free and self._failure is None: + while (not self._free or not self._bytes_available(chunk_bytes)) 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. @@ -293,6 +325,11 @@ class PinnedChunkRing: if self._failure is not None: raise self._failure slot = self._free.pop() + # Reserve before dropping the lock so concurrent producers cannot + # both pass the byte budget. Setting held_bytes here also makes a + # failed host allocation reclaim the complete reservation. + slot.held_bytes = chunk_bytes + self._pending_bytes += chunk_bytes 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 @@ -311,7 +348,8 @@ class ChunkedVideoMP4Session: 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. ``transfer_slots`` sets - how many chunks may be in flight to the host at once. + how many chunks may be in flight to the host at once, while + ``max_pending_bytes`` bounds their combined uint8 payload size. """ def __init__( @@ -326,9 +364,12 @@ class ChunkedVideoMP4Session: video_codec_options: dict[str, str] | None = None, crop: tuple[int, int] | None = None, transfer_slots: int = 2, + max_pending_bytes: int | None = _DEFAULT_MAX_PENDING_BYTES, ) -> None: if batch_frames <= 0: raise ValueError("batch_frames must be positive") + if max_pending_bytes is not None and (type(max_pending_bytes) is not int or max_pending_bytes <= 0): + raise ValueError("max_pending_bytes must be a positive integer or None") self._value_range = value_range self._fps = fps self._audio_waveforms = audio_waveforms @@ -338,6 +379,7 @@ class ChunkedVideoMP4Session: self._video_codec_options = video_codec_options self._crop = crop self._transfer_slots = transfer_slots + self._max_pending_bytes = max_pending_bytes self._encoders: list[ChunkedMP4Encoder] = [] self._pending: list[torch.Tensor] = [] self._pending_frames = 0 @@ -392,7 +434,8 @@ class ChunkedVideoMP4Session: 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: + chunk_bytes = frames.numel() * frames.element_size() + while self._inflight and ring.would_block(chunk_bytes): self._feed_lease(self._inflight.popleft()) self._inflight.append(ring.transfer(frames, readers=len(self._encoders))) while self._inflight and self._inflight[0].ready(): @@ -440,7 +483,11 @@ class ChunkedVideoMP4Session: if self._transfer_slots <= 0: return None try: - self._ring = PinnedChunkRing(depth=self._transfer_slots, device=device) + self._ring = PinnedChunkRing( + depth=self._transfer_slots, + device=device, + max_pending_bytes=self._max_pending_bytes, + ) except Exception: logger.warning( "Pinned chunk transfers unavailable on %s; falling back to blocking copies",