dcu.go

package logic

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

func have_dcu() bool {
	b, _ := utils.DetectCmd("hy-smi")
	return b
}

var (
	ReEmpty    = regexp.MustCompile(`^$`)
	ReUseless  = regexp.MustCompile(`^=.*$`)
	ReHCUIndex = regexp.MustCompile(`^(?i)\s+hcu index:\s+\[(.*)\]$`)
)

func DCUInfo() map[int]*CardInfo {
	if !have_dcu() {
		return nil
	}
	output, err := exec.Command("hy-smi", "--showpids").CombinedOutput()
	if err != nil {
		return nil
	}
	return parse_dcu_info(string(output))
}

func parse_dcu_info(input string) map[int]*CardInfo {
	lines := strings.Split(strings.Trim(input, "\n"), "\n")
	useful := make([]string, 0, len(lines)/8)
	for _, v := range lines {
		if ReEmpty.MatchString(v) || ReUseless.MatchString(v) {
			continue
		}
		v = strings.Trim(v, " ")
		if strings.HasPrefix(v, "PID:") || ReHCUIndex.MatchString(v) {
			useful = append(useful, v)
		}
	}
	result := make(map[int]*CardInfo)
	var pid int = 0
	for _, v := range useful {
		if strings.HasPrefix(v, "PID:") {
			p := strings.TrimPrefix(v, "PID: ")
			pp, err := strconv.Atoi(p)
			if err == nil {
				pid = pp
			} else {
				pid = 0
			}
			continue
		}
		if pid != 0 {
			f := ReHCUIndex.FindStringSubmatch(v)
			dcuStr := strings.Fields(strings.ReplaceAll(strings.ReplaceAll(f[1], "'", " "), ",", " "))
			dcuIndexs := make([]int, 0, len(dcuStr))
			for _, d := range dcuStr {
				pp, err := strconv.Atoi(d)
				if err == nil {
					dcuIndexs = append(dcuIndexs, pp)
				}
			}
			for _, dcuIndex := range dcuIndexs {
				cinfo, have := result[dcuIndex]
				if have {
					cinfo.Pids = append(cinfo.Pids, pid)
				} else {
					cinfo = &CardInfo{
						Type:  0,
						Index: uint8(dcuIndex),
						Pids:  make([]int, 0),
					}
					cinfo.Pids = append(cinfo.Pids, pid)
					result[dcuIndex] = cinfo
				}
			}
		}
	}
	return result
}