wehub-resource-sync
51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
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()
|