128 lines
5.7 KiB
Python
128 lines
5.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""
|
|
Analogous to test_common_offline, but for server tests. Validates the full
|
|
online serving stack (CLI arg parsing, subprocess, API routing, response
|
|
encoding) using tiny models.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from xdist import is_xdist_worker
|
|
|
|
from tests.helpers.runtime import OmniServer, OnlineOmniClient
|
|
from tests.model_tests.diffusion.case_filtering import get_parametrized_options
|
|
from tests.model_tests.diffusion.config_types import (
|
|
DiffusionAccs,
|
|
DiffusionTasks,
|
|
build_server_args_from_diff_accelerations,
|
|
)
|
|
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
|
from tests.model_tests.diffusion.task_runners import (
|
|
run_and_validate_online_determinism,
|
|
run_and_validate_online_image_edits,
|
|
run_and_validate_online_image_to_image_request,
|
|
run_and_validate_online_image_to_video_request,
|
|
run_and_validate_online_multi_output,
|
|
run_and_validate_online_text_to_image_request,
|
|
run_and_validate_online_text_to_video_request,
|
|
)
|
|
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
|
|
|
# NOTE : Hardware marks are added dynamically based on test requirements
|
|
pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _disable_global_gpu_cleanup_for_parallel_workers(
|
|
request: pytest.FixtureRequest,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Do not wait on total GPU usage while sibling xdist workers are active.
|
|
|
|
``cleanup_test_environment`` observes the whole device, not allocations
|
|
owned by the current worker. Waiting for the device to fall below its
|
|
global threshold therefore turns server teardown into a cross-worker
|
|
barrier when online tests run with xdist. ``OmniServer`` still tears down
|
|
its own subprocess tree; retain the broader cleanup for non-xdist runs.
|
|
"""
|
|
if is_xdist_worker(request):
|
|
monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
|
|
get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
|
|
)
|
|
def test_online_on_supported_tasks(
|
|
model_name: str,
|
|
accelerations: list[DiffusionAccs] | None,
|
|
supported_tasks: list[DiffusionTasks],
|
|
check_multioutput: bool,
|
|
check_determinism: bool,
|
|
tiny_model_paths: dict[str, str],
|
|
run_level: str,
|
|
subtests,
|
|
):
|
|
"""Smoke test: start a tiny model server and run each supported task via the API."""
|
|
model_path = tiny_model_paths[model_name]
|
|
server_args = build_server_args_from_diff_accelerations(accelerations)
|
|
server_args.append("--enforce-eager")
|
|
settings = DIFFUSION_TEST_SETTINGS[model_name]
|
|
if settings.checkpoint_filename is not None:
|
|
model_path = str(Path(model_path) / settings.checkpoint_filename)
|
|
server_args.extend(["--model-class-name", model_name])
|
|
|
|
with OmniServer(model_path, server_args) as server:
|
|
# TODO: We may want to revisit run_level validation here,
|
|
# because checks for things like image size etc should not
|
|
# depend on whether or not the weights are real or random
|
|
client = OnlineOmniClient(
|
|
host=server.host,
|
|
port=server.port,
|
|
api_key="EMPTY",
|
|
run_level=run_level,
|
|
log_stats=server.log_stats,
|
|
)
|
|
for task_type in supported_tasks:
|
|
with subtests.test(msg=task_type.value):
|
|
if task_type == DiffusionTasks.TEXT_TO_IMAGE:
|
|
run_and_validate_online_text_to_image_request(server, client)
|
|
elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
|
|
run_and_validate_online_image_to_image_request(server, client)
|
|
max_multimodal_image_inputs = (
|
|
get_diffusion_model_metadata(model_name).max_multimodal_image_inputs or 1
|
|
)
|
|
image_counts = [1]
|
|
if max_multimodal_image_inputs != 1:
|
|
image_counts.append(max_multimodal_image_inputs)
|
|
image_counts.append(max_multimodal_image_inputs + 1)
|
|
for num_images in image_counts:
|
|
with subtests.test(api="/v1/images/edits", num_images=num_images):
|
|
run_and_validate_online_image_edits(
|
|
server,
|
|
client,
|
|
num_images=num_images,
|
|
max_multimodal_image_inputs=max_multimodal_image_inputs,
|
|
)
|
|
elif task_type == DiffusionTasks.TEXT_TO_VIDEO:
|
|
run_and_validate_online_text_to_video_request(server, client)
|
|
elif task_type == DiffusionTasks.IMAGE_TO_VIDEO:
|
|
run_and_validate_online_image_to_video_request(server, client)
|
|
else:
|
|
raise ValueError(f"Task type {task_type} is not yet supported")
|
|
|
|
# NOTE: For now, we only check determinism + multi output for the base case,
|
|
# since checking it on every extra acceleration configuration is redundant
|
|
# (see case_filtering).
|
|
if check_determinism:
|
|
for task_type in supported_tasks:
|
|
with subtests.test(msg=f"determinism[{task_type}]"):
|
|
run_and_validate_online_determinism(server, client, task_type)
|
|
if check_multioutput:
|
|
for task_type in supported_tasks:
|
|
with subtests.test(msg=f"multi_output[{task_type}]"):
|
|
run_and_validate_online_multi_output(server, client, task_type)
|