vllm-omni/tests/diffusion/test_diffusion_ipc.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import contextlib
import multiprocessing as mp
import queue
import threading
import time
from multiprocessing import shared_memory

import numpy as np
import pytest
import torch
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue

from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
from vllm_omni.diffusion.ipc import (
    _SHM_TENSOR_THRESHOLD,
    DIFFUSION_RPC_RESULT_ENVELOPE,
    _pack_value_if_large,
    _unpack_if_shm_handle,
    pack_diffusion_output_shm,
    unpack_diffusion_output_shm,
)
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


def _large_numel(dtype: torch.dtype) -> int:
    return (_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=dtype).element_size()) + 1


def _cleanup_shm_handle(value: object) -> None:
    if isinstance(value, dict) and value.get("__tensor_shm__"):
        with contextlib.suppress(FileNotFoundError):
            _unpack_if_shm_handle(value)


def _result_queue_worker(connection, rank: int) -> None:
    """Send one nested NumPy result through a worker-owned MessageQueue."""
    result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
    try:
        connection.send(result_mq.export_handle())
        assert connection.recv() == "reader-ready"

        frames = np.full(300_000, rank, dtype=np.float32)
        audio = np.arange(300_000, dtype=np.float32) + rank
        output = DiffusionOutput(
            output={
                "rank": rank,
                "nested": {"frames": frames},
                "audio": [audio],
            }
        )
        pack_diffusion_output_shm(output)
        connection.send(
            [
                output.output["nested"]["frames"]["name"],
                output.output["audio"][0]["name"],
            ]
        )
        result_mq.enqueue(output)
        assert connection.recv() == "shutdown"
    finally:
        result_mq.shutdown()
        connection.close()


def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> None:
    """Exercise the production per-worker queue/pump/SHM lifecycle."""
    ctx = mp.get_context("spawn")
    parent_connections = []
    processes = []
    result_mqs = []
    shm_names = []
    executor = None

    try:
        for rank in range(2):
            parent_connection, child_connection = ctx.Pipe()
            process = ctx.Process(
                target=_result_queue_worker,
                args=(child_connection, rank),
                name=f"IPCResultWorker-{rank}",
            )
            process.start()
            child_connection.close()
            parent_connections.append(parent_connection)
            processes.append(process)

        for connection in parent_connections:
            handle = connection.recv()
            result_mqs.append(MessageQueue.create_from_handle(handle, 0))
            connection.send("reader-ready")
            shm_names.extend(connection.recv())

        executor = object.__new__(MultiprocDiffusionExecutor)
        executor._result_mqs = result_mqs
        executor._result_mq = result_mqs[0]
        executor._pump_running = False
        executor._pump_stop = threading.Event()
        executor._sync_result_buffer = queue.Queue()
        executor._is_failed = False
        executor._futures_lock = threading.RLock()
        executor._rpc_futures = {}
        executor._output_futures = {}
        executor._completed_outputs = {}
        executor._batch_split_map = {}
        executor._start_result_pump()

        deadline = time.monotonic() + 10.0
        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
            time.sleep(0.01)
        assert executor._sync_result_buffer.qsize() == 2

        outputs = []
        for _ in range(2):
            output = executor._sync_result_buffer.get_nowait()
            unpack_diffusion_output_shm(output)
            outputs.append(output)

        assert sorted(output.output["rank"] for output in outputs) == [0, 1]
        for output in outputs:
            rank = output.output["rank"]
            assert output.output["nested"]["frames"][0] == rank
            assert output.output["audio"][0][0] == rank

        executor._pump_stop.set()
        for thread in executor._result_pump_threads:
            thread.join(timeout=3.0)
        assert all(not thread.is_alive() for thread in executor._result_pump_threads)
    finally:
        if executor is not None:
            executor._pump_stop.set()
            for thread in getattr(executor, "_result_pump_threads", []):
                thread.join(timeout=3.0)

        for connection in parent_connections:
            with contextlib.suppress(BrokenPipeError, EOFError, OSError):
                connection.send("shutdown")
        for process in processes:
            process.join(timeout=10.0)
            if process.is_alive():
                process.terminate()
                process.join(timeout=5.0)
        for result_mq in result_mqs:
            with contextlib.suppress(Exception):
                result_mq.shutdown()
        for connection in parent_connections:
            connection.close()

    assert all(process.exitcode == 0 for process in processes)
    for name in shm_names:
        try:
            leaked = shared_memory.SharedMemory(name=name)
        except FileNotFoundError:
            continue
        leaked.close()
        pytest.fail(f"shared memory segment {name} still exists after unpack")


def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
    image = torch.arange(300_000, dtype=torch.float32)
    video = torch.arange(300_000, dtype=torch.float32) * 2
    output = DiffusionOutput(output={"image": image, "video": video, "metadata": {"keep": "inline"}})

    pack_diffusion_output_shm(output)

    assert output.output["image"]["__tensor_shm__"] is True
    assert output.output["video"]["__tensor_shm__"] is True
    assert output.output["metadata"] == {"keep": "inline"}

    unpack_diffusion_output_shm(output)

    torch.testing.assert_close(output.output["image"], image)
    torch.testing.assert_close(output.output["video"], video)
    assert output.output["metadata"] == {"keep": "inline"}


def test_diffusion_output_tuple_tensors_round_trip_through_shm() -> None:
    # LTX2 / DreamID return (video, audio) tuples as DiffusionOutput.output.
    video = torch.arange(300_000, dtype=torch.float32)
    audio = torch.arange(300_000, dtype=torch.float32) * 3
    output = DiffusionOutput(output=(video, audio))

    pack_diffusion_output_shm(output)

    assert isinstance(output.output, tuple)
    assert output.output[0]["__tensor_shm__"] is True
    assert output.output[1]["__tensor_shm__"] is True

    unpack_diffusion_output_shm(output)

    assert isinstance(output.output, tuple)
    torch.testing.assert_close(output.output[0], video)
    torch.testing.assert_close(output.output[1], audio)


def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
    frames = [torch.arange(300_000, dtype=torch.float32), torch.arange(300_000, dtype=torch.float32) + 1]
    output = DiffusionOutput(output=list(frames))

    pack_diffusion_output_shm(output)

    assert isinstance(output.output, list)
    assert all(isinstance(item, dict) and item["__tensor_shm__"] is True for item in output.output)

    unpack_diffusion_output_shm(output)

    assert isinstance(output.output, list)
    torch.testing.assert_close(output.output[0], frames[0])
    torch.testing.assert_close(output.output[1], frames[1])


def test_diffusion_output_numpy_array_round_trips_through_shm() -> None:
    frames = np.arange(300_000, dtype=np.float32)
    output = DiffusionOutput(output=frames)

    pack_diffusion_output_shm(output)

    assert output.output["__ndarray_shm__"] is True

    unpack_diffusion_output_shm(output)

    assert isinstance(output.output, np.ndarray)
    np.testing.assert_array_equal(output.output, frames)


def test_numpy_object_array_stays_inline() -> None:
    values = np.empty(200_000, dtype=object)
    values[:] = "safe-to-pickle"

    packed = _pack_value_if_large(values)

    assert packed is values


def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
    tensor = torch.arange(300_000, dtype=torch.float32)
    envelope = {
        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
        "result": DiffusionOutput(output=tensor),
        "rank_statuses": [{"rank": 0, "ok": True}],
    }

    packed = pack_diffusion_output_shm(envelope)

    assert packed is envelope
    result = packed["result"]
    assert isinstance(result, DiffusionOutput)
    assert result.output["__tensor_shm__"] is True
    assert packed["rank_statuses"] == [{"rank": 0, "ok": True}]

    unpacked = unpack_diffusion_output_shm(packed)

    assert unpacked is envelope
    result = unpacked["result"]
    assert isinstance(result, DiffusionOutput)
    torch.testing.assert_close(result.output, tensor)
    assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]


def test_rpc_result_envelope_dp_tagged_output_round_trips_through_shm() -> None:
    frames = np.arange(300_000, dtype=np.float32)
    envelope = {
        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
        "result": {"dp_rank": 1, "output": DiffusionOutput(output=frames)},
        "rank_statuses": [{"rank": 1, "ok": True}],
    }

    packed = pack_diffusion_output_shm(envelope)
    tagged = packed["result"]
    assert tagged["dp_rank"] == 1
    assert tagged["output"].output["__ndarray_shm__"] is True

    unpacked = unpack_diffusion_output_shm(packed)
    tagged = unpacked["result"]
    assert tagged["dp_rank"] == 1
    np.testing.assert_array_equal(tagged["output"].output, frames)


