nabu.resources.gpu module

gpu.py: general-purpose utilities for GPU

class nabu.resources.gpu.GPUDescription(device, vendor=None, device_id=None)[source]

Bases: object

Simple description of a Graphical Processing Unit. This class is designed to be simple to understand, and to be serializable for being used by dask.distributed.

Create a description from a device.

Parameters:

device (pycuda.driver.Device or pyopencl.Device) – Class describing a GPU device.

get_dict()[source]
nabu.resources.gpu.pick_gpus(method, cuda_gpus, opencl_platforms, n_gpus)[source]
nabu.resources.gpu.pick_gpus_auto(cuda_gpus, opencl_platforms, n_gpus)[source]

Pick n_gpus devices with the best available driver.

This function browse the visible Cuda GPUs and Opencl platforms to pick the GPUs with the best driver. A worker might see several implementations of a GPU driver. For example with Nvidia hardware, we can see:

  • The Cuda implementation (nvidia-cuda-toolkit)

  • OpenCL implementation by Nvidia (nvidia-opencl-icd)

  • OpenCL implementation by Portable OpenCL

Parameters:

  • cuda_gpu (dict) – Dictionary where each key is an ID, and the value is a dictionary describing some attributes of the GPU (obtained with GPUDescription)

  • opencl_platforms (dict) – Dictionary where each key is the platform name, and the value is a list of dictionary descriptions.

  • n_gpus (int) – Number of GPUs to pick.

nabu.resources.gpu.pick_gpus_nvidia(cuda_gpus, n_gpus)[source]

Pick one or more Nvidia GPUs.