[Bugfix] Keep the offline Omni engine alive after a non-fatal request error

tests/entrypoints/test_omni_entrypoints.py

@@ -1482,3 +1482,82 @@ def test_del_shutsdown_engine(monkeypatch: pytest.MonkeyPatch):
     del app
     gc.collect()
     assert engine.shutdown_called
+
+
+def _client_error_then_outputs(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
+    """First request is rejected with a 4xx client error; later ones complete normally."""
+    if len(engine.submitted) == 1:
+        _enqueue_client_error_message(engine, msg)
+    else:
+        _enqueue_omni_final_only_outputs(engine, msg)
+
+
+def test_omni_generate_client_error_keeps_engine_usable(monkeypatch: pytest.MonkeyPatch):
+    """Issue #8449: a non-fatal client error must fail only that generate() call.
+
+    The rejected request is aborted and the error propagates with its status
+    metadata, but the engine stays up so the next generate() succeeds, matching
+    AsyncOmni. Before the fix list-mode generate() closed the engine and the
+    second call died with janus.ShutDown.
+    """
+    sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        default_sampling_params_list=sampling_params,
+        on_add_request=_client_error_then_outputs,
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        with pytest.raises(OmniClientError) as exc_info:
+            app.generate(["blocked"], use_tqdm=False)
+
+        assert exc_info.value.status_code == 400
+        assert exc_info.value.error_type == "BadRequestError"
+        rejected_id = engine.submitted[0]["request_id"]
+        assert engine.aborted == [[rejected_id]]
+        assert rejected_id not in app.request_states
+        assert not engine.shutdown_called
+
+        outputs = app.generate(["p2"], use_tqdm=False)
+    finally:
+        app.shutdown()
+
+    assert [output.stage_id for output in outputs] == [0, 2]
+    assert len(engine.submitted) == 2
+
+
+def test_omni_generate_server_error_keeps_engine_usable(monkeypatch: pytest.MonkeyPatch):
+    """A non-fatal server-side error (RuntimeError, no 4xx status) also leaves the engine up."""
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_enqueue_error_message,
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    try:
+        with pytest.raises(RuntimeError, match="engine boom") as exc_info:
+            app.generate(["boom"], use_tqdm=False)
+        assert not isinstance(exc_info.value, OmniClientError)
+        assert not engine.shutdown_called
+        assert engine.aborted == [[engine.submitted[0]["request_id"]]]
+    finally:
+        app.shutdown()
+
+
+def test_omni_generate_fatal_error_still_closes_engine(monkeypatch: pytest.MonkeyPatch):
+    """A fatal ErrorMessage means the engine is unusable: generate() keeps closing it."""
+    engine = FakeAsyncOmniEngine(
+        stage_metadata=THREE_STAGE_META,
+        on_add_request=_enqueue_fatal_error_message,
+    )
+    _patch_engine(monkeypatch, engine)
+
+    app = Omni("dummy-model")
+    with pytest.raises(EngineDeadError, match="engine dead") as exc_info:
+        app.generate(["dead"], use_tqdm=False)
+
+    assert isinstance(exc_info.value, OmniEngineDeadError)
+    assert engine.shutdown_called

vllm_omni/entrypoints/omni.py

@@ -12,6 +12,7 @@
 from tqdm.auto import tqdm
 from vllm.logger import init_logger
 from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.messages import OutputMessage
@@ -95,9 +96,17 @@ def generate(
             if py_generator:
                 return self._run_generation_with_generator(prompts, sampling_params_list, use_tqdm)
             return list(self._run_generation(prompts, sampling_params_list, use_tqdm))
+        except EngineDeadError as e:
+            # Fatal: a stage or the orchestrator is gone, so the engine cannot
+            # serve anything else. Release the processes now.
+            logger.exception("[Omni] Engine died during generation: %s", e)
+            self.close()
+            raise
         except Exception as e:
+            # Request-scoped failure (e.g. OmniClientError -> 4xx). The in-flight
+            # requests were already aborted inside _run_generation; keep the
+            # engine alive so later generate() calls work, like AsyncOmni.
             logger.exception("[Omni] Failed to run generation: %s", e)
-            self.close()
             raise
 
     def _run_generation_with_generator(