prokbert.general_utils.count_gpus

prokbert.general_utils.count_gpus(method='clinfo')

Count the number of available GPUs using the specified method.

This function counts the number of NVIDIA and AMD GPUs using the chosen method. By default, it uses the ‘clinfo’ method for AMD GPUs.

Parameters

method (str, optional) – The method to use for GPU counting. Choose between ‘clinfo’ (default) and ‘rocm’.

Returns

The total number of GPUs detected.

Return type

int

Raises

  • ValueError – If an unknown method is provided.

  • Exception – If an error occurs while querying AMD GPUs using the specified method.

Note

  • The ‘clinfo’ method queries AMD GPUs by running the ‘clinfo’ command.

  • The ‘rocm’ method queries AMD GPUs by running ‘rocm-smi –list’ command.