func checkGpuType()

in partition_gpu/partition_gpu.go [241:261]


func checkGpuType() (string, error) {
	gpuType, err := exec.Command(*nvidiaSmiPath, "--query-gpu=gpu_name", "--format=csv,noheader").Output()
	if err != nil {
		return "", err
	}
	switch {
	case strings.HasPrefix(string(gpuType), NvidiaGB200):
		return NvidiaGB200, nil
	case strings.HasPrefix(string(gpuType), NvidiaB200):
		return NvidiaB200, nil
	case strings.HasPrefix(string(gpuType), Nvidia141gbH200):
		return Nvidia141gbH200, nil
	case strings.HasPrefix(string(gpuType), Nvidia80gbH100):
		return Nvidia80gbH100, nil
	case strings.HasPrefix(string(gpuType), Nvidia40gbA100):
		return Nvidia40gbA100, nil
	case strings.HasPrefix(string(gpuType), Nvidia80gbA100):
		return Nvidia80gbA100, nil
	}
	return "", fmt.Errorf("nvidia-smi returned invalid GPU type for MIG: %s", gpuType)
}