def GetGpuType()

in perfkitbenchmarker/linux_packages/nvidia_driver.py [0:0]
42 lines of code
16 McCabe index (conditional complexity)

def GetGpuType(vm):
  """Return the type of NVIDIA gpu(s) installed on the vm.

  Args:
    vm: Virtual machine to query.

  Returns:
    Type of gpus installed on the vm as a string.

  Raises:
    NvidiaSmiParseOutputError: If nvidia-smi output cannot be parsed.
    HeterogeneousGpuTypesError: If more than one gpu type is detected.
    UnsupportedClockSpeedError: If gpu type is not supported.

  Example:
    If 'nvidia-smi -L' returns:

    GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-1a046bb9-e456-45d3-5a35-52da392d09a5)
    GPU 1: Tesla V100-SXM2-16GB (UUID: GPU-56cf4732-054c-4e40-9680-0ec27e97d21c)
    GPU 2: Tesla V100-SXM2-16GB (UUID: GPU-4c7685ad-4b3a-8adc-ce20-f3a945127a8a)
    GPU 3: Tesla V100-SXM2-16GB (UUID: GPU-0b034e63-22be-454b-b395-382e2d324728)
    GPU 4: Tesla V100-SXM2-16GB (UUID: GPU-b0861159-4727-ef2f-ff66-73a765f4ecb6)
    GPU 5: Tesla V100-SXM2-16GB (UUID: GPU-16ccaf51-1d1f-babe-9f3d-377e900bf37e)
    GPU 6: Tesla V100-SXM2-16GB (UUID: GPU-6eba1fa6-de10-80e9-ec5f-4b8beeff7e12)
    GPU 7: Tesla V100-SXM2-16GB (UUID: GPU-cba5a243-219c-df12-013e-1dbc98a8b0de)

    GetGpuType() will return:

    ['V100-SXM2-16GB', 'V100-SXM2-16GB', 'V100-SXM2-16GB', 'V100-SXM2-16GB',
     'V100-SXM2-16GB', 'V100-SXM2-16GB', 'V100-SXM2-16GB', 'V100-SXM2-16GB']
  """
  stdout, _ = vm.RemoteCommand('nvidia-smi -L')
  try:
    gpu_types = []
    for line in stdout.splitlines():
      if not line:
        continue
      splitted = line.split()
      if splitted[2] in ('Tesla', 'NVIDIA'):
        gpu_types.append(splitted[3])
      else:
        gpu_types.append(splitted[2])
  except:
    raise NvidiaSmiParseOutputError(
        'Unable to parse gpu type from {}'.format(stdout)
    )
  if any(gpu_type != gpu_types[0] for gpu_type in gpu_types):
    raise HeterogeneousGpuTypesError('PKB only supports one type of gpu per VM')

  if 'K80' in gpu_types[0]:
    return NVIDIA_TESLA_K80
  elif 'P4' in gpu_types[0]:
    return NVIDIA_TESLA_P4
  elif 'P100' in gpu_types[0]:
    return NVIDIA_TESLA_P100
  elif 'V100' in gpu_types[0]:
    return NVIDIA_TESLA_V100
  elif 'T4' in gpu_types[0]:
    return NVIDIA_TESLA_T4
  elif 'L4' in gpu_types[0]:
    return NVIDIA_L4
  elif 'A100' in gpu_types[0]:
    return NVIDIA_TESLA_A100
  elif 'A10' in gpu_types[0]:
    return NVIDIA_TESLA_A10
  elif 'H100' in gpu_types[0]:
    return NVIDIA_H100
  elif 'H200' in gpu_types[0]:
    return NVIDIA_H200
  else:
    raise UnsupportedClockSpeedError(
        'Gpu type {} is not supported by PKB'.format(gpu_types[0])
    )