74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Cleanup fixtures for process-global test state.
|
|
|
|
Torch dtype/device restoration is autouse. GPU memory, the pipeline registry,
|
|
and the speaker cache stay opt-in.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def restore_torch_default_dtype_and_device():
|
|
"""Restore process-global dtype and device after every test.
|
|
|
|
``get_default_device()`` reports ``cpu`` both when no default device was
|
|
set and when ``cpu`` was set explicitly. Writing that value back enables
|
|
default-device dispatch, so snapshot the active context instead.
|
|
"""
|
|
import torch
|
|
|
|
default_dtype = torch.get_default_dtype()
|
|
device_context = getattr(torch._GLOBAL_DEVICE_CONTEXT, "device_context", None)
|
|
default_device = None if device_context is None else torch.get_default_device()
|
|
try:
|
|
yield
|
|
finally:
|
|
torch.set_default_dtype(default_dtype)
|
|
torch.set_default_device(default_device)
|
|
|
|
|
|
@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.clean import cleanup_test_environment
|
|
|
|
print("\n=== PRE-TEST DEVICE CLEANUP ===")
|
|
cleanup_test_environment()
|
|
yield
|
|
cleanup_test_environment()
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_pipeline_registry():
|
|
"""Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it."""
|
|
snapshot = dict(OMNI_PIPELINES)
|
|
yield
|
|
OMNI_PIPELINES.clear()
|
|
OMNI_PIPELINES.update(snapshot)
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_speaker_cache():
|
|
"""Reset the process-wide speaker cache singleton before and after the test."""
|
|
import vllm_omni.utils.speaker_cache as _sc
|
|
|
|
def _reset():
|
|
with _sc._SINGLETON_LOCK:
|
|
if _sc._SINGLETON is not None:
|
|
_sc._SINGLETON.clear()
|
|
_sc._SINGLETON = None
|
|
|
|
_reset()
|
|
yield
|
|
_reset()
|