diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 21167cb07..f37ac6ada 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/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 diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index 71d4aa95b..d6d1e0b30 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/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), + ]