[Bugfix][CI] Reuse servers across paired perf cases (#5624)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

tests/dfx/perf/scripts/run_benchmark.py

@@ -1,6 +1,8 @@
 import json
 import os
 import threading
+from collections.abc import Callable
+from contextlib import AbstractContextManager, ExitStack, contextmanager
 from pathlib import Path
 from typing import Any
 
@@ -63,35 +65,78 @@ paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_C
 _omni_server_lock = threading.Lock()
 
 
+class _SingleActiveContext:
+    """Reuse one active context while its configuration key is unchanged."""
+
+    def __init__(self) -> None:
+        self._key: Any = None
+        self._stack: ExitStack | None = None
+        self._value: Any = None
+
+    def acquire(self, key: Any, factory: Callable[[], AbstractContextManager[Any]]) -> Any:
+        if self._stack is not None and key == self._key:
+            return self._value
+
+        self.close()
+        stack = ExitStack()
+        value = stack.enter_context(factory())
+        self._key = key
+        self._stack = stack
+        self._value = value
+        return value
+
+    def close(self) -> None:
+        stack = self._stack
+        self._key = None
+        self._stack = None
+        self._value = None
+        if stack is not None:
+            stack.close()
+
+
+@contextmanager
+def _start_omni_server(server_param):
+    test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = server_param
+
+    print(f"Starting OmniServer with test: {test_name}, model: {model}")
+
+    server_args: list[str] = []
+    if use_omni:
+        server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
+    # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
+    # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
+    if stage_config_path:
+        server_args = ["--deploy-config", stage_config_path] + server_args
+    if stage_overrides:
+        server_args = ["--stage-overrides", stage_overrides] + server_args
+    if extra_cli_args:
+        server_args = list(extra_cli_args) + server_args
+    with OmniServer(model, server_args, use_omni=use_omni) as server:
+        server.test_name = test_name
+        print("OmniServer started successfully")
+        yield server
+        print("OmniServer stopping...")
+
+    print("OmniServer stopped")
+
+
 @pytest.fixture(scope="module")
-def omni_server(request):
+def omni_server_context():
     """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
+    Reuse it for adjacent benchmark cases with the same server configuration.
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = request.param
+        active_context = _SingleActiveContext()
+        try:
+            yield active_context
+        finally:
+            active_context.close()
 
-        print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args: list[str] = []
-        if use_omni:
-            server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
-        # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
-        # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
-        if stage_config_path:
-            server_args = ["--deploy-config", stage_config_path] + server_args
-        if stage_overrides:
-            server_args = ["--stage-overrides", stage_overrides] + server_args
-        if extra_cli_args:
-            server_args = list(extra_cli_args) + server_args
-        with OmniServer(model, server_args, use_omni=use_omni) as server:
-            server.test_name = test_name
-            print("OmniServer started successfully")
-            yield server
-            print("OmniServer stopping...")
-
-        print("OmniServer stopped")
+@pytest.fixture
+def omni_server(request, omni_server_context):
+    return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param))
 
 
 @pytest.fixture

tests/dfx/perf/tests/test_runner_metadata.py

@@ -1,6 +1,8 @@
 """Tests for DFX runner metadata field exclusion."""
 
 import json
+from contextlib import contextmanager
+from types import SimpleNamespace
 
 import pytest
 
@@ -254,8 +256,9 @@ def test_benchmark_param_id_suffix_from_task_eval_phase():
     ]
 
 
-def test_create_paired_omni_benchmark_pytest_params(tmp_path):
+def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_path, monkeypatch):
     from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
+    from tests.dfx.perf.scripts import run_benchmark
 
     configs = [
         {
@@ -285,3 +288,51 @@ def test_create_paired_omni_benchmark_pytest_params(tmp_path):
     assert bench_row == ("test_tts", 0)
     assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks)
     assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks)
+
+    omni_p0_server = by_id["test_omni-p0"].values[0]
+    omni_p1_server = by_id["test_omni-p1"].values[0]
+    tts_server = by_id["test_tts-p0"].values[0]
+    assert omni_p0_server == omni_p1_server
+    assert omni_p0_server != tts_server
+
+    events = []
+
+    @contextmanager
+    def fake_start(server_param):
+        events.append(("start", server_param))
+        yield object()
+        events.append(("stop", server_param))
+
+    monkeypatch.setattr(run_benchmark, "_start_omni_server", fake_start)
+    active_context = run_benchmark._SingleActiveContext()
+    try:
+        first = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=omni_p0_server),
+            active_context,
+        )
+        second = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=omni_p1_server),
+            active_context,
+        )
+        assert first is second
+        assert events == [("start", omni_p0_server)]
+
+        third = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=tts_server),
+            active_context,
+        )
+        assert third is not first
+        assert events == [
+            ("start", omni_p0_server),
+            ("stop", omni_p0_server),
+            ("start", tts_server),
+        ]
+    finally:
+        active_context.close()
+
+    assert events == [
+        ("start", omni_p0_server),
+        ("stop", omni_p0_server),
+        ("start", tts_server),
+        ("stop", tts_server),
+    ]