From 2ec1fc3c9ca071cdcb6855a9bb8f56bb9effc03d Mon Sep 17 00:00:00 2001
From: Jiangyun Zhu <riverclouds.zhu@qq.com>
Date: Tue, 3 Feb 2026 15:03:54 +0800
Subject: [PATCH] [Bugfix] exit when oom (#1163)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
---
 vllm_omni/entrypoints/async_omni.py           | 24 +++++++++++++++++++
 vllm_omni/entrypoints/async_omni_diffusion.py | 18 ++++++--------
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index ab829beaa..2fd65d81a 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -471,6 +471,29 @@ class AsyncOmni(OmniBase):
             else:
                 logger.debug(f"[{self._name}] Request {req_id} fully completed")
 
+    def handle_oom(self, result: dict[str, Any], stage: OmniStage, stage_id: int) -> None:
+        if getattr(stage, "stage_type", None) == "diffusion":
+            error_text = str(result.get("error", ""))
+            error_lower = error_text.lower()
+            is_oom = "out of memory" in error_lower or "cuda oom" in error_lower
+            if is_oom:
+                logger.critical(
+                    f"[{self._name}] Diffusion stage {stage_id} reported OOM; stopping stage worker.",
+                )
+                try:
+                    stage.stop_stage_worker()
+                except Exception as e:
+                    logger.warning(
+                        f"[{self._name}] Failed to stop diffusion stage {stage_id}: {e}",
+                    )
+                logger.critical(
+                    f"[{self._name}] Exiting API process due to diffusion stage {stage_id} OOM.",
+                )
+
+                import os as _os
+
+                _os._exit(1)
+
     def _process_single_result(
         self,
         result: dict[str, Any],
@@ -493,6 +516,7 @@ class AsyncOmni(OmniBase):
             logger.error(
                 f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
             )
+            self.handle_oom(result, stage, stage_id)  # only temporary workaround for oom
             raise RuntimeError(result)
 
         engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py
index 535f04f7d..6f447c332 100644
--- a/vllm_omni/entrypoints/async_omni_diffusion.py
+++ b/vllm_omni/entrypoints/async_omni_diffusion.py
@@ -147,17 +147,13 @@ class AsyncOmniDiffusion:
 
         # Run engine in thread pool
         loop = asyncio.get_event_loop()
-        try:
-            # In async mode, only a single request is submitted at a time
-            result = await loop.run_in_executor(
-                self._executor,
-                self.engine.step,
-                request,
-            )
-            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
+        # In async mode, only a single request is submitted at a time
+        result = await loop.run_in_executor(
+            self._executor,
+            self.engine.step,
+            request,
+        )
+        result = result[0]
 
         # Update request_id if needed
         if not result.request_id: