wehub-resource-sync
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
from typing import List
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
import ray._private.thirdparty.pynvml as pynvml
|
|
|
|
|
|
class DeviceHandleMock(dict):
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
uuid: str,
|
|
mig_devices: List["DeviceHandleMock"] = None,
|
|
**kwargs
|
|
):
|
|
super().__init__()
|
|
self["name"] = name
|
|
self["uuid"] = uuid
|
|
if mig_devices is not None:
|
|
self["mig_devices"] = mig_devices
|
|
self.update(kwargs)
|
|
|
|
|
|
# pnvml mock for gpu resources
|
|
class PyNVMLMock:
|
|
def __init__(self, mock_data, driver_version="535.104.12"):
|
|
self._mock_data = mock_data
|
|
self.driver_version = driver_version
|
|
|
|
def nvmlInit(self):
|
|
return
|
|
|
|
def nvmlShutdown(self):
|
|
return
|
|
|
|
def nvmlSystemGetDriverVersion(self):
|
|
return self.driver_version
|
|
|
|
def nvmlDeviceGetCount(self):
|
|
return len(self._mock_data)
|
|
|
|
def nvmlDeviceGetHandleByIndex(self, index):
|
|
return self._mock_data[index]
|
|
|
|
def nvmlDeviceGetName(self, handle):
|
|
return handle.get("name", "")
|
|
|
|
def nvmlDeviceGetMaxMigDeviceCount(self, handle):
|
|
if "mig_devices" in handle:
|
|
return max(7, len(handle["mig_devices"]))
|
|
else:
|
|
raise pynvml.NVMLError_NotSupported
|
|
|
|
def nvmlDeviceGetMigDeviceHandleByIndex(self, handle, mig_index):
|
|
try:
|
|
return handle["mig_devices"][mig_index]
|
|
except IndexError:
|
|
raise pynvml.NVMLError_NotFound
|
|
|
|
def nvmlDeviceGetUUID(self, handle):
|
|
return handle.get("uuid", "")
|
|
|
|
def nvmlDeviceGetComputeInstanceId(self, mig_handle):
|
|
return mig_handle["ci_id"]
|
|
|
|
def nvmlDeviceGetGpuInstanceId(self, mig_handle):
|
|
return mig_handle["gi_id"]
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_mock_pynvml(mock_nvml):
|
|
with patch("ray._private.thirdparty.pynvml.nvmlInit", mock_nvml.nvmlInit), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlShutdown", mock_nvml.nvmlShutdown
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlSystemGetDriverVersion",
|
|
mock_nvml.nvmlSystemGetDriverVersion,
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetCount",
|
|
mock_nvml.nvmlDeviceGetCount,
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetHandleByIndex",
|
|
mock_nvml.nvmlDeviceGetHandleByIndex,
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetName", mock_nvml.nvmlDeviceGetName
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetMaxMigDeviceCount",
|
|
mock_nvml.nvmlDeviceGetMaxMigDeviceCount,
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetMigDeviceHandleByIndex",
|
|
mock_nvml.nvmlDeviceGetMigDeviceHandleByIndex,
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetUUID", mock_nvml.nvmlDeviceGetUUID
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetComputeInstanceId",
|
|
mock_nvml.nvmlDeviceGetComputeInstanceId,
|
|
), patch(
|
|
"ray._private.thirdparty.pynvml.nvmlDeviceGetGpuInstanceId",
|
|
mock_nvml.nvmlDeviceGetGpuInstanceId,
|
|
):
|
|
yield
|