vllm-omni/tests/diffusion/utils/test_chunked_transfer.py

# 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, *, 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():
    """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_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)
    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_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)
    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