Skip to content

device

zeus.device

GPU device module for Zeus. Abstraction of GPU devices.

The main function of this module is get_gpus, which returns a GPU Manager object specific to the platform. To instantiate a GPU Manager object, you can do the following:

1
2
from zeus.device import get_gpus
gpus = get_gpus() # Returns NVIDIAGPUs() or AMDGPUs() depending on the platform.

There exists a 1:1 mapping between specific library functions and methods implemented in the GPU Manager object. For example, for NVIDIA systems, if you wanted to do:

1
2
handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_index)
constraints = pynvml.nvmlDeviceGetPowerManagementLimitConstraints(handle)

You can now do:

1
2
gpus = get_gpus() # returns a NVIDIAGPUs object
constraints =  gpus.getPowerManagementLimitConstraints(gpu_index)

Class hierarchy:

  • GPUs: Abstract class for GPU managers.
    • NVIDIAGPUs: GPU manager for NVIDIA GPUs, initialize NVIDIAGPU objects.
    • AMDGPUs: GPU manager for AMD GPUs, initialize AMDGPU objects.
  • GPU: Abstract class for GPU objects.

The following exceptions are defined in this module: