[Tests] Run diffusion tiny model tests in parallel (#6339)

Signed-off-by: Nick Cao <ncao@redhat.com>

.buildkite/cuda/test-ready.yml

@@ -76,7 +76,7 @@ steps:
 
       - label: "Tiny Model Tests · Diffusion (base)"
         commands:
-          - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
+          - timeout 15m pytest -sv -n 4 tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model"
         mirror_hardwares: l4_1
 
   - label: "Engine&Model Executor Test"

pyproject.toml

@@ -41,8 +41,10 @@ dev = [
     "pytest-cov==7.1.0",
     "pytest-mock==3.15.1",
     "pytest-shard==0.1.2",
+    "pytest-xdist==3.8.0",
     "websockets==16.1",
     "pytest-subtests==0.15.0",
+    "filelock>=3.16.1,<4",
     "datasets==3.6.0",
     "mypy==1.11.1",
     "pre-commit==4.0.1",
@@ -242,7 +244,8 @@ python_classes = ["Test*"]
 python_functions = ["test_*"]
 addopts = [
     "--strict-markers",
-    "--strict-config"
+    "--strict-config",
+    "--dist=loadgroup"
 ]
 markers = [
     # ci/cd required
@@ -291,6 +294,7 @@ markers = [
     "sp: Sequence parallelism tests (multi-GPU)",
     "slow: Slow tests (may skip in quick CI)",
     "benchmark: Benchmark tests",
+    "xdist: xdist-parallelizable tests",
 ]
 filterwarnings = [
     "ignore:.*does not have '__test__' attribute.*:UserWarning",

tests/conftest.py

@@ -25,6 +25,7 @@ pytest_plugins = (
     "tests.helpers.fixtures.run_args",
     "tests.helpers.fixtures.runtime",
     "tests.helpers.fixtures.speaker_cache",
+    "tests.helpers.fixtures.xdist",
 )
 
 

tests/helpers/fixtures/xdist.py

@@ -0,0 +1,14 @@
+"""
+Only tests explicitly marked ``pytest.mark.xdist`` run concurrently.
+Other test are pinned to a single xdist worker.
+"""
+
+import pytest
+
+
+def pytest_collection_modifyitems(config, items):
+    if not config.pluginmanager.hasplugin("xdist"):
+        return
+    for item in items:
+        if not item.get_closest_marker("xdist"):
+            item.add_marker(pytest.mark.xdist_group(name="serial"))

tests/model_tests/conftest.py

@@ -1,12 +1,14 @@
+import json
 from shutil import rmtree
 
 import pytest
+from filelock import ReadWriteLock
 
 from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
 
 
 @pytest.fixture(scope="session")
-def tiny_model_paths(request, run_level):
+def tiny_model_paths(request, run_level, tmp_path_factory):
     """Build or download the tiny models for the selected tests.
 
     At core_model level, builds tiny models via the builder function.
@@ -15,27 +17,43 @@ def tiny_model_paths(request, run_level):
     NOTE: this is session scoped to avoid churn in tiny model creation,
     but will ensure all the tiny models you need are created for the selected tests
     before it starts to execute them."""
+    base_path = tmp_path_factory.getbasetemp().parent / "tiny_models"
+    use_lock = ReadWriteLock(base_path.with_suffix(".use.lock"))
+    build_lock = ReadWriteLock(base_path.with_suffix(".build.lock"))
+    data_file = base_path.with_suffix(".json")
+
     model_paths = {}
     built_paths = []
-    print("Initializing tiny models...")
-    for item in request.session.items:
-        if not hasattr(item, "callspec"):
-            continue
-        model_name = item.callspec.params.get("model_name")
-        if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS:
-            continue
-        if model_name not in model_paths:
-            settings = DIFFUSION_TEST_SETTINGS[model_name]
-            if run_level == "core_model":
-                print(f"Calling tiny model builder for: {model_name}")
-                path = settings.builder()
-                built_paths.append(path)
-            else:
-                print(f"Run level is {run_level}; {model_name} will use the full model")
-                path = settings.model
-            model_paths[model_name] = path
 
-    yield model_paths
-    for path in built_paths:
-        print(f"Removing tiny model: {path}")
-        rmtree(path, ignore_errors=True)
+    with use_lock.read_lock():
+        with build_lock.write_lock():
+            if data_file.is_file():
+                model_paths = json.loads(data_file.read_text())
+            else:
+                print("Initializing tiny models...")
+                for item in request.session.items:
+                    if not hasattr(item, "callspec"):
+                        continue
+                    model_name = item.callspec.params.get("model_name")
+                    if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS:
+                        continue
+                    if model_name not in model_paths:
+                        settings = DIFFUSION_TEST_SETTINGS[model_name]
+                        if run_level == "core_model":
+                            print(f"Calling tiny model builder for: {model_name}")
+                            path = settings.builder()
+                            built_paths.append(path)
+                        else:
+                            print(f"Run level is {run_level}; {model_name} will use the full model")
+                            path = settings.model
+                        model_paths[model_name] = path
+                data_file.write_text(json.dumps(model_paths))
+
+        yield model_paths
+
+    with use_lock.write_lock():
+        if built_paths:
+            data_file.unlink(missing_ok=True)
+        for path in built_paths:
+            print(f"Removing tiny model: {path}")
+            rmtree(path, ignore_errors=True)

tests/model_tests/diffusion/test_common_offline.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 
 from tests.model_tests.diffusion.case_filtering import get_parametrized_options
@@ -17,7 +20,7 @@ from tests.model_tests.diffusion.task_runners import (
 )
 
 # NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
 
 
 @pytest.mark.parametrize(
@@ -50,7 +53,7 @@ def test_pipeline_on_supported_tasks(
     )
     try:
         for task_type in supported_tasks:
-            with subtests.test(msg=task_type):
+            with subtests.test(msg=task_type.value):
                 if task_type == DiffusionTasks.TEXT_TO_IMAGE:
                     run_and_validate_text_to_image_request(omni)
                 elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:

tests/model_tests/diffusion/test_common_online.py

@@ -24,7 +24,7 @@ from tests.model_tests.diffusion.task_runners import (
 )
 
 # NOTE : Hardware marks are added dynamically based on test requirements
-pytestmark = [pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion, pytest.mark.xdist]
 
 
 @pytest.mark.parametrize(
@@ -58,7 +58,7 @@ def test_online_on_supported_tasks(
             log_stats=server.log_stats,
         )
         for task_type in supported_tasks:
-            with subtests.test(msg=task_type):
+            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: