[release/v0.18.0.post1][Bugfix] fix oom no response to user (#2638)

Signed-off-by: bjf-frz <frz123db@gmail.com>

vllm_omni/entrypoints/async_omni.py

@@ -21,6 +21,7 @@ from vllm.pooling_params import PoolingParams
 from vllm.tasks import SupportedTask
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import OmniBase
 from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
@@ -290,14 +291,15 @@ class AsyncOmni(EngineClient, OmniBase):
             stage_id = result.get("stage_id", 0)
 
             # Check for errors
-            if "error" in result:
-                logger.error(
-                    "[AsyncOmni] Orchestrator error for req=%s stage-%s: %s",
-                    request_id,
-                    stage_id,
-                    result["error"],
+            if isinstance(result, dict) and result.get("type") == "error":
+                raise OmniRequestError(
+                    result.get("error", "Unknown orchestrator error"),
+                    status_code=result.get("status_code", 500),
+                    request_id=result.get("request_id", request_id),
+                    stage_id=result.get("stage_id", stage_id),
+                    error_type=result.get("error_type"),
+                    detail=result.get("detail") or {},
                 )
-                raise RuntimeError(result)
 
             # Process the result (constructs OmniRequestOutput)
             output_to_yield = self._process_single_result(
@@ -344,6 +346,13 @@ class AsyncOmni(EngineClient, OmniBase):
                         await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S)
                         continue
 
+                    if isinstance(msg, dict) and msg.get("type") == "error":
+                        req_id = msg.get("request_id")
+                        req_state = self.request_states.get(req_id)
+                        if req_state is not None:
+                            await req_state.queue.put(msg)
+                        continue
+
                     should_continue, _, stage_id, req_state = self._handle_output_message(msg)
                     if should_continue:
                         continue
@@ -358,7 +367,16 @@ class AsyncOmni(EngineClient, OmniBase):
             except Exception as e:
                 logger.exception("[AsyncOmni] final_output_loop failed.")
                 for req_state in list(self.request_states.values()):
-                    error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    # error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    error_msg = {
+                        "type": "error",
+                        "request_id": req_state.request_id,
+                        "status_code": 500,
+                        "error": str(e),
+                        "error_type": type(e).__name__,
+                        "detail": {},
+                        "finished": True,
+                    }
                     await req_state.queue.put(error_msg)
                 self.final_output_task = None
 

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -18,7 +18,7 @@ from typing import Any
 from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_hf_file_to_dict
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig, normalize_omni_error
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
@@ -237,7 +237,8 @@ class AsyncOmniDiffusion:
             )
         except Exception as e:
             logger.error("Batch generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
+            # raise RuntimeError(f"Diffusion batch generation failed: {e}") from e
+            raise normalize_omni_error(e, request_id=request_id) from e
 
         # Combine all per-prompt results into a single OmniRequestOutput
         all_images = []
@@ -310,7 +311,8 @@ class AsyncOmniDiffusion:
             result = result[0]
         except Exception as e:
             logger.error("Generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion generation failed: {e}") from e
+            # raise RuntimeError(f"Diffusion generation failed: {e}") from e
+            raise normalize_omni_error(e, request_id=request_id) from e
 
         if not result.request_id:
             result.request_id = request_id

vllm_omni/entrypoints/omni_base.py

@@ -12,6 +12,7 @@ import huggingface_hub
 from vllm.logger import init_logger
 from vllm.v1.engine.exceptions import EngineDeadError
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
@@ -206,7 +207,14 @@ class OmniBase:
             return True, None, None, None
 
         if msg_type == "error":
-            raise RuntimeError(msg.get("error", "Orchestrator returned an error message"))
+            raise OmniRequestError(
+                msg.get("error", "Orchestrator returned an error message"),
+                status_code=msg.get("status_code", 500),
+                request_id=msg.get("request_id"),
+                stage_id=msg.get("stage_id"),
+                error_type=msg.get("error_type"),
+                detail=msg.get("detail") or {},
+            )
 
         if msg_type != "output":
             logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)

vllm_omni/entrypoints/openai/serving_video.py

@@ -13,6 +13,7 @@ from PIL import Image
 from vllm.engine.protocol import EngineClient
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.data import OmniRequestError
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoData,
@@ -214,12 +215,24 @@ class OmniOpenAIServingVideo:
         sampling_params_list: list[OmniSamplingParams] = [gen_params for _ in stage_configs]
 
         result = None
-        async for output in engine_client.generate(
-            prompt=prompt,
-            request_id=request_id,
-            sampling_params_list=sampling_params_list,
-        ):
-            result = output
+        try:
+            async for output in engine_client.generate(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params_list=sampling_params_list,
+            ):
+                result = output
+        except OmniRequestError as e:
+            raise HTTPException(
+                status_code=e.status_code,
+                detail={
+                    "message": str(e),
+                    "request_id": e.request_id,
+                    "stage_id": e.stage_id,
+                    "error_type": e.error_type,
+                    "detail": e.detail,
+                },
+            ) from e
 
         if result is None:
             raise HTTPException(