[Bugfix] Fix lost async diffusion outputs in request-level batching (#6023)

Signed-off-by: samithuang <285365963@qq.com>

tests/diffusion/test_async_output_worker.py

@@ -3,6 +3,8 @@
 """Unit tests for WorkerProc async output logic."""
 
 import queue
+import threading
+import time
 from unittest.mock import MagicMock
 
 import pytest
@@ -29,9 +31,59 @@ def _make_worker_proc(step_execution=False):
     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."""
 
@@ -192,6 +244,87 @@ class TestAsyncOutputLoopLogic:
         assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
 
 
+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

tests/diffusion/test_diffusion_ipc.py

@@ -305,6 +305,49 @@ def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
     assert output["req-error"].result.error == "boom"
 
 
+def test_per_request_views_of_a_batch_tensor_do_not_ship_whole_storage() -> None:
+    """Regression: per-request slices used to pickle the whole batch storage.
+
+    Each view is under the packing threshold, but pickle writes the entire
+    shared storage per view, so a batch's message grew by the full batch
+    tensor once per request and overflowed the ring buffer.
+    """
+    import pickle
+
+    num_requests = 16
+    # Each row stays under the threshold; the shared storage is well over it.
+    row_numel = _SHM_TENSOR_THRESHOLD // (2 * torch.empty((), dtype=torch.float32).element_size())
+    batch = torch.arange(num_requests * row_numel, dtype=torch.float32).reshape(num_requests, row_numel)
+    per_request = batch[0]
+    assert per_request.nelement() * per_request.element_size() <= _SHM_TENSOR_THRESHOLD
+    assert per_request.untyped_storage().nbytes() > _SHM_TENSOR_THRESHOLD
+
+    output = BatchRunnerOutput.from_list(
+        [
+            RunnerOutput(
+                request_id=f"req-{i}",
+                finished=True,
+                result=DiffusionOutput(trajectory_latents=batch[i]),
+            )
+            for i in range(num_requests)
+        ]
+    )
+    unpacked_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
+
+    pack_diffusion_output_shm(output)
+    packed_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
+
+    assert output.runner_outputs[0].result.trajectory_latents["__tensor_shm__"] is True
+    # Whole storage per request before; handles only after.
+    storage_bytes = per_request.untyped_storage().nbytes()
+    assert unpacked_size > num_requests * storage_bytes / 2
+    assert packed_size < _SHM_TENSOR_THRESHOLD
+
+    unpack_diffusion_output_shm(output)
+    for i in range(num_requests):
+        torch.testing.assert_close(output[f"req-{i}"].result.trajectory_latents, batch[i])
+
+
 def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
     tensor = torch.arange(
         _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),

tests/diffusion/test_result_pump.py

@@ -6,6 +6,7 @@ import concurrent.futures
 import queue
 import threading
 import time
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -264,6 +265,119 @@ class TestResultPumpDispatch:
         fut = executor.wait_output_ready(async_output_id)
         assert fut.done()
         assert fut.result(timeout=1.0) is output
+        assert async_output_id not in executor._output_futures
+
+    def test_output_ready_atomic_resolution_when_future_already_waiting(self):
+        """When OUTPUT_READY arrives and a future is already waiting, resolve it directly."""
+        executor = _make_executor()
+        async_output_id = "abc123"
+        output = DiffusionOutput(output="waiting")
+        fut = executor.wait_output_ready(async_output_id)
+        assert not fut.done()
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            output=output,
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        assert fut.result(timeout=1.0) is output
+        assert async_output_id not in executor._output_futures
+        assert async_output_id not in executor._completed_outputs
+
+
+class _FakeBatchOutput:
+    """Batch-level output exposing per-request results."""
+
+    def __init__(self, results):
+        self._results = results
+
+    def get_request_output(self, req_id):
+        result = self._results.get(req_id)
+        if result is None:
+            return None
+        return SimpleNamespace(result=result)
+
+
+def _make_scheduler_output(req_ids):
+    return SimpleNamespace(
+        scheduled_new_reqs=[SimpleNamespace(request_id=rid, req=SimpleNamespace()) for rid in req_ids]
+    )
+
+
+class TestBatchSplitDelivery:
+    """execute_batch must resolve per-request futures in either arrival order."""
+
+    @staticmethod
+    def _run(executor, req_ids, batch_id, deliver_early):
+        # Keep execute_batch on the fused request-batch path (not DLO DP).
+        executor.od_config.parallel_config.data_parallel_size = 1
+        executor.od_config.enable_distributed_layerwise_offload = False
+        executor._ensure_open = lambda: None
+
+        outputs = {rid: DiffusionOutput(output=f"img-{rid}") for rid in req_ids}
+        ready = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=batch_id,
+            output=_FakeBatchOutput(outputs),
+        )
+
+        def fake_collective_rpc(*args, **kwargs):
+            if deliver_early:
+                # Worker's background D2H/SHM thread wins the race: OUTPUT_READY
+                # is pumped before execute_batch registers the split map.
+                _feed_one_msg_to_pump(executor, ready)
+            return AsyncDiffusionOutput(
+                kind=AsyncOutputKind.COMPUTE_DONE,
+                rpc_id="1",
+                async_output_id=batch_id,
+            )
+
+        executor.collective_rpc = fake_collective_rpc
+        batch = executor.execute_batch(_make_scheduler_output(req_ids))
+        if not deliver_early:
+            _feed_one_msg_to_pump(executor, ready)
+        return batch, outputs
+
+    def test_output_ready_after_split_map(self):
+        executor = _make_executor()
+        req_ids = ["r0", "r1", "r2"]
+        _, outputs = self._run(executor, req_ids, "batch-1", deliver_early=False)
+
+        for rid in req_ids:
+            fut = executor.wait_output_ready(f"batch-1/{rid}")
+            assert fut.done()
+            assert fut.result(timeout=1.0) is outputs[rid]
+
+    def test_output_ready_before_split_map(self):
+        """Regression: the whole batch used to be lost, hanging every request."""
+        executor = _make_executor()
+        req_ids = ["r0", "r1", "r2"]
+        _, outputs = self._run(executor, req_ids, "batch-1", deliver_early=True)
+
+        for rid in req_ids:
+            fut = executor.wait_output_ready(f"batch-1/{rid}")
+            assert fut.done(), f"request {rid} never resolved"
+            assert fut.result(timeout=1.0) is outputs[rid]
+
+        # No stale batch-level state left behind.
+        assert executor._batch_split_map == {}
+        assert executor._completed_outputs == {}
+
+    def test_engine_waiter_before_output_is_resolved(self):
+        """Consumers already blocked in wait_output_ready must be woken."""
+        executor = _make_executor()
+        req_ids = ["r0", "r1"]
+        batch_id = "batch-2"
+        futures = {rid: executor.wait_output_ready(f"{batch_id}/{rid}") for rid in req_ids}
+
+        _, outputs = self._run(executor, req_ids, batch_id, deliver_early=True)
+
+        for rid in req_ids:
+            assert futures[rid].done(), f"request {rid} never resolved"
+            assert futures[rid].result(timeout=1.0) is outputs[rid]
 
 
 class TestShutdownCleansUpFutures: