[Bugfix] Release GPU memory after diffusion execution failures (#6385)

Signed-off-by: nagisa <1434936049@qq.com>
Signed-off-by: nagisa-kun <1434936049@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/test_diffusion_worker_broadcast.py

@@ -71,6 +71,45 @@ def test_recv_message_waits_indefinitely():
     worker_proc.mq.dequeue.assert_called_once_with(indefinite=True)
 
 
+def test_busy_loop_releases_failed_rpc_traceback_and_device_cache(monkeypatch):
+    original = RuntimeError("model execution failed")
+    worker_proc = WorkerProc.__new__(WorkerProc)
+    worker_proc.gpu_id = 0
+    worker_proc.mq = Mock()
+    worker_proc.mq.dequeue.side_effect = [
+        {
+            "type": "rpc",
+            "method": "execute_model",
+            "args": (),
+            "kwargs": {},
+            "output_rank": 0,
+            "exec_all_ranks": False,
+            "collect_rank_status": False,
+        },
+        {"type": "shutdown"},
+    ]
+    worker_proc.worker = Mock()
+    worker_proc.worker.execute_method.side_effect = original
+    worker_proc.result_mq = object()
+    worker_proc._return_result = Mock()
+    worker_proc.wake_event = None
+    worker_proc._running = True
+
+    gc_collect = Mock()
+    mock_platform = Mock()
+    monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect)
+    monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform)
+
+    worker_proc._worker_busy_loop()
+
+    worker_proc._return_result.assert_called_once_with(
+        {"status": "error", "error": "model execution failed", "wave_id": None}
+    )
+    assert original.__traceback__ is None
+    gc_collect.assert_called_once_with()
+    mock_platform.empty_cache.assert_called_once_with()
+
+
 def test_busy_loop_preserves_order_when_overflow_payload_is_delayed(monkeypatch):
     """A consumed overflow marker must not be abandoned on a socket timeout."""
     writer = MessageQueue(

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -14,6 +14,7 @@ import torch
 import zmq
 from vllm.v1.engine.exceptions import EngineDeadError
 
+import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
@@ -822,6 +823,31 @@ class TestWorkerProcRpcRankStatus:
         assert status["bool_result"] is None
         assert "local boom" in status["traceback"]
 
+    def test_execute_rpc_collect_exception_releases_traceback_and_device_cache(self, monkeypatch):
+        proc = self._make_worker_proc()
+        original = RuntimeError("local boom")
+        proc.worker.execute_method = Mock(side_effect=original)
+        gc_collect = Mock()
+        mock_platform = Mock()
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+        monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect)
+        monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform)
+
+        proc._execute_rpc(
+            {
+                "method": "add_lora",
+                "args": (),
+                "kwargs": {},
+                "output_rank": 0,
+                "exec_all_ranks": True,
+                "collect_rank_status": True,
+            }
+        )
+
+        assert original.__traceback__ is None
+        gc_collect.assert_called_once_with()
+        mock_platform.empty_cache.assert_called_once_with()
+
     def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self):
         proc = self._make_worker_proc()
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -84,6 +84,16 @@ _ASYNC_OUTPUT_DRAIN_TIMEOUT_S = 10.0
 _MEMORY_RELEASING_METHODS = frozenset({"sleep", "handle_sleep_task"})
 
 
+def _cleanup_after_execution_error(exc: Exception) -> None:
+    """Release device tensors retained by a failed execution traceback."""
+    exc.__traceback__ = None
+    try:
+        gc.collect()
+        current_omni_platform.empty_cache()
+    except Exception:
+        logger.warning("Failed to release device memory after an execution error", exc_info=True)
+
+
 def _all_gather_rank_values(value: Any) -> list[Any]:
     if not dist.is_available() or not dist.is_initialized():
         return [value]
@@ -1187,7 +1197,6 @@ class WorkerProc:
             return None, False
 
         result = None
-        rpc_exception: Exception | None = None
         status: dict[str, Any] = {
             "rank": self.gpu_id,
             "ok": True,
@@ -1204,7 +1213,6 @@ class WorkerProc:
             result = self.worker.execute_method(method, *args, **kwargs)
         except Exception as e:
             logger.error(f"Error executing RPC: {e}", exc_info=True)
-            rpc_exception = e
             status.update(
                 {
                     "ok": False,
@@ -1213,6 +1221,9 @@ class WorkerProc:
                     "traceback": traceback.format_exc(),
                 }
             )
+            if not collect_rank_status:
+                raise
+            _cleanup_after_execution_error(e)
 
         if isinstance(result, bool):
             status["bool_result"] = result
@@ -1232,8 +1243,6 @@ class WorkerProc:
                 )
             return None, False
 
-        if rpc_exception is not None:
-            raise rpc_exception
         if isinstance(result, dict) and wave_id is not None:
             result["wave_id"] = wave_id
         return result, should_reply
@@ -1282,6 +1291,8 @@ class WorkerProc:
                         self._return_result(result, rpc_id=rpc_id)
                 except Exception as e:
                     logger.error(f"Error processing RPC: {e}", exc_info=True)
+                    error = str(e)
+                    _cleanup_after_execution_error(e)
                     # Apply the same reply gate as the success path so
                     # non-output ranks don't enqueue stale error replies
                     # that compete with the expected responder's message.
@@ -1296,7 +1307,7 @@ class WorkerProc:
                                 AsyncDiffusionOutput(
                                     kind=AsyncOutputKind.RPC_RESULT,
                                     rpc_id=rpc_id,
-                                    error=str(e),
+                                    error=error,
                                 )
                             )
                         elif output_rank is None and exec_all_ranks:
@@ -1324,11 +1335,11 @@ class WorkerProc:
                                 except Exception:
                                     dp_rank = self.gpu_id
                                 self._return_result(
-                                    {"status": "error", "error": str(e), "dp_rank": dp_rank, "wave_id": wave_id}
+                                    {"status": "error", "error": error, "dp_rank": dp_rank, "wave_id": wave_id}
                                 )
                         elif output_rank is None or output_rank == self.gpu_id:
                             # Normal RPC: only the expected rank replies
-                            self._return_result({"status": "error", "error": str(e), "wave_id": wave_id})
+                            self._return_result({"status": "error", "error": error, "wave_id": wave_id})
 
             elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                 logger.info("Worker %s: Received shutdown message", self.gpu_id)
@@ -1345,6 +1356,7 @@ class WorkerProc:
                         exc_info=True,
                     )
                     output = DiffusionOutput.from_exception(e)
+                    _cleanup_after_execution_error(e)
 
                 try:
                     self._return_result(output)