[Misc] Migrate L1 tests to use pytest-mock (#1315)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com> Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed
1 parent
9913c00144
commit
ed08a216c1
20 files changed
+977
-954
No files matched your search
@@ -1,52 +1,54 @@
|
||||
"""Tests for process-scoped GPU memory accounting."""
|
||||
|
||||
import os
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TestParseCudaVisibleDevices:
|
||||
def test_empty(self):
|
||||
def test_empty(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
|
||||
|
||||
with mock.patch.dict(os.environ, {}, clear=True):
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
assert parse_cuda_visible_devices() == []
|
||||
mocker.patch.dict(os.environ, {}, clear=True)
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
assert parse_cuda_visible_devices() == []
|
||||
|
||||
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
|
||||
assert parse_cuda_visible_devices() == []
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""})
|
||||
assert parse_cuda_visible_devices() == []
|
||||
|
||||
def test_integer_indices(self):
|
||||
def test_integer_indices(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
|
||||
|
||||
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}):
|
||||
assert parse_cuda_visible_devices() == [2, 3, 5]
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"})
|
||||
assert parse_cuda_visible_devices() == [2, 3, 5]
|
||||
|
||||
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}):
|
||||
assert parse_cuda_visible_devices() == [0]
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"})
|
||||
assert parse_cuda_visible_devices() == [0]
|
||||
|
||||
def test_uuids(self):
|
||||
def test_uuids(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
|
||||
|
||||
uuid1 = "GPU-12345678-1234-1234-1234-123456789abc"
|
||||
uuid2 = "GPU-87654321-4321-4321-4321-cba987654321"
|
||||
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}):
|
||||
assert parse_cuda_visible_devices() == [uuid1, uuid2]
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"})
|
||||
assert parse_cuda_visible_devices() == [uuid1, uuid2]
|
||||
|
||||
def test_mig_ids(self):
|
||||
def test_mig_ids(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
|
||||
|
||||
mig1 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/0/0"
|
||||
mig2 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/1/0"
|
||||
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}):
|
||||
assert parse_cuda_visible_devices() == [mig1, mig2]
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"})
|
||||
assert parse_cuda_visible_devices() == [mig1, mig2]
|
||||
|
||||
def test_spaces(self):
|
||||
def test_spaces(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
|
||||
|
||||
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}):
|
||||
assert parse_cuda_visible_devices() == [2, 3, 5]
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "})
|
||||
assert parse_cuda_visible_devices() == [2, 3, 5]
|
||||
|
||||
|
||||
class TestGetProcessGpuMemory:
|
||||
@@ -68,87 +70,79 @@ class TestGetProcessGpuMemory:
|
||||
del tensor
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def test_raises_on_invalid_device(self):
|
||||
def test_raises_on_invalid_device(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
||||
|
||||
with (
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
|
||||
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="Invalid GPU device"):
|
||||
get_process_gpu_memory(5)
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
|
||||
mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1)
|
||||
|
||||
def test_returns_zero_when_process_not_found(self):
|
||||
with pytest.raises(RuntimeError, match="Invalid GPU device"):
|
||||
get_process_gpu_memory(5)
|
||||
|
||||
def test_returns_zero_when_process_not_found(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
||||
|
||||
with (
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
|
||||
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex"),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]),
|
||||
):
|
||||
memory = get_process_gpu_memory(0)
|
||||
assert memory == 0
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
|
||||
mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8)
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex")
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[])
|
||||
|
||||
def test_uses_uuid_when_provided(self):
|
||||
memory = get_process_gpu_memory(0)
|
||||
assert memory == 0
|
||||
|
||||
def test_uses_uuid_when_provided(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
||||
|
||||
uuid = "GPU-12345678-1234-1234-1234-123456789abc"
|
||||
mock_handle = mock.MagicMock()
|
||||
mock_handle = mocker.MagicMock()
|
||||
|
||||
with (
|
||||
mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
|
||||
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) as mock_by_uuid,
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]),
|
||||
):
|
||||
memory = get_process_gpu_memory(0)
|
||||
assert memory == 0
|
||||
mock_by_uuid.assert_called_once_with(uuid)
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid})
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
|
||||
mock_by_uuid = mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle)
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[])
|
||||
|
||||
def test_raises_on_invalid_uuid(self):
|
||||
memory = get_process_gpu_memory(0)
|
||||
assert memory == 0
|
||||
mock_by_uuid.assert_called_once_with(uuid)
|
||||
|
||||
def test_raises_on_invalid_uuid(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
||||
|
||||
uuid = "GPU-invalid-uuid"
|
||||
|
||||
with (
|
||||
mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
|
||||
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")),
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="Failed to get NVML handle"):
|
||||
get_process_gpu_memory(0)
|
||||
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid})
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
|
||||
mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID"))
|
||||
|
||||
def test_returns_none_on_nvml_init_failure(self):
|
||||
with pytest.raises(RuntimeError, match="Failed to get NVML handle"):
|
||||
get_process_gpu_memory(0)
|
||||
|
||||
def test_returns_none_on_nvml_init_failure(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
|
||||
|
||||
with (
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
|
||||
):
|
||||
result = get_process_gpu_memory(0)
|
||||
assert result is None
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable"))
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
|
||||
result = get_process_gpu_memory(0)
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestIsProcessScopedMemoryAvailable:
|
||||
def test_returns_true_when_nvml_works(self):
|
||||
def test_returns_true_when_nvml_works(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
|
||||
|
||||
with (
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
|
||||
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
|
||||
):
|
||||
assert is_process_scoped_memory_available() is True
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
|
||||
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
|
||||
assert is_process_scoped_memory_available() is True
|
||||
|
||||
def test_returns_false_when_nvml_fails(self):
|
||||
def test_returns_false_when_nvml_fails(self, mocker: MockerFixture):
|
||||
from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
|
||||
|
||||
with mock.patch(
|
||||
mocker.patch(
|
||||
"vllm_omni.worker.gpu_memory_utils.nvmlInit",
|
||||
side_effect=Exception("NVML unavailable"),
|
||||
):
|
||||
assert is_process_scoped_memory_available() is False
|
||||
)
|
||||
assert is_process_scoped_memory_available() is False
|
||||
Reference in new issue
Block a user