[CI][Bugfix] Refactor the test case to add support for increasing init timeout and stage init timeout in order to resolve the CI timeout error. (#2711)

Signed-off-by: wangyu <410167048@qq.com>

tests/e2e/offline_inference/test_cache_dit.py

@@ -8,27 +8,15 @@ This test verifies that cache-dit acceleration works correctly with diffusion mo
 It uses minimal settings to keep test time short for CI.
 """
 
-import os
-import sys
-from pathlib import Path
-
 import pytest
 import torch
 
+from tests.conftest import OmniRunner
 from tests.utils import hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from vllm_omni import Omni
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
 # Use random weights model for testing
 models = ["riverclouds/qwen_image_random"]
 
@@ -48,20 +36,17 @@ def test_cache_dit(model_name: str):
         "residual_diff_threshold": 0.24,
         "max_continuous_cached_steps": 3,
     }
-    m = None
-    try:
-        m = Omni(
-            model=model_name,
-            cache_backend="cache_dit",
-            cache_config=cache_config,
-        )
-
+    with OmniRunner(
+        model_name,
+        cache_backend="cache_dit",
+        cache_config=cache_config,
+    ) as runner:
         # Use minimal settings for fast testing
         height = 256
         width = 256
         num_inference_steps = 4  # Minimal steps for fast test
 
-        outputs = m.generate(
+        outputs = runner.omni.generate(
             "a photo of a cat sitting on a laptop keyboard",
             OmniDiffusionSamplingParams(
                 height=height,
@@ -90,9 +75,3 @@ def test_cache_dit(model_name: str):
         # Check image size
         assert images[0].width == width
         assert images[0].height == height
-    except Exception as e:
-        print(f"Test failed with error: {e}")
-        raise
-    finally:
-        if m is not None and hasattr(m, "close"):
-            m.close()