67 lines
2.7 KiB
Python
67 lines
2.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
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, tmp_path_factory):
|
|
"""Build or download the tiny models for the selected tests.
|
|
|
|
At core_model level, builds tiny models via the builder function.
|
|
At advanced_model / full_model level, uses the real HF model path or asset resolver.
|
|
|
|
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 = []
|
|
|
|
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")
|
|
if callable(settings.model):
|
|
path = settings.model()
|
|
built_paths.append(path)
|
|
else:
|
|
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)
|