vllm-omni/tests/diffusion/test_async_output_worker.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for WorkerProc async output logic."""

import queue
import threading
import time
from unittest.mock import MagicMock

import pytest
import torch

import vllm_omni.diffusion.ipc as diffusion_ipc
from vllm_omni.diffusion.data import (
    AsyncDiffusionOutput,
    AsyncOutputKind,
    DiffusionOutput,
    OmniACK,
)
from vllm_omni.diffusion.media import (
    DiffusionMediaOutput,
    VideoMediaOutput,
    VideoTensorEncoding,
    VideoTensorLayout,
    VideoTensorSpec,
    VideoValueRange,
)
from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

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


def _make_worker_proc(step_execution=False):
    """Create a minimal WorkerProc-like object with async output state."""
    proc = object.__new__(WorkerProc)
    proc.od_config = MagicMock()
    proc.od_config.step_execution = step_execution
    proc.gpu_id = 0
    proc.result_mq = MagicMock()
    proc._running = True
    proc._async_output_queue = queue.Queue() if not step_execution else None
    proc._async_output_thread = None
    proc._async_output_done = threading.Condition()
    proc._async_output_pending = 0
    proc._result_mq_lock = threading.Lock()
    return proc


class _OverlapDetectingQueue:
    """Fails if two threads are inside enqueue() at the same time.

    Mirrors vLLM's MessageQueue, whose acquire_write() reads current_idx,
    writes that block, then advances it. Overlapping writers can target the
    same ring block and silently drop a message.
    """

    def __init__(self):
        self._active = 0
        self._guard = threading.Lock()
        self.overlaps = 0
        self.messages = []

    def enqueue(self, msg, timeout=None):
        with self._guard:
            self._active += 1
            if self._active > 1:
                self.overlaps += 1
        time.sleep(0.001)  # widen the window a real ring-buffer write would have
        with self._guard:
            self.messages.append(msg)
            self._active -= 1


class TestResultQueueSingleWriter:
    """Worker must serialize writes to the single-writer result queue."""

    def test_concurrent_enqueues_do_not_overlap(self):
        proc = _make_worker_proc(step_execution=False)
        detector = _OverlapDetectingQueue()
        proc.result_mq = detector

        def writer(tag):
            for i in range(20):
                proc._enqueue_result(f"{tag}-{i}")

        threads = [threading.Thread(target=writer, args=(tag,)) for tag in ("compute", "output")]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert detector.overlaps == 0, f"{detector.overlaps} overlapping writes to result_mq"
        assert len(detector.messages) == 40


class TestGenerateAsyncOutputId:
    """Test _generate_async_output_id uniqueness."""

    def test_returns_hex_string(self):
        id1 = WorkerProc._generate_async_output_id()
        assert isinstance(id1, str)
        assert len(id1) == 32  # uuid4 hex is 32 chars

    def test_ids_are_unique(self):
        ids = {WorkerProc._generate_async_output_id() for _ in range(100)}
        assert len(ids) == 100


class TestReturnResultAsyncPath:
    """Test _return_result in request-mode (step_execution=False)."""

    def test_omni_ack_bypasses_async_path(self):
        """OmniACK should always use sync path regardless of step_execution."""
        proc = _make_worker_proc(step_execution=False)
        ack = OmniACK(task_id="task1", status="done")

        proc._return_result(ack)
        proc.result_mq.enqueue.assert_called_once_with(ack)
        assert proc._async_output_queue.empty()  # nothing enqueued for bg thread

    def test_diffusion_output_triggers_async_path(self, mocker):
        """DiffusionOutput in request-mode should enqueue compute_done + bg work."""
        proc = _make_worker_proc(step_execution=False)

        mock_platform = mocker.MagicMock()
        mock_platform.record_device_event.return_value = MagicMock()
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
            mock_platform,
        )

        output = DiffusionOutput(output="data")
        proc._return_result(output, rpc_id="1")

        # Should have enqueued COMPUTE_DONE to result_mq
        enqueued_msg = proc.result_mq.enqueue.call_args[0][0]
        assert isinstance(enqueued_msg, AsyncDiffusionOutput)
        assert enqueued_msg.kind == AsyncOutputKind.COMPUTE_DONE
        assert enqueued_msg.rpc_id == "1"
        assert enqueued_msg.async_output_id is not None

        # Should have placed work item in async_output_queue
        assert not proc._async_output_queue.empty()
        item = proc._async_output_queue.get_nowait()
        assert item[0] is output
        assert item[1] == enqueued_msg.async_output_id

    def test_batch_runner_output_triggers_async_path(self, mocker):
        """BatchRunnerOutput in request-mode should also go through async path."""
        proc = _make_worker_proc(step_execution=False)

        mock_platform = mocker.MagicMock()
        mock_platform.record_device_event.return_value = MagicMock()
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
            mock_platform,
        )

        batch = BatchRunnerOutput(
            runner_outputs=[RunnerOutput(request_id="r1", finished=True)],
        )
        proc._return_result(batch)

        enqueued_msg = proc.result_mq.enqueue.call_args[0][0]
        assert isinstance(enqueued_msg, AsyncDiffusionOutput)
        assert enqueued_msg.kind == AsyncOutputKind.COMPUTE_DONE

    def test_none_result_mq_skips_all_paths(self):
        """If result_mq is None, _return_result should do nothing."""
        proc = _make_worker_proc(step_execution=False)
        proc.result_mq = None

        proc._return_result(DiffusionOutput())
        # No crash, no enqueue attempted


class TestReturnResultSyncPath:
    """Test _return_result in step-execution mode (sync)."""

    def test_diffusion_output_uses_sync_path_when_step_execution(self, mocker):
        """In step_execution mode, DiffusionOutput goes through sync SHM pack."""
        proc = _make_worker_proc(step_execution=True)
        proc._async_output_queue = None  # step mode doesn't have async queue

        mock_pack = mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
        )
        output = DiffusionOutput(output="data")
        proc._return_result(output)

        mock_pack.assert_called_once_with(output)
        proc.result_mq.enqueue.assert_called_once_with(output)

    def test_media_contract_error_is_not_sent_as_an_unpacked_result(self, mocker):
        proc = _make_worker_proc(step_execution=True)
        proc._async_output_queue = None
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
            side_effect=ValueError("media was not prepared"),
        )

        with pytest.raises(ValueError, match="media was not prepared"):
            proc._return_result(DiffusionOutput())

        proc.result_mq.enqueue.assert_not_called()

    def test_media_packing_memory_failure_is_reraised_for_typed_media(self, mocker):
        """A non-ValueError packing failure on typed media must re-raise, not
        swallow it and enqueue a half-packed payload."""
        proc = _make_worker_proc(step_execution=True)
        proc._async_output_queue = None
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
            side_effect=RuntimeError("SHM pack out of memory"),
        )

        output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput))
        with pytest.raises(RuntimeError, match="SHM pack out of memory"):
            proc._return_result(output)

        proc.result_mq.enqueue.assert_not_called()

    def test_media_packing_failure_is_reraised_for_dp_container(self, mocker):
        """Container dispatch and typed-media failure detection share a walker."""
        proc = _make_worker_proc(step_execution=True)
        proc._async_output_queue = None
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
            side_effect=RuntimeError("SHM pack out of memory"),
        )

        output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput))
        tagged_output = {"dp_rank": 0, "output": output}
        with pytest.raises(RuntimeError, match="SHM pack out of memory"):
            proc._return_result(tagged_output)

        proc.result_mq.enqueue.assert_not_called()

    def test_legacy_output_packing_failure_is_still_swallowed(self, mocker):
        """Non-typed (legacy output) failures keep the historic swallow-and-
        enqueue behavior; only typed media re-raises."""
        proc = _make_worker_proc(step_execution=True)
        proc._async_output_queue = None
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
            side_effect=RuntimeError("SHM pack out of memory"),
        )

        proc._return_result(DiffusionOutput(output="data"))

        # Legacy path: warning logged, output still enqueued once.
        proc.result_mq.enqueue.assert_called_once_with(DiffusionOutput(output="data"))

    def test_non_diffusion_output_skips_async_in_request_mode(self, mocker):
        """Non-DiffusionOutput/BatchRunnerOutput in request mode uses sync path."""
        proc = _make_worker_proc(step_execution=False)

        mock_pack = mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
        )
        # A plain dict is neither DiffusionOutput nor BatchRunnerOutput
        output = {"type": "other"}
        proc._return_result(output)

        mock_pack.assert_called_once_with(output)
        proc.result_mq.enqueue.assert_called_once_with(output)


class TestAsyncOutputLoopLogic:
    """Test _async_output_loop background thread logic (mocked, no device)."""

    def test_bg_thread_enqueues_output_ready_after_pack(self, mocker):
        """The real _async_output_loop packs via SHM with d2h_stream and
        enqueues an OUTPUT_READY message."""
        import threading
        import time

        proc = _make_worker_proc(step_execution=False)

        mock_pack = mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
        )
        # Mock torch to avoid device initialization
        mock_torch = mocker.MagicMock()
        mock_stream = mocker.MagicMock()
        mock_torch.Stream.return_value = mock_stream
        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
        mock_torch.device.return_value = mocker.MagicMock()
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.torch",
            mock_torch,
        )

        output = DiffusionOutput(output="data")
        gpu_event = mocker.MagicMock()
        proc._async_output_queue.put((output, "abc123", gpu_event))

        # Run the real _async_output_loop in a background thread.
        t = threading.Thread(target=proc._async_output_loop, daemon=True)
        t.start()

        # Give the loop time to process, then send the shutdown sentinel.
        time.sleep(0.5)
        proc._async_output_queue.put(None)
        t.join(timeout=2.0)

        mock_pack.assert_called()
        # Verify an OUTPUT_READY with the correct async_output_id was enqueued.
        output_ready_enqueued = any(
            c[0][0].kind == AsyncOutputKind.OUTPUT_READY and c[0][0].async_output_id == "abc123"
            for c in proc.result_mq.enqueue.call_args_list
        )
        assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"

    def test_late_field_failure_sends_error_without_partial_payload_or_shm_leak(
        self,
        mocker,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        from multiprocessing import shared_memory

        proc = _make_worker_proc(step_execution=False)
        media = DiffusionMediaOutput(
            video=VideoMediaOutput(
                tensor=torch.zeros(1, 3, 2, 4, 5),
                spec=VideoTensorSpec(
                    layout=VideoTensorLayout.BCTHW,
                    encoding=VideoTensorEncoding.NORMALIZED_FLOAT,
                    value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE,
                ),
            ),
            prepared_for_transport=True,
        )
        output = DiffusionOutput(media=media, trajectory_latents=torch.ones(30))
        original_latents = output.trajectory_latents

        real_array_to_shm = diffusion_ipc._array_to_shm
        created_handles: list[dict[str, object]] = []

        def _fail_on_second(array):
            if created_handles:
                raise RuntimeError("late-field packing failure")
            handle = real_array_to_shm(array)
            created_handles.append(handle)
            return handle

        monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1)
        monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second)

        def _pack_without_stream(value, d2h_stream=None):
            # The worker loop behavior is under test here. CPU tensors avoid
            # introducing a device dependency into this unit test.
            del d2h_stream
            return diffusion_ipc.pack_diffusion_output_shm(value)

        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
            side_effect=_pack_without_stream,
        )
        mock_torch = mocker.MagicMock()
        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)

        proc._async_output_queue.put((output, "late-failure", None))
        proc._async_output_queue.put(None)
        proc._async_output_loop()

        messages = [call.args[0] for call in proc.result_mq.enqueue.call_args_list]
        assert len(messages) == 1
        assert messages[0].kind is AsyncOutputKind.OUTPUT_READY
        assert messages[0].async_output_id == "late-failure"
        assert messages[0].error == "Background D2H/SHM packing failed"
        assert messages[0].output is None
        assert output.media is media
        assert output.trajectory_latents is original_latents
        with pytest.raises(FileNotFoundError):
            shared_memory.SharedMemory(name=str(created_handles[0]["name"]))


class TestDrainAsyncOutputs:
    """Async outputs must finish packing before device memory is released."""

    def test_drain_returns_immediately_when_nothing_in_flight(self):
        proc = _make_worker_proc(step_execution=False)

        assert proc.drain_async_outputs(timeout=0.01) is True

    def test_return_result_keeps_output_pending_until_loop_completes(self, mocker):
        """_return_result marks work pending; the bg loop clears it."""
        proc = _make_worker_proc(step_execution=False)

        mock_platform = mocker.MagicMock()
        mock_platform.record_device_event.return_value = MagicMock()
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
            mock_platform,
        )
        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm")
        mock_torch = mocker.MagicMock()
        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)

        proc._return_result(DiffusionOutput(output="data"), rpc_id="1")

        # Packing has not run yet, so a sleep here would free tensors in flight.
        assert proc._async_output_pending == 1
        assert proc.drain_async_outputs(timeout=0.05) is False

        thread = threading.Thread(target=proc._async_output_loop, daemon=True)
        thread.start()
        try:
            assert proc.drain_async_outputs(timeout=5.0) is True
        finally:
            proc._async_output_queue.put(None)
            thread.join(timeout=2.0)

    def test_sleep_rpc_drains_before_releasing_memory(self, mocker):
        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()

        call_order: list[str] = []
        mocker.patch.object(
            proc,
            "drain_async_outputs",
            side_effect=lambda *a, **kw: call_order.append("drain") or True,
        )
        proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append("sleep")

        proc._execute_rpc(
            {
                "method": "handle_sleep_task",
                "args": (),
                "kwargs": {},
                "output_rank": 0,
                "exec_all_ranks": False,
                "collect_rank_status": False,
            }
        )

        assert call_order == ["drain", "sleep"]

    def test_non_sleep_rpc_does_not_drain(self, mocker):
        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()
        drain = mocker.patch.object(proc, "drain_async_outputs", return_value=True)

        proc._execute_rpc(
            {
                "method": "execute_model",
                "args": (),
                "kwargs": {},
                "output_rank": 0,
                "exec_all_ranks": False,
                "collect_rank_status": False,
            }
        )

        drain.assert_not_called()


class TestWorkerProcShutdown:
    def test_shutdown_stops_async_thread_and_releases_queues(self):
        import threading

        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()
        proc.mq = MagicMock()
        broadcast_mq = proc.mq
        result_mq = proc.result_mq

        thread = threading.Thread(target=proc._async_output_queue.get)
        proc._async_output_thread = thread
        thread.start()

        proc.shutdown()

        assert not thread.is_alive()
        proc.worker.shutdown.assert_called_once_with()
        broadcast_mq.shutdown.assert_called_once_with()
        result_mq.shutdown.assert_called_once_with()
        assert proc.mq is None
        assert proc.result_mq is None

    def test_shutdown_releases_queues_when_worker_shutdown_fails(self):
        proc = _make_worker_proc(step_execution=True)
        proc.worker = MagicMock()
        proc.worker.shutdown.side_effect = RuntimeError("shutdown failed")
        proc.mq = MagicMock()
        broadcast_mq = proc.mq
        result_mq = proc.result_mq

        with pytest.raises(RuntimeError, match="shutdown failed"):
            proc.shutdown()

        broadcast_mq.shutdown.assert_called_once_with()
        result_mq.shutdown.assert_called_once_with()
        assert proc.mq is None
        assert proc.result_mq is None


class TestSynchronizeDeviceBarrier:
    """The pause barrier must not report a quiet device while outputs still pack."""

    @staticmethod
    def _rpc(timeout):
        return {
            "method": "synchronize_device",
            "args": (),
            "kwargs": {"timeout": timeout},
            "output_rank": 0,
            "exec_all_ranks": False,
            "collect_rank_status": False,
        }

    def test_rejects_undrained_output(self, mocker):
        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()
        mocker.patch.object(proc, "drain_async_outputs", return_value=False)

        with pytest.raises(TimeoutError, match="did not drain"):
            proc._execute_rpc(self._rpc(0.01))

        proc.worker.execute_method.assert_not_called()

    @pytest.mark.parametrize("timeout", [0.5, None])
    def test_drains_with_the_call_timeout_then_synchronizes(self, mocker, timeout):
        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()
        call_order: list = []
        mocker.patch.object(
            proc,
            "drain_async_outputs",
            side_effect=lambda **kw: call_order.append(("drain", kw)) or True,
        )
        proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append(("sync", a, kw))

        _, should_reply = proc._execute_rpc(self._rpc(timeout))

        assert call_order == [
            ("drain", {"timeout": timeout}),
            ("sync", ("synchronize_device",), {"timeout": timeout}),
        ]
        assert should_reply is True

    def test_completes_only_after_output_ready_is_published(self, mocker):
        proc = _make_worker_proc(step_execution=False)
        proc.worker = MagicMock()

        mock_platform = mocker.MagicMock()
        mock_platform.record_device_event.return_value = MagicMock()
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
            mock_platform,
        )
        mock_torch = mocker.MagicMock()
        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
        release = threading.Event()
        mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm",
            side_effect=lambda *a, **kw: release.wait(timeout=5.0),
        )

        proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
        loop_thread = threading.Thread(target=proc._async_output_loop, daemon=True)
        loop_thread.start()

        barrier_done = threading.Event()

        def run_barrier():
            proc._execute_rpc(self._rpc(5.0))
            barrier_done.set()

        barrier_thread = threading.Thread(target=run_barrier, daemon=True)
        barrier_thread.start()
        try:
            assert not barrier_done.wait(timeout=0.2)
            proc.worker.execute_method.assert_not_called()

            release.set()
            assert barrier_done.wait(timeout=5.0)

            kinds = [
                call.args[0].kind
                for call in proc.result_mq.enqueue.call_args_list
                if isinstance(call.args[0], AsyncDiffusionOutput)
            ]
            assert kinds == [AsyncOutputKind.COMPUTE_DONE, AsyncOutputKind.OUTPUT_READY]
            proc.worker.execute_method.assert_called_once_with("synchronize_device", timeout=5.0)
        finally:
            release.set()
            proc._async_output_queue.put(None)
            loop_thread.join(timeout=2.0)
            barrier_thread.join(timeout=2.0)