Source code for draugr.os_utilities.cuda_device

import os
from enum import Enum
from typing import Iterable, Union

__all__ = ["DeviceOrderEnum", "set_cuda_device_order", "set_cuda_visible_devices"]


[docs]class DeviceOrderEnum(Enum): fastest_first = "FASTEST_FIRST" pci_bus_id = "PCI_BUS_ID"
[docs]def set_cuda_device_order(order: DeviceOrderEnum = DeviceOrderEnum.pci_bus_id) -> None: """ :param order: :type order: """ os.environ["CUDA_DEVICE_ORDER"] = DeviceOrderEnum(order).value
[docs]def set_cuda_visible_devices(devices: Union[str, int, Iterable[int]]) -> None: """ :param devices: :type devices: """ if isinstance(devices, int): devices = str(devices) elif isinstance(devices, Iterable): devices = ",".join(str(d) for d in devices) if devices is None: # TODO: Nix specific, choose the least utilised device devices = "$(nvidia-smi --query-gpu=memory.free,index --format=csv,nounits,noheader | sort -nr | head -1 | awk '{ print $NF }')" os.environ["CUDA_VISIBLE_DEVICES"] = devices