wehub-resource-sync
80 lines
2.9 KiB
Python
80 lines
2.9 KiB
Python
import os
|
|
from typing import List, Union
|
|
|
|
import torch
|
|
|
|
import ray
|
|
from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
|
|
|
|
|
|
class CUDATorchDeviceManager(TorchDeviceManager):
|
|
"""CUDA device manager"""
|
|
|
|
def is_available(self) -> bool():
|
|
return torch.cuda.is_available()
|
|
|
|
def get_devices(self) -> List[torch.device]:
|
|
"""Gets the correct torch device list configured for this process.
|
|
|
|
Returns a list of torch CUDA devices allocated for the current worker.
|
|
If no GPUs are assigned, then it returns a list with a single CPU device.
|
|
|
|
Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
|
|
superset of the `ray.get_gpu_ids()`.
|
|
"""
|
|
|
|
# GPU IDs are assigned by Ray after you specify "use_gpu"
|
|
# GPU `ray.get_gpu_ids()` may return ints or may return strings.
|
|
# We should always convert to strings.
|
|
gpu_ids = [str(id) for id in ray.get_gpu_ids()]
|
|
|
|
device_ids = []
|
|
|
|
if len(gpu_ids) > 0:
|
|
cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
|
|
if cuda_visible_str and cuda_visible_str != "NoDevFiles":
|
|
cuda_visible_list = cuda_visible_str.split(",")
|
|
else:
|
|
cuda_visible_list = []
|
|
|
|
# By default, there should only be one GPU ID if `use_gpu=True`.
|
|
# If there are multiple GPUs, return a list of devices.
|
|
# If using fractional GPUs, these IDs are not guaranteed
|
|
# to be unique across different processes.
|
|
for gpu_id in gpu_ids:
|
|
try:
|
|
device_ids.append(cuda_visible_list.index(gpu_id))
|
|
except IndexError:
|
|
raise RuntimeError(
|
|
"CUDA_VISIBLE_DEVICES set incorrectly. "
|
|
f"Got {cuda_visible_str}, expected to include {gpu_id}. "
|
|
"Did you override the `CUDA_VISIBLE_DEVICES` environment"
|
|
" variable? If not, please help file an issue on Github."
|
|
)
|
|
|
|
else:
|
|
# If called on the driver or outside of Ray Train, return the
|
|
# 0th device.
|
|
device_ids.append(0)
|
|
|
|
return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
|
|
|
|
def set_device(self, device: Union[torch.device, int, str, None]):
|
|
torch.cuda.set_device(device)
|
|
|
|
def supports_stream(self) -> bool:
|
|
"""Validate if the device type support create a stream"""
|
|
return True
|
|
|
|
def create_stream(self, device: torch.device) -> torch.cuda.Stream:
|
|
"""Create a stream on cuda device"""
|
|
return torch.cuda.Stream(device)
|
|
|
|
def get_stream_context(self, stream):
|
|
"""Get a stream context for cuda device"""
|
|
return torch.cuda.stream(stream)
|
|
|
|
def get_current_stream(self) -> torch.cuda.Stream:
|
|
"""Get current stream for cuda device"""
|
|
return torch.cuda.current_stream()
|