Enable Diffusion Image models asynchronous output (#4981)

Signed-off-by: zzh <943967662@qq.com>
Co-authored-by: dengyunyang <584797741@qq.com>

tests/diffusion/test_async_output_worker.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for WorkerProc async output logic."""
+
+import queue
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.data import (
+    AsyncDiffusionOutput,
+    AsyncOutputKind,
+    DiffusionOutput,
+    OmniACK,
+)
+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
+    return proc
+
+
+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_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 unblock and stop.
+        time.sleep(0.5)
+        proc._running = False
+        proc._async_output_queue.put((DiffusionOutput(), "sentinel", 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"

tests/diffusion/test_ipc_async.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for IPC async D2H stream packing."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.ipc import (
+    _pack_diffusion_fields,
+    _pack_tensor_if_large,
+    _pack_value_if_large,
+    pack_diffusion_output_shm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestPackTensorIfLargeWithD2hStream:
+    """Test _pack_tensor_if_large with mock d2h_stream parameter."""
+
+    def test_small_tensor_not_packed_with_stream(self, mocker):
+        """Small tensors pass through even when d2h_stream is provided."""
+        small_tensor = torch.zeros(10)  # way below threshold
+        mock_stream = mocker.MagicMock()
+        result = _pack_tensor_if_large(small_tensor, d2h_stream=mock_stream)
+        # Small tensor should be returned as-is (not a SHM handle dict)
+        assert isinstance(result, torch.Tensor)
+
+    def test_large_tensor_packed_without_stream(self, mocker):
+        """Large tensor is packed via SHM when no d2h_stream."""
+        large_tensor = torch.randn(500_000)  # > 1MB for float32
+        result = _pack_tensor_if_large(large_tensor, d2h_stream=None)
+        assert isinstance(result, dict)
+        assert result.get("__tensor_shm__") is True
+
+    def test_pack_value_if_large_passes_stream_to_tensor_pack(self, mocker):
+        """_pack_value_if_large propagates d2h_stream to _pack_tensor_if_large."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        tensor = torch.zeros(1)
+        d2h = mocker.MagicMock()
+        _pack_value_if_large(tensor, d2h_stream=d2h)
+        mock_pack.assert_called_once_with(tensor, d2h_stream=d2h)
+
+    def test_pack_value_if_large_passes_stream_through_dict(self, mocker):
+        """d2h_stream propagates through dict container packing."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        data = {"image": torch.zeros(1), "video": torch.zeros(1)}
+        _pack_value_if_large(data, d2h_stream=d2h)
+        assert mock_pack.call_count == 2
+        for call in mock_pack.call_args_list:
+            assert call.kwargs.get("d2h_stream") is d2h or call.args[1] is d2h
+
+    def test_pack_value_if_large_passes_stream_through_list(self, mocker):
+        """d2h_stream propagates through list container packing."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        data = [torch.zeros(1), torch.zeros(1)]
+        _pack_value_if_large(data, d2h_stream=d2h)
+        assert mock_pack.call_count == 2
+
+    def test_pack_value_if_large_passes_stream_through_tuple(self, mocker):
+        """d2h_stream propagates through tuple container packing."""
+        mock_pack = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        data = (torch.zeros(1), torch.zeros(1))
+        result = _pack_value_if_large(data, d2h_stream=d2h)
+        assert isinstance(result, tuple)
+        assert mock_pack.call_count == 2
+
+
+class TestPackDiffusionFieldsWithD2hStream:
+    """Test _pack_diffusion_fields propagates d2h_stream."""
+
+    def test_output_field_packed_with_stream(self, mocker):
+        mock_pack_value = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_value_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        out = DiffusionOutput(output=torch.zeros(1))
+        _pack_diffusion_fields(out, d2h_stream=d2h)
+        mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h)
+
+    def test_trajectory_fields_packed_with_stream(self, mocker):
+        mock_pack_tensor = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_tensor_if_large",
+            return_value="packed",
+        )
+        d2h = mocker.MagicMock()
+        out = DiffusionOutput(
+            trajectory_latents=torch.zeros(1),
+            trajectory_log_probs=torch.zeros(1),
+        )
+        _pack_diffusion_fields(out, d2h_stream=d2h)
+        # Should pack trajectory_latents and trajectory_log_probs
+        assert mock_pack_tensor.call_count == 2
+
+
+class TestPackDiffusionOutputShmWithD2hStream:
+    """Test pack_diffusion_output_shm top-level dispatch with d2h_stream."""
+
+    def test_diffusion_output_path_with_stream(self, mocker):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+        out = DiffusionOutput()
+        pack_diffusion_output_shm(out, d2h_stream=d2h)
+        mock_pack_fields.assert_called_once_with(out, d2h_stream=d2h)
+
+    def test_rpc_envelope_path_with_stream(self, mocker):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+        result = DiffusionOutput()
+        envelope = {"type": "diffusion_rpc_result", "result": result}
+        pack_diffusion_output_shm(envelope, d2h_stream=d2h)
+        # envelope["result"] is replaced in-place by pack_diffusion_output_shm;
+        # assert with the original reference, not the replaced one.
+        mock_pack_fields.assert_called_once_with(
+            result,
+            d2h_stream=d2h,
+        )
+
+    def test_runner_output_path_with_stream(self, mocker):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+
+        from vllm_omni.diffusion.worker.utils import RunnerOutput
+
+        result = DiffusionOutput()
+        runner_out = RunnerOutput(
+            request_id="req1",
+            result=result,
+        )
+        pack_diffusion_output_shm(runner_out, d2h_stream=d2h)
+        # runner_out.result is replaced in-place by pack_diffusion_output_shm;
+        # assert with the original reference, not the replaced one.
+        mock_pack_fields.assert_called_once_with(
+            result,
+            d2h_stream=d2h,
+        )

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -56,7 +56,7 @@ def _make_executor(num_gpus: int = 1):
 
     Returns ``(executor, request_queue, result_queue)``.
     """
-    od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False)
+    od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False, step_execution=True)
     executor = object.__new__(MultiprocDiffusionExecutor)
     executor.od_config = od_cfg
 
@@ -616,6 +616,7 @@ class TestMultiprocExecutorRaisesEngineDeadError:
 
     def test_collective_rpc_raises_when_is_failed(self):
         executor = object.__new__(MultiprocDiffusionExecutor)
+        executor.od_config = SimpleNamespace(step_execution=True)
         executor._closed = False
         executor._broadcast_mq = MagicMock()
         executor._result_mq = MagicMock()
@@ -661,7 +662,7 @@ class TestMultiprocExecutorStepStreamingOutput:
 
     def test_execute_step_allows_streaming_output_mode(self):
         executor, req_q, res_q = _make_executor()
-        executor.od_config = SimpleNamespace(streaming_output=True)  # pyright: ignore[reportAttributeAccessIssue]
+        executor.od_config = SimpleNamespace(streaming_output=True, step_execution=True)  # pyright: ignore[reportAttributeAccessIssue]
         runner_outputs = [
             RunnerOutput(
                 request_id="sched-stream",
@@ -939,6 +940,14 @@ class TestMultiprocExecutorWorkerMonitor:
         executor._shutdown_cleaner = None
         # Use a no-op so shutdown() doesn't crash on None resources.
         executor._finalizer = lambda: None
+        # ------------------------------------------------------------------
+        # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates
+        # over them, so they need to exist even when constructed via __new__.
+        executor._pump_stop = threading.Event()
+        executor._futures_lock = threading.RLock()
+        executor._rpc_futures = {}
+        executor._output_futures = {}
+        executor._batch_split_map = {}
 
         proc = _make_short_lived_process()
         executor._processes = [proc]

tests/diffusion/test_result_pump.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MultiprocDiffusionExecutor async result pump and wait_output_ready."""
+
+import concurrent.futures
+import queue
+import threading
+import time
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.diffusion.data import AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_executor(step_execution=False):
+    """Create a minimal MultiprocDiffusionExecutor-like object with pump state."""
+    from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+
+    od_config = MagicMock()
+    od_config.step_execution = step_execution
+
+    executor = object.__new__(MultiprocDiffusionExecutor)
+    executor.od_config = od_config
+    executor._rpc_id_counter = 0
+    executor._rpc_id_lock = threading.Lock()
+    executor._rpc_futures = {}
+    executor._output_futures = {}
+    executor._completed_outputs = {}
+    executor._batch_split_map = {}
+    executor._futures_lock = threading.RLock()
+    executor._pump_running = False
+    executor._pump_stop = threading.Event()
+    executor._sync_result_buffer = queue.Queue()
+    executor._result_mq = MagicMock()
+    executor._broadcast_mq = MagicMock()
+    executor._closed = False
+    executor._is_failed = False
+    executor._finalizer = MagicMock()  # no-op in tests
+    executor._shutdown_cleaner = None
+    executor._processes = []
+    return executor
+
+
+def _feed_one_msg_to_pump(executor, msg):
+    """Run _result_pump in a daemon thread, feed one *msg*, then stop."""
+    call_count = [0]
+
+    def mock_dequeue(timeout=None):
+        call_count[0] += 1
+        if call_count[0] == 1:
+            return msg
+        executor._pump_stop.set()
+        time.sleep(0.05)
+        raise TimeoutError
+
+    executor._result_mq.dequeue = mock_dequeue
+    t = threading.Thread(target=executor._result_pump, daemon=True)
+    t.start()
+    t.join(timeout=2.0)
+
+
+@pytest.fixture(autouse=True)
+def _mock_unpack(mocker):
+    """Real _result_pump calls unpack_diffusion_output_shm; mock it away."""
+    mocker.patch(
+        "vllm_omni.diffusion.executor.multiproc_executor.unpack_diffusion_output_shm",
+    )
+
+
+class TestNextRpcId:
+    """Test _next_rpc_id counter."""
+
+    def test_counter_increments(self):
+        executor = _make_executor()
+        id1 = executor._next_rpc_id()
+        id2 = executor._next_rpc_id()
+        id3 = executor._next_rpc_id()
+        assert id1 == "1"
+        assert id2 == "2"
+        assert id3 == "3"
+
+    def test_counter_is_threadsafe(self):
+        executor = _make_executor()
+        ids = []
+
+        def get_id():
+            ids.append(executor._next_rpc_id())
+
+        threads = [threading.Thread(target=get_id) for _ in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # All IDs should be unique
+        assert len(set(ids)) == 10
+
+
+class TestWaitOutputReady:
+    """Test wait_output_ready future creation and caching."""
+
+    def test_returns_new_future_when_not_cached(self):
+        executor = _make_executor()
+        fut = executor.wait_output_ready("abc123")
+        assert isinstance(fut, concurrent.futures.Future)
+        assert not fut.done()
+
+    def test_future_resolves_when_output_arrives(self):
+        executor = _make_executor()
+        fut = executor.wait_output_ready("abc123")
+        output = DiffusionOutput(output="data")
+
+        # Simulate pump resolving the future
+        with executor._futures_lock:
+            executor._output_futures.pop("abc123")
+        if not fut.done():
+            fut.set_result(output)
+
+        assert fut.result(timeout=1.0) is output
+
+    def test_returns_cached_future_when_already_completed(self):
+        executor = _make_executor()
+        output = DiffusionOutput(output="cached_data")
+        fut = concurrent.futures.Future()
+        fut.set_result(output)
+        with executor._futures_lock:
+            executor._completed_outputs["abc123"] = fut
+
+        fut = executor.wait_output_ready("abc123")
+        assert fut.done()
+        assert fut.result(timeout=1.0) is output
+
+    def test_removes_from_cache_after_retrieval(self):
+        executor = _make_executor()
+        output = DiffusionOutput(output="cached_data")
+        with executor._futures_lock:
+            executor._completed_outputs["abc123"] = output
+
+        executor.wait_output_ready("abc123")
+        # Second call should not find cached result
+        with executor._futures_lock:
+            assert "abc123" not in executor._completed_outputs
+
+
+class TestResultPumpDispatch:
+    """Test _result_pump message routing (running the real pump in a thread)."""
+
+    def test_non_async_message_placed_in_sync_buffer(self):
+        executor = _make_executor()
+        msg = DiffusionOutput(output="sync_result")
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert not executor._sync_result_buffer.empty()
+        retrieved = executor._sync_result_buffer.get_nowait()
+        assert isinstance(retrieved, DiffusionOutput)
+
+    def test_compute_done_routes_to_rpc_future(self):
+        executor = _make_executor()
+        rpc_id = "42"
+        fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._rpc_futures[rpc_id] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.COMPUTE_DONE,
+            rpc_id=rpc_id,
+            async_output_id="abc",
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        result = fut.result(timeout=1.0)
+        assert result.kind == AsyncOutputKind.COMPUTE_DONE
+
+    def test_output_ready_routes_to_output_future(self):
+        executor = _make_executor()
+        async_output_id = "abc123"
+        output = DiffusionOutput(output="final")
+        fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._output_futures[async_output_id] = fut
+
+        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
+
+    def test_output_ready_with_error_routes_to_future_as_exception(self):
+        executor = _make_executor()
+        async_output_id = "abc123"
+        fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._output_futures[async_output_id] = fut
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            error="Background D2H/SHM packing failed",
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        assert fut.done()
+        with pytest.raises(RuntimeError, match="Background D2H/SHM packing failed"):
+            fut.result(timeout=1.0)
+
+    def test_output_ready_caches_when_no_future_waiting(self):
+        """When OUTPUT_READY arrives but no future is waiting, result is cached."""
+        executor = _make_executor()
+        async_output_id = "abc123"
+        output = DiffusionOutput(output="orphan")
+
+        msg = AsyncDiffusionOutput(
+            kind=AsyncOutputKind.OUTPUT_READY,
+            async_output_id=async_output_id,
+            output=output,
+        )
+        _feed_one_msg_to_pump(executor, msg)
+
+        # Later call to wait_output_ready should find it cached
+        fut = executor.wait_output_ready(async_output_id)
+        assert fut.done()
+        assert fut.result(timeout=1.0) is output
+
+
+class TestShutdownCleansUpFutures:
+    """Test that shutdown cancels pending futures."""
+
+    def test_shutdown_sets_exception_on_pending_futures(self):
+        executor = _make_executor()
+
+        rpc_fut = concurrent.futures.Future()
+        output_fut = concurrent.futures.Future()
+        with executor._futures_lock:
+            executor._rpc_futures["1"] = rpc_fut
+            executor._output_futures["abc"] = output_fut
+
+        executor.shutdown()
+
+        assert rpc_fut.done()
+        with pytest.raises(RuntimeError, match="Executor shut down"):
+            rpc_fut.result(timeout=1.0)
+
+        assert output_fut.done()
+        with pytest.raises(RuntimeError, match="Executor shut down"):
+            output_fut.result(timeout=1.0)
+
+        assert len(executor._rpc_futures) == 0
+        assert len(executor._output_futures) == 0