Class DeviceManager

checkpoint_engine/device_utils.py:198–257

class DeviceManager:
    def __init__(self):
        self.device_type = self._detect_device_type()
        self._setup_device_module()

    def _is_torch_npu_available(self) -> bool:
        try:
            if hasattr(torch, "npu") and callable(getattr(torch.npu, "is_available", None)):
                return torch.npu.is_available()
            else:
                return False
        except ImportError:
            return False

    def _detect_device_type(self) -> str:
        if self._is_torch_npu_available():
            return "npu"
        elif torch.cuda.is_available():
            return "cuda"
        else:
            raise TypeError("The current device type is not supported")

    def _setup_device_module(self):
        if self.device_type == "npu":
            import torch_npu

            self.device_module = torch_npu.npu
        elif self.device_type == "cuda":
            self.device_module = torch.cuda
        else:
            raise TypeError("The current device type is not supported")

    @property
    def backend(self) -> str:
        if self.device_type == "npu":
            return "hccl"
        elif self.device_type == "cuda":
            return "nccl"
        else:
            raise TypeError("The current device type is not supported")

    @property
    def transfer_engine_protocol(self) -> str:
        if self.device_type == "npu":
            return "ascend_direct"
        elif self.device_type == "cuda":
            if has_efa_pci():
                return "efa"
            else:
                return "rdma"
        else:
            raise TypeError("The current device type is not supported")

    def rdma_device(self, rank: int) -> str:
        if self.transfer_engine_protocol == "ascend_direct":
            return ""
        elif self.transfer_engine_protocol in ["rdma", "efa"]:
            return _get_my_rdma_device(rank, self.device_module.device_count(), _get_rdma_devices())