diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py index 11165c6d2..e912c420d 100644 --- a/tests/e2e/online_serving/test_images_generations_lora.py +++ b/tests/e2e/online_serving/test_images_generations_lora.py @@ -12,10 +12,6 @@ This validates: import base64 import json import os -import signal -import subprocess -import sys -import time from io import BytesIO from pathlib import Path @@ -25,7 +21,8 @@ import requests import torch from PIL import Image from safetensors.torch import save_file -from vllm.utils.network_utils import get_open_port + +from tests.conftest import OmniServer os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -37,87 +34,6 @@ SIZE = "256x256" SEED = 42 -class OmniServer: - """Omniserver for vLLM-Omni tests.""" - - def __init__( - self, - model: str, - serve_args: list[str], - *, - env_dict: dict[str, str] | None = None, - ) -> None: - self.model = model - self.serve_args = serve_args - self.env_dict = env_dict - self.proc: subprocess.Popen | None = None - self.host = "127.0.0.1" - self.port = get_open_port() - - def _start_server(self) -> None: - env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - if self.env_dict is not None: - env.update(self.env_dict) - - cmd = [ - sys.executable, - "-m", - "vllm_omni.entrypoints.cli.main", - "serve", - self.model, - "--omni", - "--host", - self.host, - "--port", - str(self.port), - ] + self.serve_args - - print(f"Launching OmniServer with: {' '.join(cmd)}") - self.proc = subprocess.Popen( - cmd, - env=env, - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), # vllm-omni root - start_new_session=True, - ) - - # Wait for server to be ready. - max_wait = 1200 - url = f"http://{self.host}:{self.port}/v1/models" - start_time = time.time() - while time.time() - start_time < max_wait: - try: - resp = requests.get(url, headers={"Authorization": "Bearer EMPTY"}, timeout=10) - if resp.status_code == 200: - print(f"Server ready on {self.host}:{self.port}") - return - except Exception: - pass - time.sleep(2) - - raise RuntimeError(f"Server failed to become ready within {max_wait} seconds") - - def __enter__(self): - self._start_server() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.proc is None: - return - try: - os.killpg(self.proc.pid, signal.SIGTERM) - except ProcessLookupError: - pass - try: - self.proc.wait(timeout=30) - except subprocess.TimeoutExpired: - try: - os.killpg(self.proc.pid, signal.SIGKILL) - except ProcessLookupError: - pass - self.proc.wait() - - @pytest.fixture(scope="module") def omni_server(): with OmniServer(MODEL, ["--num-gpus", "1"]) as server: @@ -183,17 +99,30 @@ def _image_blue_tail_slice(img: Image.Image) -> np.ndarray: return tail -def _assert_slice_close(actual: np.ndarray, expected: np.ndarray, *, label: str) -> None: +def _slice_diff_stats(actual: np.ndarray, expected: np.ndarray) -> tuple[float, float]: + diff = np.abs(actual - expected) + return float(diff.max()), float(diff.mean()) + + +def _assert_slice_close( + actual: np.ndarray, + expected: np.ndarray, + *, + label: str, + base_max: float, + base_mean: float, +) -> None: assert actual.shape == (3, 3) assert expected.shape == (3, 3) - diff = np.abs(actual - expected) - max_diff = float(diff.max()) - mean_diff = float(diff.mean()) + max_diff, mean_diff = _slice_diff_stats(actual, expected) # NOTE: Different attention backends / torch.compile can introduce small # floating-point drift that shows up as a few LSBs in uint8 pixels. Keep # the reset check tolerant but bounded to avoid flaky CI. - assert max_diff <= 5.0 and mean_diff <= 3.0, ( - f"{label} slice mismatch (max={max_diff:.1f}, mean={mean_diff:.1f}): {actual.tolist()}" + max_thresh = max(10.0, base_max + 4.0) + mean_thresh = max(6.0, base_mean + 4.0) + assert max_diff <= max_thresh and mean_diff <= mean_thresh, ( + f"{label} slice mismatch (max={max_diff:.1f} > {max_thresh:.1f} or " + f"mean={mean_diff:.1f} > {mean_thresh:.1f}): {actual.tolist()}" ) @@ -219,6 +148,9 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer, # Base generation. base_img = _post_images(omni_server, _basic_payload()) base_slice = _image_blue_tail_slice(base_img) + base_ref_img = _post_images(omni_server, _basic_payload()) + base_ref_slice = _image_blue_tail_slice(base_ref_img) + base_ref_max, base_ref_mean = _slice_diff_stats(base_ref_slice, base_slice) # Adapter A: apply delta to V slice only. lora_a_dir = tmp_path / "zimage_lora_a" @@ -245,11 +177,17 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer, # Ensure switching back to no-LoRA restores the base output. base_img_2 = _post_images(omni_server, _basic_payload()) base_slice_2 = _image_blue_tail_slice(base_img_2) - _assert_slice_close(base_slice_2, base_slice, label="base_after_reset") - base_reset = float(np.abs(base_slice_2 - base_slice).mean()) + _, base_reset_mean = _slice_diff_stats(base_slice_2, base_slice) + _assert_slice_close( + base_slice_2, + base_slice, + label="base_after_reset", + base_max=base_ref_max, + base_mean=base_ref_mean, + ) # Ensure LoRA effects are clearly above the baseline drift. - min_delta = base_reset + 0.5 + min_delta = max(base_reset_mean + 1.0, 1.5) assert a_vs_base > min_delta, f"lora_a_vs_base drift too small: {a_vs_base} <= {min_delta}" assert b_vs_base > min_delta, f"lora_b_vs_base drift too small: {b_vs_base} <= {min_delta}" assert b_vs_a > min_delta, f"lora_b_vs_lora_a drift too small: {b_vs_a} <= {min_delta}"