def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
    first = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
    second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1
    output = BatchRunnerOutput.from_list(
        [
            RunnerOutput(request_id="req-0", finished=True, result=DiffusionOutput(output=first)),
            RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output={"image": second})),
            RunnerOutput(request_id="req-error", finished=True, result=DiffusionOutput(error="boom")),
        ]
    )

    pack_diffusion_output_shm(output)

    assert output.runner_outputs[0].result.output["__tensor_shm__"] is True
    assert output.runner_outputs[1].result.output["image"]["__tensor_shm__"] is True
    assert output.runner_outputs[2].result.error == "boom"

    unpack_diffusion_output_shm(output)

    torch.testing.assert_close(output["req-0"].result.output, first)
    torch.testing.assert_close(output["req-1"].result.output["image"], second)
    assert output["req-error"].result.error == "boom"


def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
    tensor = torch.arange(
        _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
        dtype=torch.float32,
    )

    packed = _pack_value_if_large(tensor)

    assert packed is tensor


def test_pack_value_packs_large_tensor_and_round_trips() -> None:
    tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
    packed = _pack_value_if_large(tensor)

    try:
        assert isinstance(packed, dict)
        assert packed["__tensor_shm__"] is True
        assert packed["shape"] == [tensor.numel()]
        assert packed["torch_dtype"] == "torch.float32"

        unpacked = _unpack_if_shm_handle(packed)
        assert isinstance(unpacked, torch.Tensor)
        torch.testing.assert_close(unpacked, tensor)
    finally:
        _cleanup_shm_handle(packed)


def test_pack_value_recurses_nested_dicts_and_lists_without_mutating_inline_values() -> None:
    large = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
    small = torch.arange(8, dtype=torch.float32)
    list_tensor = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
    payload = {
        "media": {
            "large": large,
            "small": small,
        },
        "list_value": [list_tensor],
        "metadata": {"prompt": "keep inline"},
    }

    packed = _pack_value_if_large(payload)

    try:
        assert packed is not payload
        assert packed["media"] is not payload["media"]
        assert packed["media"]["large"]["__tensor_shm__"] is True
        assert packed["media"]["small"] is small
        # Lists are recursed too: the large tensor inside is packed and a new
        # list is returned, while the input payload is left untouched.
        assert packed["list_value"] is not payload["list_value"]
        assert packed["list_value"][0]["__tensor_shm__"] is True
        assert payload["list_value"][0] is list_tensor
        assert packed["metadata"] == {"prompt": "keep inline"}

        torch.testing.assert_close(_unpack_if_shm_handle(packed["media"]["large"]), large)
        torch.testing.assert_close(_unpack_if_shm_handle(packed["list_value"][0]), list_tensor)
    finally:
        if isinstance(packed, dict):
            _cleanup_shm_handle(packed.get("media", {}).get("large"))
            list_value = packed.get("list_value")
            if isinstance(list_value, list) and list_value:
                _cleanup_shm_handle(list_value[0])


def test_pack_value_preserves_dtype_shape_and_values_for_bfloat16() -> None:
    tensor = torch.arange(_large_numel(torch.bfloat16), dtype=torch.float32).to(torch.bfloat16).reshape(1, -1)
    packed = _pack_value_if_large(tensor)

    try:
        assert isinstance(packed, dict)
        assert packed["__tensor_shm__"] is True
        assert packed["shape"] == list(tensor.shape)
        assert packed["torch_dtype"] == "torch.bfloat16"
        assert packed["numpy_dtype"] == "float32"

        unpacked = _unpack_if_shm_handle(packed)
        assert isinstance(unpacked, torch.Tensor)
        assert unpacked.dtype == torch.bfloat16
        torch.testing.assert_close(unpacked, tensor)
    finally:
        _cleanup_shm_handle(packed)


def test_pack_value_packs_non_contiguous_large_tensor_values() -> None:
    tensor = torch.arange(_large_numel(torch.float32) * 2, dtype=torch.float32).reshape(-1, 2)[:, 0]
    assert not tensor.is_contiguous()

    packed = _pack_value_if_large(tensor)

    try:
        assert isinstance(packed, dict)
        assert packed["__tensor_shm__"] is True
        assert packed["shape"] == list(tensor.shape)

        unpacked = _unpack_if_shm_handle(packed)
        assert isinstance(unpacked, torch.Tensor)
        torch.testing.assert_close(unpacked, tensor)
    finally:
        _cleanup_shm_handle(packed)