[Test] Restore stage-client class identities after the module reload test

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/e2e/online_serving/test_qwen2_5_omni_expansion.py

@@ -21,13 +21,7 @@ models = ["Qwen/Qwen2.5-Omni-7B"]
 
 # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
 # platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml.
-stage_configs = [
-    modify_stage_config(
-        get_deploy_config_path("ci/qwen2_5_omni.yaml"),
-        # This suite requests two cards; Code2Wav shares the Talker's card.
-        updates={"stages": {2: {"devices": "1"}}},
-    )
-]
+stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))]
 
 # Create parameter combinations for model and stage config
 test_params = [

tests/engine/test_async_omni_engine_stage_init.py

@@ -268,11 +268,21 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
 
-    importlib.reload(client_mod)
+    # Snapshot the module namespace so a reload cannot strand other test
+    # modules that captured the pre-reload class objects (isinstance checks
+    # against the stale class would then fail). ``importlib.reload`` mutates
+    # the module dict in place; restoring it leaves the runtime exactly as it
+    # was before this test ran.
+    saved = dict(client_mod.__dict__)
+    try:
+        importlib.reload(client_mod)
 
-    assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
-        "StageEngineCoreClient | DPLBStageEngineCoreClient"
-    )
+        assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
+            "StageEngineCoreClient | DPLBStageEngineCoreClient"
+        )
+    finally:
+        client_mod.__dict__.clear()
+        client_mod.__dict__.update(saved)
 
 
 def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):