nvidia.go

package logic

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

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

/*
nvidia-smi --query-gpu index,pci.bus_id --format=csv
*/

// NVIDIAInfo 返回的map中，key为显卡index
func NVIDIAInfo() map[int]*CardInfo {
	if !have_nvidia() {
		return nil
	}
	pid2pcie, err := exec.Command("nvidia-smi", "--query-compute-apps=pid,gpu_bus_id", "--format=csv,noheader").CombinedOutput()
	if err != nil {
		return nil
	}
	index2pcie, err := exec.Command("nvidia-smi", "--query-gpu=index,pci.bus_id", "--format=csv,noheader").CombinedOutput()
	if err != nil {
		return nil
	}
	mapPCI2I := make(map[string]int)
	mapPID2PCI := make(map[int]string)
	for _, v := range strings.Split(strings.Trim(string(index2pcie), "\n"), "\n") {
		f := strings.Fields(strings.ReplaceAll(v, ",", " "))
		if len(f) != 2 {
			continue
		}
		index, err := strconv.Atoi(f[0])
		if err != nil {
			continue
		}
		mapPCI2I[f[1]] = index
	}
	for _, v := range strings.Split(strings.Trim(string(pid2pcie), "\n"), "\n") {
		f := strings.Fields(strings.ReplaceAll(v, ",", " "))
		if len(f) != 2 {
			continue
		}
		pid, err := strconv.Atoi(f[0])
		if err != nil {
			continue
		}
		mapPID2PCI[pid] = f[1]
	}
	result := make(map[int]*CardInfo)
	for pid, pci := range mapPID2PCI {
		index, have := mapPCI2I[pci]
		if have {
			cinfo, h := result[index]
			if h {
				cinfo.Pids = append(cinfo.Pids, pid)
			} else {
				cinfo = &CardInfo{
					Index: uint8(index),
					Type:  1,
					Pids:  make([]int, 0, 1),
				}
				cinfo.Pids = append(cinfo.Pids, pid)
				result[index] = cinfo
			}
		}
	}
	return result
}