From 04ccb062fcf3fb48a4247451b968aadf076adf0d Mon Sep 17 00:00:00 2001
From: tzhouam <tzhouam@connect.ust.hk>
Date: Tue, 24 Feb 2026 06:58:39 +0000
Subject: [PATCH] Add signature compatibility tests for OmniLLM class

- Introduced tests to ensure method signatures in OmniLLM match those of the base LLM class, preventing runtime TypeErrors due to parameter mismatches.
- Updated the _run_engine method in OmniLLM to include an optional output_type parameter.

Signed-off-by: [Your Name] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
---
 tests/entrypoints/test_omni_llm.py | 57 ++++++++++++++++++++++++++++++
 vllm_omni/entrypoints/omni_llm.py  |  4 ++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py
index 4f05575ca..b3ad4ce01 100644
--- a/tests/entrypoints/test_omni_llm.py
+++ b/tests/entrypoints/test_omni_llm.py
@@ -1,3 +1,4 @@
+import inspect
 import uuid
 import warnings
 from queue import Empty, Queue
@@ -997,3 +998,59 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
 
     # Verify stop_stage_worker was called (process should be set)
     assert omni.stage_list[0]._proc is not None
+
+
+# ---------------------------------------------------------------------------
+# Signature compatibility tests — catch upstream API drift early
+# ---------------------------------------------------------------------------
+
+_INHERITANCE_PAIRS: list[tuple[str, str, str, str]] = [
+    (
+        "vllm.entrypoints.llm",
+        "LLM",
+        "vllm_omni.entrypoints.omni_llm",
+        "OmniLLM",
+    ),
+]
+
+
+def _import_class(module_path: str, class_name: str):
+    mod = __import__(module_path, fromlist=[class_name])
+    return getattr(mod, class_name)
+
+
+@pytest.mark.parametrize(
+    "up_mod,up_cls,omni_mod,omni_cls",
+    _INHERITANCE_PAIRS,
+    ids=[f"{pair[3]}({pair[1]})" for pair in _INHERITANCE_PAIRS],
+)
+def test_overridden_method_signatures_compatible(up_mod: str, up_cls: str, omni_mod: str, omni_cls: str):
+    """All params accepted by a base-class method must also be accepted by
+    the overriding method in the vllm-omni subclass.  This prevents
+    TypeError at runtime when upstream callers pass new arguments."""
+    BaseCls = _import_class(up_mod, up_cls)
+    OmniCls = _import_class(omni_mod, omni_cls)
+
+    failures: list[str] = []
+    for name in dir(OmniCls):
+        if name.startswith("__") and name != "__init__":
+            continue
+        omni_method = getattr(OmniCls, name, None)
+        base_method = getattr(BaseCls, name, None)
+        if not (callable(omni_method) and callable(base_method)):
+            continue
+        if omni_method is base_method:
+            continue
+        try:
+            base_sig = inspect.signature(base_method)
+            omni_sig = inspect.signature(omni_method)
+        except (ValueError, TypeError):
+            continue
+
+        base_params = set(base_sig.parameters.keys())
+        omni_params = set(omni_sig.parameters.keys())
+        missing = base_params - omni_params
+        if missing:
+            failures.append(f"{omni_cls}.{name}() missing params {sorted(missing)}; base={base_sig}, omni={omni_sig}")
+
+    assert not failures, f"Signature mismatches found ({len(failures)}):\n" + "\n".join(f"  - {f}" for f in failures)
diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py
index faf9bfce0..fbb38860f 100644
--- a/vllm_omni/entrypoints/omni_llm.py
+++ b/vllm_omni/entrypoints/omni_llm.py
@@ -222,7 +222,9 @@ class OmniLLM(LLM):
         except Exception as e:
             logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
 
-    def _run_engine(self, *, use_tqdm: bool | Callable[..., tqdm] = True) -> list[RequestOutput | PoolingRequestOutput]:
+    def _run_engine(
+        self, output_type=None, *, use_tqdm: bool | Callable[..., tqdm] = True
+    ) -> list[RequestOutput | PoolingRequestOutput]:
         # Initialize tqdm.
         if use_tqdm:
             num_requests = self.llm_engine.get_num_unfinished_requests()