[Diffusion] Fix DLO DP concurrent request execution (#5864)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/test_async_output_worker.py

@@ -178,10 +178,9 @@ class TestAsyncOutputLoopLogic:
         t = threading.Thread(target=proc._async_output_loop, daemon=True)
         t.start()
 
-        # Give the loop time to process, then unblock and stop.
+        # Give the loop time to process, then send the shutdown sentinel.
         time.sleep(0.5)
-        proc._running = False
-        proc._async_output_queue.put((DiffusionOutput(), "sentinel", None))
+        proc._async_output_queue.put(None)
         t.join(timeout=2.0)
 
         mock_pack.assert_called()
@@ -191,3 +190,43 @@ class TestAsyncOutputLoopLogic:
             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"
+
+
+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

tests/diffusion/test_diffusion_engine.py

@@ -331,6 +331,55 @@ class TestRequestBatchCapability:
             DiffusionEngine(od_config)
         fake_executor_cls.assert_not_called()
 
+    def test_engine_allows_independent_dlo_dp_requests_for_single_request_pipeline(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            streaming_output=False,
+            max_num_seqs=2,
+            request_batch_max_wait_ms=0,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        fake_executor = SimpleNamespace(
+            execute_request=mocker.Mock(return_value="per-request"),
+            execute_batch=mocker.Mock(return_value="batch"),
+            execute_step=mocker.Mock(return_value="step"),
+        )
+        fake_executor_cls = mocker.Mock(return_value=fake_executor)
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+
+        engine = DiffusionEngine(od_config)
+
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
+        assert engine.supports_request_batch is False
+        assert engine.dp_concurrent is True
+        assert engine.scheduler.max_num_running_reqs == 2
+        assert od_config.request_batch_max_wait_ms == 0
+        fake_executor_cls.assert_called_once_with(od_config)
+
     @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
     def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
         self,

tests/diffusion/test_diffusion_ipc.py

@@ -2,11 +2,19 @@
 # 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,
@@ -30,6 +38,129 @@ def _cleanup_shm_handle(value: object) -> None:
             _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
@@ -83,6 +214,29 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
     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 = {
@@ -108,6 +262,25 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
     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

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -75,6 +75,32 @@ async def test_inline_dispatch_request_success(client, mock_engine):
     assert output.request_id == "req-1"
 
 
+@pytest.mark.asyncio
+async def test_inline_requests_clone_shared_sampling_params(client, mock_engine):
+    requests = []
+
+    async def _step_streaming(request):
+        requests.append(request)
+        yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
+
+    mock_engine.step_streaming = _step_streaming
+
+    shared = OmniDiffusionSamplingParams()
+    await client.add_request_async("req-1", "First prompt", shared)
+    await client.add_request_async("req-2", "Second prompt", shared)
+
+    for _ in range(10):
+        if len(requests) == 2:
+            break
+        await asyncio.sleep(0.01)
+
+    assert len(requests) == 2
+    assert requests[0].sampling_params is not requests[1].sampling_params
+    assert requests[0].sampling_params.guidance_scale_provided is False
+    assert requests[1].sampling_params.guidance_scale_provided is False
+    assert shared.guidance_scale is None
+
+
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_streaming_success(client, mock_engine):
     """Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
@@ -137,6 +163,20 @@ def test_inline_shutdown(client, mock_engine):
     mock_engine.close.assert_called_once()
 
 
+def test_inline_shutdown_runs_after_orchestrator_premark(client, mock_engine):
+    # The orchestrator pre-marks clients to suppress false worker-death errors
+    # before StagePool invokes the actual teardown.
+    client._shutting_down = True
+
+    client.shutdown()
+
+    mock_engine.close.assert_called_once()
+    assert client._shutdown_complete
+
+    client.shutdown()
+    mock_engine.close.assert_called_once()
+
+
 def test_inline_registers_executor_failure_callback(client, mock_engine):
     mock_engine.executor.register_failure_callback.assert_called_once()
 

tests/diffusion/test_ipc_async.py

@@ -137,6 +137,18 @@ class TestPackDiffusionOutputShmWithD2hStream:
             d2h_stream=d2h,
         )
 
+    @pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True])
+    def test_dp_tagged_output_path_with_stream(self, mocker, wrapped_in_rpc_envelope):
+        mock_pack_fields = mocker.patch(
+            "vllm_omni.diffusion.ipc._pack_diffusion_fields",
+        )
+        d2h = mocker.MagicMock()
+        result = DiffusionOutput()
+        tagged = {"dp_rank": 1, "output": result}
+        value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged
+        pack_diffusion_output_shm(value, d2h_stream=d2h)
+        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",

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -276,6 +276,8 @@ class TestRequestModeDispatch:
         def worker():
             request = req_q.get(timeout=2)
             assert "rpc_id" not in request
+            assert request["output_rank"] is None
+            assert request["collect_rank_status"] is False
             wave_id = request["wave_id"]
             for dp_rank in range(2):
                 res_q.put(
@@ -297,6 +299,181 @@ class TestRequestModeDispatch:
 
         assert [result.error for result in results] == ["0", "1"]
 
+    @pytest.mark.parametrize(
+        ("tp_size", "sp_size", "pp_size", "cfg_size", "expected_primary_ranks"),
+        [
+            (2, 1, 1, 1, [0, 2]),
+            (2, 2, 1, 1, [0, 4]),
+            (1, 2, 2, 1, [0, 4]),
+            (1, 1, 2, 2, [0, 4]),
+        ],
+    )
+    def test_step_execution_reads_each_dp_primary_result_queue(
+        self,
+        tp_size,
+        sp_size,
+        pp_size,
+        cfg_size,
+        expected_primary_ranks,
+    ):
+        dp_size = 2
+        num_gpus = dp_size * tp_size * sp_size * pp_size * cfg_size
+        executor, _, _ = _make_executor(num_gpus=num_gpus)
+        executor.od_config = SimpleNamespace(
+            step_execution=True,
+            parallel_config=SimpleNamespace(
+                data_parallel_size=dp_size,
+                tensor_parallel_size=tp_size,
+                sequence_parallel_size=sp_size,
+                pipeline_parallel_size=pp_size,
+                cfg_parallel_size=cfg_size,
+            ),
+        )
+        request: dict = {}
+        executor._broadcast_mq = SimpleNamespace(enqueue=lambda value: request.update(value))
+
+        def result_queue(global_rank):
+            result_mq = MagicMock()
+            result_mq.dequeue.side_effect = lambda timeout=None: {
+                "dp_rank": global_rank,
+                "output": _tagged_output(str(global_rank)),
+                "wave_id": request["wave_id"],
+            }
+            return result_mq
+
+        # Every queue returns a valid, rank-tagged result. If the executor
+        # selects a non-primary queue, the returned rank exposes the mistake.
+        executor._result_mqs = [result_queue(rank) for rank in range(num_gpus)]
+        executor._result_mq = executor._result_mqs[0]
+
+        results = executor.collective_rpc(
+            "execute_model",
+            unique_reply_rank=None,
+            exec_all_ranks=True,
+        )
+
+        assert [result.error for result in results] == [str(rank) for rank in expected_primary_ranks]
+        for rank, result_mq in enumerate(executor._result_mqs):
+            assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
+
+    @pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
+    def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("invalid", "valid")
+        scheduler_output.scheduled_new_reqs[0].req.prompt = empty_prompt
+
+        with pytest.raises(ValueError, match="non-empty prompt"):
+            executor.execute_request(scheduler_output)
+
+        executor.collective_rpc.assert_not_called()
+
+    def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        invalid_wave = _make_sched_output("invalid", "valid")
+        invalid_wave.scheduled_new_reqs[0].req.prompt = ""
+
+        with pytest.raises(ValueError, match="non-empty prompt"):
+            executor.execute_request(invalid_wave)
+
+        result = executor.execute_request(_make_sched_output("A", "B"))
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        executor.collective_rpc.assert_called_once()
+
+    def test_dlo_dp_allows_shared_default_denoise_steps(self):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
+        scheduler_output = _make_sched_output("A", "B")
+        for new_req in scheduler_output.scheduled_new_reqs:
+            new_req.req.sampling_params.num_inference_steps = None
+
+        result = executor.execute_request(scheduler_output)
+
+        assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
+        executor.collective_rpc.assert_called_once()
+
+    @pytest.mark.parametrize(
+        ("field", "value"),
+        [
+            ("num_inference_steps", 2),
+            ("guidance_scale", 7.5),
+            ("width", 1024),
+        ],
+    )
+    def test_dlo_dp_rejects_incompatible_collective_wave(self, field, value):
+        executor, _, _ = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor.collective_rpc = Mock()
+        scheduler_output = _make_sched_output("A", "B")
+        setattr(scheduler_output.scheduled_new_reqs[1].req.sampling_params, field, value)
+
+        with pytest.raises(ValueError, match="compatible shape, CFG"):
+            executor.execute_request(scheduler_output)
+
+        executor.collective_rpc.assert_not_called()
+
+    def test_dlo_dp_partial_reply_times_out_and_fails_closed(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        executor, req_q, res_q = _make_executor(num_gpus=2)
+        executor.od_config = SimpleNamespace(
+            step_execution=False,
+            parallel_config=SimpleNamespace(data_parallel_size=2),
+            enable_distributed_layerwise_offload=True,
+            dlo_use_allgather=True,
+        )
+        executor._sync_result_buffer = res_q
+        executor._fail_closed_on_dp_wave_timeout = Mock()
+        monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05)
+
+        def reply_from_only_one_dp_rank():
+            request = req_q.get(timeout=2.0)
+            res_q.put(
+                {
+                    "dp_rank": 0,
+                    "output": _tagged_output("rank-0"),
+                    "wave_id": request["wave_id"],
+                }
+            )
+
+        worker = threading.Thread(target=reply_from_only_one_dp_rank, daemon=True)
+        worker.start()
+        started = time.monotonic()
+
+        result = executor.execute_request(_make_sched_output("A", "B"))
+
+        elapsed = time.monotonic() - started
+        worker.join(timeout=2.0)
+        assert elapsed < 1.0
+        assert all("timed out" in output.result.error for output in result.runner_outputs)
+        executor._fail_closed_on_dp_wave_timeout.assert_called_once()
+        assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
+
 
 # ───────────────── concurrent collective RPC ─────────────────
 
@@ -948,6 +1125,42 @@ class TestStageDiffusionClientErrorPropagation:
         assert rpc_id not in client._pending_rpcs
 
 
+class TestExecutorShutdownCleaner:
+    def test_worker_joins_share_one_global_deadline(self, monkeypatch):
+        from vllm_omni.diffusion.executor import multiproc_executor as executor_module
+
+        class FakeProcess:
+            def __init__(self, name):
+                self.name = name
+                self.alive = True
+                self.terminated = False
+                self.join_timeouts = []
+
+            def is_alive(self):
+                return self.alive
+
+            def join(self, timeout):
+                self.join_timeouts.append(timeout)
+                if self.terminated:
+                    self.alive = False
+
+            def terminate(self):
+                self.terminated = True
+
+        monotonic = Mock(side_effect=[100.0, 100.0, 110.0, 120.0, 120.0, 124.0])
+        monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
+        first = FakeProcess("worker-0")
+        second = FakeProcess("worker-1")
+        cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
+
+        cleaner()
+
+        assert first.join_timeouts == [15.0, 5.0]
+        assert second.join_timeouts == [5.0, 1.0]
+        assert first.terminated and second.terminated
+        assert not first.is_alive() and not second.is_alive()
+
+
 # ───────── monitor thread & death sentinel integration tests ─────────
 
 

tests/diffusion/test_result_pump.py

@@ -35,6 +35,7 @@ def _make_executor(step_execution=False):
     executor._pump_stop = threading.Event()
     executor._sync_result_buffer = queue.Queue()
     executor._result_mq = MagicMock()
+    executor._result_mqs = []
     executor._broadcast_mq = MagicMock()
     executor._closed = False
     executor._is_failed = False
@@ -157,6 +158,41 @@ class TestResultPumpDispatch:
         retrieved = executor._sync_result_buffer.get_nowait()
         assert isinstance(retrieved, DiffusionOutput)
 
+    def test_start_result_pump_reads_every_worker_queue(self):
+        executor = _make_executor()
+        messages = [DiffusionOutput(output="rank0"), DiffusionOutput(output="rank1")]
+        result_mqs = [MagicMock(), MagicMock()]
+
+        def make_dequeue(message):
+            emitted = False
+
+            def dequeue(timeout=None):
+                nonlocal emitted
+                if not emitted:
+                    emitted = True
+                    return message
+                executor._pump_stop.wait(0.01)
+                raise TimeoutError
+
+            return dequeue
+
+        for result_mq, message in zip(result_mqs, messages, strict=True):
+            result_mq.dequeue = make_dequeue(message)
+
+        executor._result_mqs = result_mqs
+        executor._result_mq = result_mqs[0]
+        executor._start_result_pump()
+        deadline = time.monotonic() + 2.0
+        while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
+            time.sleep(0.01)
+        executor._pump_stop.set()
+        for thread in executor._result_pump_threads:
+            thread.join(timeout=2.0)
+
+        assert len(executor._result_pump_threads) == 2
+        received = [executor._sync_result_buffer.get_nowait().output for _ in range(2)]
+        assert sorted(received) == ["rank0", "rank1"]
+
     def test_compute_done_routes_to_rpc_future(self):
         executor = _make_executor()
         rpc_id = "42"
@@ -233,6 +269,20 @@ class TestResultPumpDispatch:
 class TestShutdownCleansUpFutures:
     """Test that shutdown cancels pending futures."""
 
+    def test_shutdown_joins_result_pump_threads(self):
+        executor = _make_executor()
+        pump = threading.Thread(
+            target=executor._pump_stop.wait,
+            name="test-result-pump",
+        )
+        executor._result_pump_threads = [pump]
+        pump.start()
+
+        executor.shutdown()
+
+        assert not pump.is_alive()
+        assert executor._result_pump_threads == []
+
     def test_shutdown_sets_exception_on_pending_futures(self):
         executor = _make_executor()