[Bugfix] Keep client error status in the single-GPU diffusion executor (#8474)

Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>

tests/diffusion/test_diffusion_model_runner.py

@@ -1452,3 +1452,27 @@ def test_execute_model_batch_cancellation_preserves_live_peer(monkeypatch, cance
             assert pipeline.last_batch is None
     finally:
         registry.close()
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_non_step_fallback_keeps_client_error_status():
+    """Step mode's full-forward fallback must not drop a request error's 4xx status."""
+    from vllm_omni.errors import OmniClientError
+
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.execute_model = Mock(side_effect=OmniClientError("bad request option", status_code=400))
+    request = _make_request()
+    scheduler_output = SimpleNamespace(
+        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request, diffusion_kv_metadata=None)],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+        kv_prefetch_job=None,
+    )
+
+    result = DiffusionModelRunner._execute_non_step_requests(runner, scheduler_output)
+
+    output = result.get_request_output(request.request_id)
+    assert output.finished is True
+    assert output.result.error == "bad request option"
+    assert output.result.error_status_code == 400
+    assert output.result.error_type == "BadRequestError"

tests/diffusion/test_uniproc_executor.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
 
 """Single-GPU (uniproc) diffusion executor: selection and RPC behaviour."""
 
@@ -260,3 +260,52 @@ def test_shutdown_is_idempotent_and_closes_executor(executor, monkeypatch):
     empty_cache.assert_called_once_with()
     with pytest.raises(RuntimeError, match="closed"):
         exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+
+def _single_request_wave(request_id: str = "req-1"):
+    return SimpleNamespace(
+        scheduled_new_reqs=[
+            SimpleNamespace(
+                request_id=request_id,
+                req=SimpleNamespace(request_id=request_id),
+                diffusion_kv_metadata=None,
+            )
+        ],
+        kv_prefetch_job=None,
+    )
+
+
+def test_execute_request_keeps_client_error_status(executor, monkeypatch):
+    """A pipeline's OmniClientError must keep its 4xx status on a single GPU.
+
+    The worker runs inline, so the exception reaches execute_request directly.
+    Flattening it to DiffusionOutput(error=str(exc)) made the engine rebuild a
+    RuntimeError and the API answer HTTP 500 for a request validation error.
+    """
+    from vllm_omni.errors import OmniClientError
+
+    exec_, worker = executor
+    worker.execute_method.side_effect = OmniClientError("clip exceeds the frame budget", status_code=422)
+    monkeypatch.setattr(exec_, "_device_is_usable", lambda: True)
+
+    batch = exec_.execute_request(_single_request_wave())
+
+    (runner_output,) = batch.runner_outputs
+    assert runner_output.finished is True
+    assert runner_output.result.error == "clip exceeds the frame budget"
+    assert runner_output.result.error_status_code == 422
+    assert runner_output.result.error_type == "BadRequestError"
+    assert exec_.is_dead is False
+
+
+def test_execute_request_server_error_has_no_client_status(executor, monkeypatch):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("CUDA out of memory")
+    monkeypatch.setattr(exec_, "_device_is_usable", lambda: True)
+
+    batch = exec_.execute_request(_single_request_wave())
+
+    (runner_output,) = batch.runner_outputs
+    assert runner_output.result.error == "CUDA out of memory"
+    assert runner_output.result.error_status_code is None
+    assert runner_output.result.error_type is None

vllm_omni/diffusion/executor/uniproc_executor.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
 
 """In-process diffusion executor for single-GPU deployments.
 
@@ -143,12 +143,14 @@ class UniProcDiffusionExecutor(DiffusionExecutor):
                     )
                 )
             except Exception as exc:
+                # The worker runs inline, so a pipeline's OmniClientError lands
+                # here; from_exception keeps its 4xx status for the engine/API.
                 runner_outputs.append(
                     RunnerOutput(
                         request_id=new_req.request_id,
                         step_index=None,
                         finished=True,
-                        result=DiffusionOutput(error=str(exc)),
+                        result=DiffusionOutput.from_exception(exc),
                     )
                 )
         return BatchRunnerOutput.from_list(runner_outputs)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -1275,7 +1275,8 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
                     new_req.request_id,
                     exc_info=True,
                 )
-                result = DiffusionOutput(error=str(exc))
+                # Keep client-error metadata, matching the stepwise path below.
+                result = DiffusionOutput.from_exception(exc)
 
             step_index = getattr(new_req.req.sampling_params, "step_index", None)
             runner_outputs.append(