wehub-resource-sync
72 lines
2.1 KiB
Python
72 lines
2.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_cuda_platform():
|
|
"""
|
|
Fixture that returns a factory for creating mocked CUDA platforms.
|
|
|
|
Usage:
|
|
def test_something(mock_cuda_platform):
|
|
with mock_cuda_platform(is_cuda=True, capability=(9, 0)):
|
|
# test code
|
|
"""
|
|
|
|
@contextmanager
|
|
def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None):
|
|
mock_platform = MagicMock()
|
|
mock_platform.is_cuda.return_value = is_cuda
|
|
mock_platform.is_xpu.return_value = False
|
|
device_capability = (
|
|
DeviceCapability(*capability) if capability is not None else None
|
|
)
|
|
mock_platform.get_device_capability.return_value = device_capability
|
|
|
|
def is_device_capability_family(
|
|
requested_capability: int, device_id: int = 0
|
|
) -> bool:
|
|
current_capability = mock_platform.get_device_capability(
|
|
device_id=device_id
|
|
)
|
|
if current_capability is None:
|
|
return False
|
|
return current_capability.major == (requested_capability // 10)
|
|
|
|
mock_platform.is_device_capability_family.side_effect = (
|
|
is_device_capability_family
|
|
)
|
|
with patch("vllm.platforms.current_platform", mock_platform):
|
|
yield mock_platform
|
|
|
|
return _mock_platform
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_xpu_platform():
|
|
"""
|
|
Fixture that returns a factory for creating mocked XPU platforms.
|
|
|
|
Usage:
|
|
def test_something(mock_xpu_platform):
|
|
with mock_xpu_platform():
|
|
# test code
|
|
"""
|
|
|
|
@contextmanager
|
|
def _mock_platform():
|
|
mock_platform = MagicMock()
|
|
mock_platform.is_cuda.return_value = False
|
|
mock_platform.is_xpu.return_value = True
|
|
with patch("vllm.platforms.current_platform", mock_platform):
|
|
yield mock_platform
|
|
|
|
return _mock_platform
|