diff --git a/e2e_test/bindings_go/conftest.py b/e2e_test/bindings_go/conftest.py index 64e2c0dc40..8184577e63 100644 --- a/e2e_test/bindings_go/conftest.py +++ b/e2e_test/bindings_go/conftest.py @@ -9,10 +9,7 @@ import logging import os -import signal -import socket import subprocess -import time from collections.abc import Generator from pathlib import Path from typing import TYPE_CHECKING @@ -22,6 +19,9 @@ if TYPE_CHECKING: from infra import ModelInstance, ModelPool +from infra import get_open_port, release_port, terminate_process +from infra.process_utils import wait_for_health + logger = logging.getLogger(__name__) # Paths @@ -30,25 +30,6 @@ _GO_OAI_SERVER = _GO_BINDINGS / "examples" / "oai_server" -def _find_free_port() -> int: - """Find an available port.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - return s.getsockname()[1] - - -def _wait_for_server(host: str, port: int, timeout: float = 30.0) -> bool: - """Wait for server to become available.""" - start = time.time() - while time.time() - start < timeout: - try: - with socket.create_connection((host, port), timeout=1.0): - return True - except (TimeoutError, ConnectionRefusedError, OSError): - time.sleep(0.5) - return False - - @pytest.fixture(scope="session") def go_ffi_library() -> Path: """Build the Go FFI library and return its directory path.""" @@ -241,7 +222,7 @@ def go_oai_server( grpc_endpoint = f"grpc://localhost:{grpc_worker.port}" # Find a free port for the Go OAI server - oai_port = _find_free_port() + oai_port = get_open_port() # Set up environment - the Go OAI server uses env vars for config env = os.environ.copy() @@ -261,17 +242,26 @@ def go_oai_server( cmd = [str(go_oai_binary)] - process = subprocess.Popen( - cmd, - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - try: + process = subprocess.Popen( + cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + # Wait for server to start - if not _wait_for_server("localhost", oai_port, timeout=30.0): - stdout, stderr = process.communicate(timeout=5) + try: + wait_for_health(f"http://localhost:{oai_port}", timeout=30.0, check_interval=0.5) + except TimeoutError: + try: + stdout, stderr = process.communicate(timeout=5) + except subprocess.TimeoutExpired: + terminate_process(process, timeout=10) + pytest.fail( + f"Go OAI server failed to start and did not exit cleanly.\n" + f"Command: {' '.join(cmd)}" + ) pytest.fail( f"Go OAI server failed to start.\n" f"Command: {' '.join(cmd)}\n" @@ -283,14 +273,9 @@ def go_oai_server( yield ("localhost", oai_port, grpc_worker.model_path) finally: - # Shutdown the server logger.info("Shutting down Go OAI server...") - process.send_signal(signal.SIGTERM) - try: - process.wait(timeout=10) - except subprocess.TimeoutExpired: - process.kill() - process.wait() + terminate_process(process, timeout=10) + release_port(oai_port) @pytest.fixture(scope="class") @@ -318,7 +303,7 @@ def go_oai_server_multi( grpc_endpoints = ",".join(f"grpc://localhost:{w.port}" for w in grpc_workers) # Find a free port for the Go OAI server - oai_port = _find_free_port() + oai_port = get_open_port() # Set up environment env = os.environ.copy() @@ -350,17 +335,26 @@ def go_oai_server_multi( cmd = [str(go_oai_binary)] - process = subprocess.Popen( - cmd, - env=env, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - try: + process = subprocess.Popen( + cmd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + # Wait for server to start - if not _wait_for_server("localhost", oai_port, timeout=60.0): - stdout, stderr = process.communicate(timeout=5) + try: + wait_for_health(f"http://localhost:{oai_port}", timeout=60.0, check_interval=0.5) + except TimeoutError: + try: + stdout, stderr = process.communicate(timeout=5) + except subprocess.TimeoutExpired: + terminate_process(process, timeout=10) + pytest.fail( + f"Go OAI server failed to start and did not exit cleanly.\n" + f"Command: {' '.join(cmd)}" + ) pytest.fail( f"Go OAI server failed to start.\n" f"Command: {' '.join(cmd)}\n" @@ -376,12 +370,8 @@ def go_oai_server_multi( finally: logger.info("Shutting down Go OAI server...") - process.send_signal(signal.SIGTERM) - try: - process.wait(timeout=10) - except subprocess.TimeoutExpired: - process.kill() - process.wait() + terminate_process(process, timeout=10) + release_port(oai_port) @pytest.fixture(scope="class") diff --git a/e2e_test/conftest.py b/e2e_test/conftest.py index 6c32dcd7d1..8d0b6198f0 100644 --- a/e2e_test/conftest.py +++ b/e2e_test/conftest.py @@ -213,8 +213,6 @@ def pytest_runtest_logstart(nodeid: str, location: tuple) -> None: from smg_client import SmgClient from smg_client._errors import SmgError -logger = logging.getLogger(__name__) - @pytest.fixture def smg(setup_backend): diff --git a/e2e_test/fixtures/__init__.py b/e2e_test/fixtures/__init__.py index a3bbeae76f..3b4d06e751 100644 --- a/e2e_test/fixtures/__init__.py +++ b/e2e_test/fixtures/__init__.py @@ -6,9 +6,6 @@ - setup_backend.py: Backend setup fixtures (class/function-scoped) - markers.py: Helper utilities for marker extraction -Legacy modules (to be removed during e2e_response_api migration): -- ports.py: Use infra.get_open_port() instead -- router_manager.py: Use infra.Gateway instead """ # Pytest hooks (imported by conftest.py via pytest_plugins) diff --git a/e2e_test/fixtures/pool.py b/e2e_test/fixtures/pool.py index ff8ae0c1c6..c5c01ae5b0 100644 --- a/e2e_test/fixtures/pool.py +++ b/e2e_test/fixtures/pool.py @@ -16,7 +16,7 @@ import pytest if TYPE_CHECKING: - from infra import ModelPool + from infra import ModelInstance, ModelPool from .hooks import get_pool_requirements @@ -170,44 +170,60 @@ def model_pool(request: pytest.FixtureRequest) -> ModelPool: return _model_pool -@pytest.fixture -def model_client( - request: pytest.FixtureRequest, model_pool: ModelPool -) -> Generator[object, None, None]: - """Get OpenAI client for the model specified by @pytest.mark.model(). +def _get_model_instance( + request: pytest.FixtureRequest, model_pool: ModelPool, fixture_name: str +) -> ModelInstance: + """Extract model from marker and acquire instance from pool. - Usage: - @pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct") - def test_chat(model_client): - response = model_client.chat.completions.create(...) + Args: + request: Pytest fixture request. + model_pool: The model pool. + fixture_name: Name of the calling fixture (for error messages). + + Returns: + Acquired ModelInstance. """ - import openai - from infra import PARAM_MODEL + from infra import PARAM_MODEL, ConnectionMode marker = request.node.get_closest_marker(PARAM_MODEL) - if marker is None: + if marker is None or not marker.args: pytest.fail( f"Test must be marked with @pytest.mark.{PARAM_MODEL}('model-id') " - "to use model_client fixture" + f"to use {fixture_name} fixture" ) model_id = marker.args[0] try: - # get() auto-acquires the returned instance - instance = model_pool.get(model_id) + return model_pool.get(model_id, ConnectionMode.HTTP) except KeyError: pytest.skip(f"Model {model_id} not available in model pool") + +@pytest.fixture +def model_client( + request: pytest.FixtureRequest, model_pool: ModelPool +) -> Generator[object, None, None]: + """Get OpenAI client for the model specified by @pytest.mark.model(). + + Usage: + @pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct") + def test_chat(model_client): + response = model_client.chat.completions.create(...) + """ + import openai + + instance = _get_model_instance(request, model_pool, "model_client") + client = openai.OpenAI( base_url=f"{instance.base_url}/v1", api_key="not-used", ) - yield client - - # Release reference to allow eviction - instance.release() + try: + yield client + finally: + instance.release() @pytest.fixture @@ -221,24 +237,9 @@ def model_base_url( def test_direct_http(model_base_url): response = httpx.get(f"{model_base_url}/health") """ - from infra import PARAM_MODEL - - marker = request.node.get_closest_marker(PARAM_MODEL) - if marker is None: - pytest.fail( - f"Test must be marked with @pytest.mark.{PARAM_MODEL}('model-id') " - "to use model_base_url fixture" - ) - - model_id = marker.args[0] + instance = _get_model_instance(request, model_pool, "model_base_url") try: - # get() auto-acquires the returned instance - instance = model_pool.get(model_id) - except KeyError: - pytest.skip(f"Model {model_id} not available in model pool") - - yield instance.base_url - - # Release reference to allow eviction - instance.release() + yield instance.base_url + finally: + instance.release() diff --git a/e2e_test/fixtures/ports.py b/e2e_test/fixtures/ports.py deleted file mode 100644 index b04703e68a..0000000000 --- a/e2e_test/fixtures/ports.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Legacy port utilities. - -DEPRECATED: This module will be removed during e2e_response_api migration. -Use infra.get_open_port() instead. -""" - -import socket - - -def find_free_port() -> int: - """Return an available TCP port on localhost.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index f49cf027f3..84e7af3bb3 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -59,6 +59,15 @@ def __init__(self, gen, value, cls, param): _cache_lock = threading.Lock() +def _release_workers(workers: list) -> None: + """Release a list of workers, logging warnings on failure.""" + for w in workers: + try: + w.release() + except Exception as e: + logger.warning("Failed to release worker during cleanup: %s", e) + + def _create_backend(request: pytest.FixtureRequest, model_pool: ModelPool): """Extract configuration from request and return the appropriate backend generator. @@ -96,10 +105,10 @@ def _create_backend(request: pytest.FixtureRequest, model_pool: ModelPool): # PD disaggregation backends - explicit connection modes if backend_name == "pd_http": - return _setup_pd_http_backend(request, model_pool, model_id, workers_config, gateway_config) + return _setup_pd_http_backend(model_pool, model_id, workers_config, gateway_config) if backend_name == "pd_grpc": - return _setup_pd_grpc_backend(request, model_pool, model_id, workers_config, gateway_config) + return _setup_pd_grpc_backend(model_pool, model_id, workers_config, gateway_config) # Check if this is a local backend (grpc, http) try: @@ -256,7 +265,6 @@ def cleanup_all_cached_backends() -> None: def _setup_pd_http_backend( - request: pytest.FixtureRequest, model_pool: ModelPool, model_id: str, workers_config: dict, @@ -274,7 +282,6 @@ def _setup_pd_http_backend( def _setup_pd_grpc_backend( - request: pytest.FixtureRequest, model_pool: ModelPool, model_id: str, workers_config: dict, @@ -413,12 +420,7 @@ def _setup_pd_backend_common( f"{len(decodes)} decode (need {num_prefill} prefill, {num_decode} decode)" ) except Exception: - # Release all acquired workers on any failure - for w in acquired_workers: - try: - w.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(acquired_workers) raise model_path = prefills[0].model_path @@ -435,12 +437,7 @@ def _setup_pd_backend_common( log_dir=gateway_config.get("log_dir"), ) except Exception: - # Release workers if gateway fails to start - for w in acquired_workers: - try: - w.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(acquired_workers) raise client = openai.OpenAI( @@ -463,11 +460,7 @@ def _setup_pd_backend_common( finally: logger.info("Tearing down %s PD gateway", runtime_label) gateway.shutdown() - for worker in acquired_workers: - try: - worker.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(acquired_workers) def _setup_grpc_backend( @@ -554,11 +547,7 @@ def _setup_grpc_backend( worker_urls = [instance.worker_url] model_path = instance.model_path except Exception as e: - for inst in instances: - try: - inst.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(instances) if isinstance(e, RuntimeError): pytest.fail(str(e)) raise @@ -575,11 +564,7 @@ def _setup_grpc_backend( log_dir=gateway_config.get("log_dir"), ) except Exception: - for inst in instances: - try: - inst.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(instances) raise client = openai.OpenAI( @@ -601,11 +586,7 @@ def _setup_grpc_backend( finally: logger.info("Tearing down %s gRPC gateway", runtime_label) gateway.shutdown() - for inst in instances: - try: - inst.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(instances) def _setup_local_backend( @@ -665,12 +646,7 @@ def _setup_local_backend( worker_urls = [instance.worker_url] model_path = instance.model_path except Exception as e: - # Release any acquired instances on failure - for inst in instances: - try: - inst.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(instances) if isinstance(e, RuntimeError): pytest.fail(str(e)) raise @@ -688,12 +664,7 @@ def _setup_local_backend( log_dir=gateway_config.get("log_dir"), ) except Exception: - # Release workers if gateway fails to start - for inst in instances: - try: - inst.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(instances) raise client = openai.OpenAI( @@ -715,12 +686,7 @@ def _setup_local_backend( finally: logger.info("Tearing down gateway for %s backend", backend_name) gateway.shutdown() - # Release references to allow eviction - for inst in instances: - try: - inst.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers(instances) def _setup_cloud_backend( @@ -808,21 +774,13 @@ def test_router_state(backend_router): model_path=instance.model_path, ) except Exception: - # Release worker if gateway fails to start if instance is not None: - try: - instance.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers([instance]) raise try: yield gateway finally: gateway.shutdown() - # Release reference to allow eviction if instance is not None: - try: - instance.release() - except Exception as release_err: - logger.warning("Failed to release worker during cleanup: %s", release_err) + _release_workers([instance]) diff --git a/e2e_test/infra/gpu_allocator.py b/e2e_test/infra/gpu_allocator.py index 7b1c648e80..0a42644aff 100644 --- a/e2e_test/infra/gpu_allocator.py +++ b/e2e_test/infra/gpu_allocator.py @@ -191,6 +191,21 @@ def wait_for_gpu_memory_to_clear( physical_devices = get_physical_device_indices(devices) start_time = time.time() + # Define threshold check once (loop-invariant) + if threshold_bytes is not None: + threshold_gb = threshold_bytes / (1024**3) + + def is_free(used: float, total: float) -> bool: + return used <= threshold_gb + + threshold_desc = f"{threshold_gb:.1f} GiB" + else: + + def is_free(used: float, total: float) -> bool: + return used / total <= threshold_ratio # type: ignore[operator] + + threshold_desc = f"{threshold_ratio:.2%}" # type: ignore[str-format] + with nvml_context(): while True: output: dict[int, str] = {} @@ -209,19 +224,6 @@ def wait_for_gpu_memory_to_clear( " ".join(f"{k}={v}" for k, v in output.items()), ) - if threshold_bytes is not None: - - def is_free(used: float, total: float) -> bool: - return used <= threshold_bytes / (1024**3) - - threshold_desc = f"{threshold_bytes / (1024**3):.1f} GiB" - else: - - def is_free(used: float, total: float) -> bool: - return used / total <= threshold_ratio # type: ignore[operator] - - threshold_desc = f"{threshold_ratio:.2%}" # type: ignore[str-format] - dur_s = time.time() - start_time if all(is_free(used, total) for used, total in output_raw.values()): logger.info( diff --git a/e2e_test/infra/process_utils.py b/e2e_test/infra/process_utils.py index 1b20e68f76..5baae0f0ca 100644 --- a/e2e_test/infra/process_utils.py +++ b/e2e_test/infra/process_utils.py @@ -104,22 +104,23 @@ def wait_for_workers_ready( start = time.perf_counter() headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} - while time.perf_counter() - start < timeout: - try: - resp = requests.get(f"{router_url}/workers", headers=headers, timeout=5) - if resp.status_code == 200: - data = resp.json() - total = data.get("total", len(data.get("workers", []))) - if total >= expected_workers: - logger.info( - "All %d workers connected after %.1fs", - expected_workers, - time.perf_counter() - start, - ) - return - except requests.RequestException: - pass - time.sleep(2) + with requests.Session() as session: + while time.perf_counter() - start < timeout: + try: + resp = session.get(f"{router_url}/workers", headers=headers, timeout=5) + if resp.status_code == 200: + data = resp.json() + total = data.get("total", len(data.get("workers", []))) + if total >= expected_workers: + logger.info( + "All %d workers connected after %.1fs", + expected_workers, + time.perf_counter() - start, + ) + return + except requests.RequestException: + pass + time.sleep(2) raise TimeoutError( f"Router at {router_url} did not get {expected_workers} workers within {timeout}s"