From 855719e3318bcc96ed070a6d331a19532f47c7d0 Mon Sep 17 00:00:00 2001 From: hurukawa <61525444+nagisa-kunhah@users.noreply.github.com> Date: Sun, 23 Aug 2026 02:32:48 +0800 Subject: [PATCH] [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 --- .../test_diffusion_worker_broadcast.py | 39 +++++++++++++++++++ .../test_multiproc_engine_concurrency.py | 26 +++++++++++++ .../diffusion/worker/diffusion_worker.py | 26 +++++++++---- 3 files changed, 84 insertions(+), 7 deletions(-) diff --git a/tests/diffusion/test_diffusion_worker_broadcast.py b/tests/diffusion/test_diffusion_worker_broadcast.py index 03f84b819..750d4c30d 100644 --- a/tests/diffusion/test_diffusion_worker_broadcast.py +++ b/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( diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index d5c36a093..d3298844f 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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() diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 9344cb9b4..25bca55e1 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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)