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

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

tests/e2e/accuracy/sana_video/conftest.py

@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-import pytest
-
-
-@pytest.fixture
-def sana_model_parallel(tmp_path):
-    """Construct native parallel layers with real single-rank GPU groups."""
-    from vllm_omni.diffusion.distributed.parallel_state import (
-        destroy_distributed_env,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
-    )
-    try:
-        initialize_model_parallel()
-        yield
-    finally:
-        destroy_distributed_env()

tests/e2e/accuracy/sana_video/test_sana_video_golden.py

@@ -202,7 +202,6 @@ def test_sana_video_transformer_matches_frozen_golden(
     variant: str,
     _hardware,
     accuracy_artifact_root: Path,
-    sana_model_parallel,
 ) -> None:
     del _hardware
     from vllm.config import LoadConfig

tests/e2e/accuracy/sana_video/test_sana_video_pipeline_alignment.py

@@ -189,7 +189,6 @@ def _cfg_prediction(
 def test_sana_video_quick_pipeline_stages_match_diffusers(
     variant: str,
     _hardware,
-    sana_model_parallel,
 ) -> None:
     """Compare the first denoising step without replacing native attention.
 

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):