anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Fix] stabilize diffusion images LoRA E2E across CI drift (#1075)

Browse Source 
Signed-off-by: dongbo910220 <1275604947@qq.com>
This commit is contained in:
dongbo910220 authored and GitHub committed 2026-01-30 00:28:29 +08:00
1 parent bfbf3e5b69
commit 3eecd2057b
1 file changed
+33 -95
Download Patch File Download Diff File Expand all files Collapse all files
tests/e2e/online_serving/test_images_generations_lora.py
+33 -95
View File
@@ -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}"