wangyuandamy-why-3459
4af0ddc1c2
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] <your.email@example.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Session-scoped config and environment fixtures."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
|
|
|
|
|
@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", 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).
|
|
"""
|
|
# 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
|