chore: import upstream snapshot with attribution

python/ray/air/_internal/device_manager/npu.py

@@ -0,0 +1,104 @@
+import os
+from importlib.util import find_spec
+from typing import List, Union
+
+import torch
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+
+def is_package_present(package_name: str) -> bool:
+    try:
+        return find_spec(package_name) is not None
+    except ModuleNotFoundError:
+        return False
+
+
+NPU_TORCH_PACKAGE_AVAILABLE = is_package_present("torch_npu")
+
+
+if NPU_TORCH_PACKAGE_AVAILABLE:
+    import torch_npu  # noqa: F401
+
+
+class NPUTorchDeviceManager(TorchDeviceManager):
+    """Ascend NPU device manager"""
+
+    @staticmethod
+    def register_custom_torch_dist_backend():
+        if NPU_TORCH_PACKAGE_AVAILABLE:
+            import torch_npu  # noqa: F401, F811
+
+    def is_available(self) -> bool:
+        if not NPU_TORCH_PACKAGE_AVAILABLE:
+            return False
+
+        return torch.npu.is_available()
+
+    def get_devices(self) -> List[torch.device]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch NPU devices allocated for the current worker.
+        If no NPUs are assigned, then it returns a list with a single CPU device.
+        """
+        if NPU_TORCH_PACKAGE_AVAILABLE and torch.npu.is_available():
+            npu_ids = [
+                str(id)
+                for id in ray.get_runtime_context().get_accelerator_ids()[
+                    ray_constants.NPU
+                ]
+            ]
+
+            device_ids = []
+
+            if len(npu_ids) > 0:
+                npu_visible_str = os.environ.get(ASCEND_RT_VISIBLE_DEVICES_ENV_VAR, "")
+                if npu_visible_str and npu_visible_str != "NoDevFiles":
+                    npu_visible_list = npu_visible_str.split(",")
+                else:
+                    npu_visible_list = []
+
+                for npu_id in npu_ids:
+                    try:
+                        device_ids.append(npu_visible_list.index(npu_id))
+                    except IndexError:
+                        raise RuntimeError(
+                            "ASCEND_RT_VISIBLE_DEVICES set incorrectly. "
+                            f"Got {npu_visible_str}, expected to include {npu_id}. "
+                            "Did you override the `ASCEND_RT_VISIBLE_DEVICES` "
+                            "environment variable?"
+                        )
+            else:
+                # If called on the driver or outside of Ray Train, return the
+                # 0th device.
+                device_ids.append(0)
+
+            devices = [torch.device(f"npu:{device_id}") for device_id in device_ids]
+        else:
+            raise RuntimeError(
+                "Using NPUTorchDeviceManager but torch npu is not available."
+            )
+
+        return devices
+
+    def set_device(self, device: Union[torch.device, int]):
+        torch.npu.set_device(device)
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support to create a stream"""
+        return True
+
+    def create_stream(self, device):
+        """Create a stream on NPU device"""
+        return torch.npu.Stream(device)
+
+    def get_stream_context(self, stream):
+        """Get a torch.stream context on NPU device"""
+        return torch.npu.stream(stream)
+
+    def get_current_stream(self):
+        """Get current stream for NPU device"""
+        return torch.npu.current_stream()