chore: import upstream snapshot with attribution

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

@@ -0,0 +1,50 @@
+from contextlib import contextmanager
+from typing import List, Union
+
+import torch
+
+from ray._private.accelerators.hpu import HPU_PACKAGE_AVAILABLE
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+if HPU_PACKAGE_AVAILABLE:
+    import habana_frameworks.torch.hpu as torch_hpu
+
+
+class HPUTorchDeviceManager(TorchDeviceManager):
+    """HPU device manager"""
+
+    @staticmethod
+    def register_custom_torch_dist_backend():
+        if HPU_PACKAGE_AVAILABLE:
+            import habana_frameworks.torch.core  # noqa: F401
+            import habana_frameworks.torch.distributed.hccl  # noqa: F401
+
+    def is_available(self) -> bool():
+        if not HPU_PACKAGE_AVAILABLE:
+            return False
+
+        return torch_hpu.is_available()
+
+    def get_devices(self) -> List[torch.device]:
+        if not self.is_available():
+            raise RuntimeError(
+                "Using HPUTorchDeviceManager but torch hpu is not available."
+            )
+
+        return [torch.device("hpu")]
+
+    def set_device(self, device: Union[torch.device, int, str, None]):
+        torch_hpu.set_device(device)
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support create a stream"""
+        return False
+
+    def get_stream_context(self, stream):
+        """Get HPU stream context manager, empty so far."""
+
+        @contextmanager
+        def default_context_manager():
+            yield
+
+        return default_context_manager()