75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
import os
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.distributed import parallel_state
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def default_env():
|
|
# Keep behavior but avoid import-time side effects (RFC #2299).
|
|
keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE")
|
|
previous = {key: os.environ.get(key) for key in keys}
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn"
|
|
os.environ["VLLM_TARGET_DEVICE"] = previous["VLLM_TARGET_DEVICE"] or (
|
|
"cuda" if torch.cuda.is_available() and torch.accelerator.device_count() > 0 else "cpu"
|
|
)
|
|
yield
|
|
for key, value in previous.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def model_prefix() -> str:
|
|
prefix = os.environ.get("MODEL_PREFIX", "")
|
|
return f"{prefix.rstrip('/')}/" if prefix else ""
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_gpu_memory_between_tests():
|
|
"""Opt-in GPU pre/post hooks for a test (no environment-variable gate).
|
|
|
|
Use as a test parameter or ``@pytest.mark.usefixtures("clean_gpu_memory_between_tests")``.
|
|
"""
|
|
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
|
|
|
print("\n=== PRE-TEST GPU CLEANUP ===")
|
|
run_pre_test_cleanup()
|
|
yield
|
|
run_post_test_cleanup()
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def default_vllm_config():
|
|
"""Set a default VllmConfig for the whole test session.
|
|
|
|
Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and
|
|
deferred imports of ``tests.helpers.runtime`` both see the same context.
|
|
Function-scoped autouse ran too late for ``OmniRunner`` setup and could
|
|
desynchronize vLLM init vs request preprocessing (e.g. renderer state).
|
|
"""
|
|
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
|
|
|
# Use CPU device if no GPU is available (e.g., in CI environments)
|
|
has_gpu = torch.cuda.is_available() and torch.accelerator.device_count() > 0
|
|
device = "cuda" if has_gpu else "cpu"
|
|
device_config = DeviceConfig(device=device)
|
|
|
|
with set_current_vllm_config(VllmConfig(device_config=device_config)):
|
|
yield
|
|
|
|
|
|
@pytest.fixture()
|
|
def init_fake_tp_group(mocker):
|
|
"""Provide a fake TP group so vllm linear layers can be instantiated."""
|
|
mock_tp = mocker.MagicMock()
|
|
mock_tp.world_size = 1
|
|
mock_tp.rank_in_group = 0
|
|
old = parallel_state._TP
|
|
parallel_state._TP = mock_tp
|
|
yield
|
|
parallel_state._TP = old
|