chore: import upstream snapshot with attribution

python/ray/tests/accelerators/mock_pynvml.py

@@ -0,0 +1,101 @@
+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