chore: import upstream snapshot with attribution

python/ray/_private/accelerators/metax_gpu.py

@@ -0,0 +1,90 @@
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+
+logger = logging.getLogger(__name__)
+
+CUDA_VISIBLE_DEVICES_ENV_VAR = "CUDA_VISIBLE_DEVICES"
+NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
+
+
+class MetaxGPUAcceleratorManager(AcceleratorManager):
+    """Metax GPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "GPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return CUDA_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        cuda_visible_devices = os.environ.get(
+            MetaxGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+        if cuda_visible_devices is None:
+            return None
+
+        if cuda_visible_devices == "":
+            return []
+
+        return list(cuda_visible_devices.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        try:
+            import pymxsml.mxsml_extension as pymxsml
+
+            try:
+                pymxsml.mxSmlExInit()
+            except pymxsml.MXSMLEXError:
+                return 0
+            device_count = pymxsml.mxSmlExDeviceGetCount()
+            pymxsml.mxSmlExShutdown()
+            return device_count
+        except Exception as e:
+            logger.debug("Could not import pymxsml: %s", e)
+        return 0
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        try:
+            import pymxsml.mxsml_extension as pymxsml
+
+            try:
+                pymxsml.mxSmlExInit()
+            except pymxsml.MXSMLEXError:
+                return None
+            device_name = None
+            device_count = pymxsml.mxSmlExDeviceGetCount()
+            if device_count > 0:
+                handle = pymxsml.mxSmlExDeviceGetHandleByIndex(0)
+                device_name = pymxsml.mxSmlExDeviceGetName(handle)
+                if isinstance(device_name, bytes):
+                    device_name = device_name.decode("utf-8")
+            pymxsml.mxSmlExShutdown()
+            return device_name
+        except Exception:
+            logger.warning("Failed to detect GPU type.", exc_info=True)
+        return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_cuda_devices: List[str],
+    ) -> None:
+        if os.environ.get(NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR):
+            return
+
+        os.environ[
+            MetaxGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join(visible_cuda_devices)