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

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

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 ─────────