def _generate_gpu_devices()

in azure-slurm/slurmcc/cli.py [0:0]

• 12 lines of code
• 4 McCabe index (conditional complexity)

def _generate_gpu_devices(partition: partitionlib.Partition) -> str:
    
    # Override from node configuration
    if partition.gpu_device_config:
        return partition.gpu_device_config
    
    try:
        # TODO: azure sku should eventually indicate the GPU type
        #       (currently, attribute is not available, so gpu_family is not available to the partition)   
        #       We'll fall back to parsing machine type until it is available
        has_amd_gpu = "amd" in partition.gpu_family.lower()
    except AttributeError:           
        has_amd_gpu = "mi300" in partition.machine_type.lower()

    if has_amd_gpu:
        gpu_devices = _generate_amd_devices(partition.gpu_count)
    else:
        gpu_devices = _generate_nvidia_devices(partition.gpu_count)
    return gpu_devices