[Core][Diffusion] Add a typed pre-D2H video media contract (#6615)

Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>
Co-authored-by: Samit <285365963@qq.com>

tests/diffusion/test_async_output_worker.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for WorkerProc async output logic."""
 
 import queue
@@ -8,13 +8,23 @@ 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
 
@@ -182,6 +192,66 @@ class TestReturnResultSyncPath:
         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)
@@ -243,6 +313,70 @@ class TestAsyncOutputLoopLogic:
         )
         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."""