src/cmd/cos_gpu_installer/deviceinfo/gpu_devices.go

// Package deviceinfo provides the devices information for cos-gpu-installer
package deviceinfo

import (
	"fmt"
	"os/exec"
	"strings"
)

type GPUType int

const (
	K80 GPUType = iota
	P4
	P100
	V100
	L4
	T4
	H100
	A100_40GB
	A100_80GB
	NO_GPU
	Others
)

var AvailableGPUTypesList = []GPUType{
	P4,
	P100,
	V100,
	L4,
	T4,
	H100,
	A100_40GB,
	A100_80GB,
}

var AllGPUTypeStrings = []string{
	"NVIDIA_TESLA_K80",
	"NVIDIA_TESLA_P4",
	"NVIDIA_TESLA_P100",
	"NVIDIA_TESLA_V100",
	"NVIDIA_L4",
	"NVIDIA_H100_80GB",
	"NVIDIA_TESLA_A100",
	"NVIDIA_A100_80GB",
	"NVIDIA_TESLA_T4",
}

// ParseGPUType converts a string to a GPUType enum.
func ParseGPUType(gpu string) (GPUType, error) {
	processedGPU := strings.ToUpper(strings.TrimSpace(gpu))
	switch processedGPU {
	case "NVIDIA_TESLA_K80":
		return K80, nil
	case "NVIDIA_TESLA_P4":
		return P4, nil
	case "NVIDIA_TESLA_P100":
		return P100, nil
	case "NVIDIA_TESLA_V100":
		return V100, nil
	case "NVIDIA_L4":
		return L4, nil
	case "NVIDIA_H100_80GB":
		return H100, nil
	case "NVIDIA_TESLA_A100":
		return A100_40GB, nil
	case "NVIDIA_A100_80GB":
		return A100_80GB, nil
	case "NVIDIA_TESLA_T4":
		return T4, nil
	default:
		return 0, fmt.Errorf("invalid GPU type string. Available GPU types are: %s", strings.Join(AllGPUTypeStrings, ", "))
	}
}

func (g GPUType) String() string {
	switch g {
	case K80:
		return "NVIDIA_TESLA_K80"
	case P4:
		return "NVIDIA_TESLA_P4"
	case P100:
		return "NVIDIA_TESLA_P100"
	case V100:
		return "NVIDIA_TESLA_V100"
	case L4:
		return "NVIDIA_L4"
	case H100:
		return "NVIDIA_H100_80GB"
	case A100_40GB:
		return "NVIDIA_TESLA_A100"
	case A100_80GB:
		return "NVIDIA_A100_80GB"
	case T4:
		return "NVIDIA_TESLA_T4"
	case NO_GPU:
		return "NO_GPU"
	case Others:
		return "OTHERS"
	default:
		return "UNKNOWN"
	}
}
func (g GPUType) OpenSupported() bool {
	switch g {
	case NO_GPU, K80, P4, P100, V100:
		return false
	default:
		return true
	}
}
func GetGPUTypeInfo() (GPUType, error) {
	cmd := "lspci | grep -i \"nvidia\""
	outBytes, err := exec.Command("/bin/bash", "-c", cmd).Output()
	if err != nil {
		return NO_GPU, err
	}
	out := string(outBytes)
	switch {
	case strings.Contains(out, "[Tesla K80]"):
		return K80, nil
	case strings.Contains(out, "NVIDIA Corporation Device 15f8"), strings.Contains(out, "NVIDIA Corporation GP100GL"), strings.Contains(out, "[Tesla P100"):
		return P100, nil
	case strings.Contains(out, "NVIDIA Corporation Device 1db1"), strings.Contains(out, "NVIDIA Corporation GV100GL"), strings.Contains(out, "[Tesla V100"):
		return V100, nil
	case strings.Contains(out, "NVIDIA Corporation Device 1bb3"), strings.Contains(out, "NVIDIA Corporation GP104GL"), strings.Contains(out, "[Tesla P4"):
		return P4, nil
	case strings.Contains(out, "NVIDIA Corporation Device 27b8"), strings.Contains(out, "NVIDIA Corporation AD104GL [L4]"):
		return L4, nil
	case strings.Contains(out, "NVIDIA Corporation Device 2330"), strings.Contains(out, "NVIDIA Corporation GH100[H100"):
		return H100, nil
	case strings.Contains(out, "NVIDIA Corporation Device 20b0"), strings.Contains(out, "NVIDIA Corporation GA100 [A100 SXM4 40GB"):
		return A100_40GB, nil
	case strings.Contains(out, "NVIDIA Corporation Device 20b2"), strings.Contains(out, "NVIDIA Corporation GA100 [A100 SXM4 80GB"):
		return A100_80GB, nil
	case strings.Contains(out, "NVIDIA Corporation Device 1eb8"), strings.Contains(out, "[Tesla T4]"):
		return T4, nil
	default:
		return Others, nil
	}
}