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