From 01f4479138c308a9ccf0ac84403b8b27fab0d3ae Mon Sep 17 00:00:00 2001
From: Alex Brooks <albrooks@redhat.com>
Date: Mon, 14 Sep 2026 16:33:38 -0600
Subject: [PATCH] [Bugfix] Don't shutdown Engine on py_generator=True (#6334)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
---
 tests/entrypoints/test_omni_entrypoints.py | 61 ++++++++++++++++++++--
 vllm_omni/entrypoints/omni.py              | 17 ++++--
 2 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index be7a0ffac..95059a06e 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import gc
 import queue
 from collections.abc import Callable
 from types import SimpleNamespace
@@ -766,7 +767,7 @@ def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkey
         f"{engine.submitted[1]['request_id']}-stage0-0",
         f"{engine.submitted[1]['request_id']}-stage2-final",
     ]
-    assert engine.shutdown_called is True
+    assert engine.shutdown_called is not True
 
 
 def test_omni_generate_returns_list_when_not_using_generator(monkeypatch: pytest.MonkeyPatch):
@@ -812,7 +813,7 @@ def test_omni_generate_diffusion_only_yields_single_image_per_request(monkeypatc
         [f"{engine.submitted[0]['request_id']}-image"],
         [f"{engine.submitted[1]['request_id']}-image"],
     ]
-    assert engine.shutdown_called is True
+    assert engine.shutdown_called is not True
 
 
 def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request(
@@ -845,7 +846,7 @@ def test_omni_generate_llm_diffusion_yields_final_text_then_image_per_request(
         [f"{engine.submitted[1]['request_id']}-image"],
     ]
     assert engine.submitted[0]["sampling_params_list"][0].output_kind == RequestOutputKind.FINAL_ONLY
-    assert engine.shutdown_called is True
+    assert engine.shutdown_called is not True
 
 
 def test_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
@@ -1326,3 +1327,57 @@ def test_omni_errored_property_dead_stage(monkeypatch: pytest.MonkeyPatch):
         assert app.errored is False
     finally:
         app.shutdown()
+
+
+def test_omni_pygenerator_does_not_kill_engine(monkeypatch: pytest.MonkeyPatch):
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_enqueue_async_three_stage_outputs,
+    )
+
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    # Evaluating the generator should not shut down the engine
+    list(app.generate(["hello"], py_generator=True))
+    assert not engine.shutdown_called
+
+
+def test_omni_generator_close_cleans_up(monkeypatch: pytest.MonkeyPatch):
+    """Ensure that a closed generator cleans things up properly."""
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_enqueue_async_three_stage_outputs,
+    )
+
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+
+    # Create a generator and start to evaluate it to make sure the request isn't aborted yet
+    my_gen = app.generate(["hello"], py_generator=True)
+    next(my_gen)
+    request_id = engine.submitted[0]["request_id"]
+    assert engine.aborted == []
+    assert request_id in app.request_states
+
+    # Close it and make sure the sure it's aborted, but without killing engine
+    my_gen.close()
+    assert engine.aborted == [[request_id]]
+    assert request_id not in app.request_states
+    assert not engine.shutdown_called
+
+
+def test_del_shutsdown_engine(monkeypatch: pytest.MonkeyPatch):
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_enqueue_async_three_stage_outputs,
+    )
+
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    assert not engine.shutdown_called
+    del app
+    gc.collect()
+    assert engine.shutdown_called
diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py
index 28bbe9219..cf4da3b7a 100644
--- a/vllm_omni/entrypoints/omni.py
+++ b/vllm_omni/entrypoints/omni.py
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import copy
@@ -99,11 +102,11 @@ class Omni(OmniBase):
         sampling_params_list: Sequence[OmniSamplingParams],
         use_tqdm: bool | Callable[..., tqdm] = True,
     ) -> Generator[OmniRequestOutput, None, None]:
-        gen = self._run_generation(prompts, sampling_params_list, use_tqdm)
-        try:
-            yield from gen
-        finally:
-            self.close()
+        yield from self._run_generation(
+            prompts,
+            sampling_params_list,
+            use_tqdm,
+        )
 
     def _run_generation(
         self,
@@ -198,6 +201,10 @@ class Omni(OmniBase):
                     if pbar is not None:
                         pbar.update(1)
                     self._log_summary_and_cleanup(req_id)
+        except GeneratorExit:
+            if "active_reqs" in locals() and active_reqs:
+                self.abort(list(active_reqs))
+            raise
         except Exception:
             if "active_reqs" in locals() and active_reqs:
                 for req_id in active_reqs: