[Bugfix] Restore torch default device between tests (#8157)
Signed-off-by: wangyu <410167048@qq.com>
This commit is contained in:
1 parent
4be6bbe0e4
commit
43e507117f
2 files changed
+25
-16
No files matched your search
@@ -1,15 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def restore_torch_default_dtype():
|
||||
"""Keep process-global dtype changes from leaking between diffusion tests."""
|
||||
default_dtype = torch.get_default_dtype()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.set_default_dtype(default_dtype)
|
||||
@@ -1,7 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Opt-in cleanup fixtures for GPU memory, pipeline registry, and speaker cache."""
|
||||
"""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
|
||||
|
||||
@@ -10,6 +14,26 @@ 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).
|
||||
|
||||
Reference in new issue
Block a user