hy-smi.go

package gpu

import (
	"encoding/json"
	"get-container/utils"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
)

type HYVersionInfo struct {
	SMIVersion    string
	LibVersion    string
	DriverVersion string
}

type DCUInfo struct {
}

type DCUPidInfo struct {
	Pid             uint64
	PASId           uint64
	HCUNode         []string
	HCUIndex        []string
	GPUID           []string
	PCIBus          []string
	VRamUsed        utils.MemorySize
	VRamUsedPercent int
	SDMAUsed        int
}

var (
	ReEmptyLine   = regexp.MustCompile(`^\s*$`)
	ReUselessLine = regexp.MustCompile(`^=[ =a-zA-Z0-9]*=$`)
)

const (
	PIDHeader             = "PID"
	PASIDHeader           = "PASID"
	HCUNodeHeader         = "HCU Node(Include CPU sort)"
	HCUIndexHeader        = "HCU Index"
	GPUIDHeader           = "GPUID"
	PCIBusHeader          = "PCI BUS"
	VRamUsedHeader        = "VRAM USED(MiB)"
	VRamUsedPercentHeader = "VRAM USED(%)"
	SDMAUsedHeader        = "SDMA USED"
)

// GetDCUPidInfo 获取Pid相关信息
func GetDCUPidInfo() ([]DCUPidInfo, error) {
	output, err := exec.Command("hy-smi", "--showpids").Output()
	if err != nil {
		return nil, err
	}
	return parseDCUPidInfo(string(output))
}

func parseDCUPidInfo(s string) ([]DCUPidInfo, error) {
	lines := strings.Split(strings.Trim(string(s), "\n"), "\n")
	linesArray := make([][]string, 0)
	for _, line := range lines {
		if ReEmptyLine.MatchString(line) || ReUselessLine.MatchString(line) {
			continue
		}
		linesArray = append(linesArray, strings.SplitN(strings.TrimSpace(line), ":", 2))
	}
	infosArray := make([]map[string]string, 0)
	index := -1
	for _, line := range linesArray {
		if line == nil || len(line) != 2 {
			continue
		}
		line[1] = strings.TrimSpace(line[1])
		switch line[0] {
		case PIDHeader:
			index += 1
			infosArray = append(infosArray, make(map[string]string))
			infosArray[index][PIDHeader] = line[1]
			break
		case PASIDHeader:
			infosArray[index][PASIDHeader] = line[1]
			break
		case HCUNodeHeader:
			infosArray[index][HCUNodeHeader] = line[1]
			break
		case HCUIndexHeader:
			infosArray[index][HCUIndexHeader] = line[1]
			break
		case GPUIDHeader:
			infosArray[index][GPUIDHeader] = line[1]
			break
		case PCIBusHeader:
			infosArray[index][PCIBusHeader] = line[1]
			break
		case VRamUsedHeader:
			infosArray[index][VRamUsedHeader] = line[1]
			break
		case VRamUsedPercentHeader:
			infosArray[index][VRamUsedPercentHeader] = line[1]
			break
		case SDMAUsedHeader:
			infosArray[index][SDMAUsedHeader] = line[1]
			break
		default:
			break
		}
	}
	result := make([]DCUPidInfo, 0)
	if len(infosArray) == 0 {
		return result, nil
	}
	for _, info := range infosArray {
		i := DCUPidInfo{}
		pid, innerErr := strconv.ParseUint(info[PIDHeader], 10, 64)
		if innerErr != nil {
			return nil, innerErr
		}
		i.Pid = pid
		i.PASId, innerErr = strconv.ParseUint(info[PASIDHeader], 10, 64)
		if innerErr != nil {
			return nil, innerErr
		}
		_ = json.Unmarshal([]byte(strings.ReplaceAll(info[HCUNodeHeader], "'", `"`)), &i.HCUNode)
		_ = json.Unmarshal([]byte(strings.ReplaceAll(info[HCUIndexHeader], "'", `"`)), &i.HCUIndex)
		_ = json.Unmarshal([]byte(strings.ReplaceAll(info[GPUIDHeader], "'", `"`)), &i.GPUID)
		_ = json.Unmarshal([]byte(strings.ReplaceAll(info[PCIBusHeader], "'", `"`)), &i.PCIBus)
		s, innerErr := strconv.ParseUint(info[VRamUsedHeader], 10, 64)
		if innerErr != nil {
			return nil, innerErr
		}
		i.VRamUsed = utils.MemorySize{Num: s, Unit: utils.MiB}
		if info[VRamUsedPercentHeader] == "" || info[VRamUsedPercentHeader] == "inf" {
			i.VRamUsedPercent = -1
		} else {
			i.VRamUsedPercent, innerErr = strconv.Atoi(info[VRamUsedPercentHeader])
			if innerErr != nil {
				return nil, innerErr
			}
		}
		i.SDMAUsed, innerErr = strconv.Atoi(info[SDMAUsedHeader])
		if innerErr != nil {
			return nil, innerErr
		}
		result = append(result, i)
	}
	return result, nil
}