Commit 4ff719ad authored by liming6's avatar liming6
Browse files

feature 添加了一个查看容器使用显卡的工具

parent d653d9ff
...@@ -43,6 +43,7 @@ func main() { ...@@ -43,6 +43,7 @@ func main() {
} }
model := tui.NewModelMain(w, h) model := tui.NewModelMain(w, h)
if _, err := tea.NewProgram(&model, tea.WithAltScreen()).Run(); err != nil { if _, err := tea.NewProgram(&model, tea.WithAltScreen()).Run(); err != nil {
backend.Shutdown()
log.Fatalf("error create program; %v", err) log.Fatalf("error create program; %v", err)
} }
backend.Shutdown() backend.Shutdown()
......
package backend
type DockerContainerInfo struct {
Name string `json:"name"` // 容器名
CreateUser string `json:"createUser"` // 容器创建者
Id string `json:"id"` // 容器id
CreateTime string `json:"createTime"` // 容器创建时间
UsingUser []string `json:"usingUser"` // 使用docker exec操作容器的用户
Status string `json:"status"` // 容器状态
Image string `json:"image"` // 容器使用的镜像
}
// // ListContainer
// func ListContainer() ([]*DockerContainerInfo, error) {
// cli, err := docker.GetDockerClient()
// if err != nil {
// return nil, err
// }
// defer cli.Close()
// csum, err := cli.ContainerList(context.Background(), client.ContainerListOptions{
// All: true,
// })
// if err != nil {
// return nil, err
// }
// result := make(map[string]*DockerContainerInfo)
// for _, v := range csum {
// }
// return result, nil
// }
...@@ -73,6 +73,7 @@ func ReturnGinList[T any](ctx *gin.Context, data []T, err error) { ...@@ -73,6 +73,7 @@ func ReturnGinList[T any](ctx *gin.Context, data []T, err error) {
}) })
} }
// webAuth 接口认证中间件
func webAuth(ctx *gin.Context) { func webAuth(ctx *gin.Context) {
authHeader := ctx.GetHeader("Authorization") authHeader := ctx.GetHeader("Authorization")
if authHeader == "" { if authHeader == "" {
......
package logic
import (
"fmt"
"strings"
)
type CardInfo struct {
Type uint8 // 0 dcu 1 nvidia
Index uint8
Pids []int
}
func (ci *CardInfo) String() string {
var cardType string
switch ci.Type {
case 0:
cardType = "dcu"
case 1:
cardType = "nvidia"
default:
cardType = "unknow"
}
is := make([]string, 0, len(ci.Pids))
for _, v := range ci.Pids {
is = append(is, fmt.Sprintf("%d", v))
}
return fmt.Sprintf("%s:%d:%s", cardType, ci.Index, strings.Join(is, ","))
}
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
}
================================= System Management Interface ==================================
================================================================================================
PIDs for KFD processes:
PID: 1523390
PASID: 32811
HCU Node(Include CPU sort): ['6']
HCU Index: ['4']
GPUID: ['1577']
PCI BUS: ['0000:9c:00.0']
VRAM USED(MiB): 59318
VRAM USED(%): 91
SDMA USED: 0
PID: 1523314
PASID: 32810
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
PID: 3646979
PASID: 32844
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 1
VRAM USED(%): inf
SDMA USED: 0
PID: 3208307
PASID: 32841
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
PID: 3647557
PASID: 32869
HCU Node(Include CPU sort): ['3']
HCU Index: ['1']
GPUID: ['34260']
PCI BUS: ['0000:54:00.0']
VRAM USED(MiB): 56595
VRAM USED(%): 86
SDMA USED: 0
PID: 1724730
PASID: 32777
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
PID: 1724483
PASID: 32768
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 1
VRAM USED(%): inf
SDMA USED: 0
PID: 3647298
PASID: 32868
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 1
VRAM USED(%): inf
SDMA USED: 0
PID: 3208460
PASID: 32842
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
PID: 3209594
PASID: 32853
HCU Node(Include CPU sort): ['2']
HCU Index: ['0']
GPUID: ['29238']
PCI BUS: ['0000:49:00.0']
VRAM USED(MiB): 64123
VRAM USED(%): 98
SDMA USED: 0
PID: 1523076
PASID: 32801
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 1
VRAM USED(%): inf
SDMA USED: 0
PID: 3208948
PASID: 32823
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
PID: 3646817
PASID: 32840
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 1
VRAM USED(%): inf
SDMA USED: 0
PID: 3871394
PASID: 32788
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
PID: 1724804
PASID: 32778
HCU Node(Include CPU sort): ['7']
HCU Index: ['5']
GPUID: ['24410']
PCI BUS: ['0000:bc:00.0']
VRAM USED(MiB): 60046
VRAM USED(%): 92
SDMA USED: 0
PID: 3871241
PASID: 32787
HCU Node(Include CPU sort):
HCU Index:
GPUID:
PCI BUS:
VRAM USED(MiB): 0
VRAM USED(%): inf
SDMA USED: 0
================================================================================================
======================================== End of SMI Log ========================================
package logic
import (
"fmt"
"os"
"strings"
"testing"
"time"
)
func TestRegExt(t *testing.T) {
test := " HCU Index: ['5','6']"
if ReHCUIndex.MatchString(test) {
t.Log("match")
f := ReHCUIndex.FindStringSubmatch(test)
t.Logf("%v", f[1])
}
}
func TestParse(t *testing.T) {
content, err := os.ReadFile("./hysmi-pid.log")
if err != nil {
t.Error(err)
}
s := parse_dcu_info(string(content))
for _, v := range s {
t.Logf("%s", v.String())
}
}
func TestDCUInfo(t *testing.T) {
s := DCUInfo()
for _, v := range s {
t.Logf("%s", v.String())
}
}
func TestNVIDIAInfo(t *testing.T) {
start := time.Now
s := NVIDIAInfo()
dd := time.Since(start())
t.Logf("%d ms", dd.Milliseconds())
for _, v := range s {
t.Logf("%s", v.String())
}
}
func formatStr(raw string, l int) string {
lstr := len(raw)
if l >= lstr {
return fmt.Sprintf("%s%s", raw, strings.Repeat(" ", l-lstr))
}
return raw[:l]
}
func TestFormatStr(t *testing.T) {
t.Logf("|%s|", formatStr("hello world", 2))
t.Logf("|%s|", formatStr("hello world", 20))
}
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
}
package main
import (
"context"
"fmt"
"get-container/docker"
"log"
"os"
"regexp"
"slices"
"sort"
"strconv"
"strings"
"get-container/cmd/pid_of_docker/logic"
"github.com/charmbracelet/lipgloss"
"github.com/moby/moby/client"
)
/*
容器名 容器id 创建者 使用的卡 容器中的进程数
ContainerName ContainerID CreateUser UsedCard PidsInContainer
*/
type Cinfo struct {
Name string
Id string
User []string
Pid []int
Cards []*logic.CardInfo
}
// Format 格式化容器信息
// 容器名,容器ID,创建者,使用的卡,容器中的进程数
func (c *Cinfo) Format() ([5]string, [5]int) {
var user, card string = " - ", " - "
if len(c.User) != 0 {
user = strings.Join(RemoveDuplicates(c.User), ",")
}
if len(c.Cards) != 0 {
ss := make([]string, 0, len(c.Cards))
for _, v := range c.Cards {
ss = append(ss, fmt.Sprintf("%d", v.Index))
}
slices.Sort(ss)
card = strings.Join(ss, ",")
}
rs := [5]string{c.Name, c.Id[:12], user, card, fmt.Sprintf("%d", len(c.Pid))}
rl := [5]int{}
for k, v := range rs {
rl[k] = len(v)
}
return rs, rl
}
var RegUser = regexp.MustCompile(`^(?i)/public[0-9]*/home/([0-9a-z_]+)(?:|/.*)$`)
func main() {
cli, err := docker.GetDockerClient()
if err != nil {
log.Fatalf("can't connect to docker daemon: %v", err)
}
csum, err := cli.ContainerList(context.Background(), client.ContainerListOptions{})
if err != nil {
log.Printf("error get container list: %v \n", err)
cli.Close()
os.Exit(1)
}
result := make(map[string]*Cinfo)
for _, i := range csum {
c := Cinfo{}
result[i.ID] = &c
c.Id = i.ID
c.Cards = make([]*logic.CardInfo, 0)
if len(i.Names) > 0 {
c.Name = strings.TrimPrefix(i.Names[0], "/")
}
if name, ok := i.Labels["com.sugon.username"]; ok {
c.User = make([]string, 0, 1)
c.User = append(c.User, name)
} else {
c.User = make([]string, 0, 4)
for _, v := range i.Mounts {
if RegUser.MatchString(v.Source) {
f := RegUser.FindStringSubmatch(v.Source)
if len(f) >= 2 {
c.User = append(c.User, f[1])
}
}
}
}
top, err := cli.ContainerTop(context.Background(), i.ID, nil)
if err != nil {
continue
}
index := slices.Index(top.Titles, "PID")
if index == -1 {
continue
}
c.Pid = make([]int, 0, len(top.Processes))
for _, v := range top.Processes {
pid, err := strconv.Atoi(v[index])
if err == nil {
c.Pid = append(c.Pid, pid)
}
}
}
cli.Close()
var cardInfo map[int]*logic.CardInfo
if a := logic.DCUInfo(); a != nil {
cardInfo = a
} else if a := logic.NVIDIAInfo(); a != nil {
cardInfo = a
}
if cardInfo != nil {
for _, v := range result {
for _, vin := range cardInfo {
if contain(vin.Pids, v.Pid) {
v.Cards = append(v.Cards, vin)
}
}
}
}
printInfo(result)
}
func contain(s1, s2 []int) bool {
if len(s1) == 0 || len(s2) == 0 {
return false
}
for _, v := range s1 {
if slices.Contains(s2, v) {
return true
}
}
return false
}
/*
容器名 容器id 创建者 使用的卡 容器中的进程数
ContainerName ContainerID CreateUser UsedCard PidsInContainer
*/
// printInfo 打印信息
func printInfo(info map[string]*Cinfo) {
titleStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("#6071f1"))
s1 := lipgloss.NewStyle().Foreground(lipgloss.Color("#82ff60"))
s2 := lipgloss.NewStyle().Foreground(lipgloss.Color("#54fff6"))
s3 := lipgloss.NewStyle().Foreground(lipgloss.Color("#eef867"))
title := []string{"ContainerName", "ContainerID", "CreateUser", "UsedCard", "PidsInContainer"}
maxWidth := []int{0, 0, 0, 0, 0}
for k, v := range title {
maxWidth[k] = len(v)
}
out := make([][5]string, 0, len(info))
for _, v := range info {
a, b := v.Format()
out = append(out, a)
for x, y := range b {
maxWidth[x] = max(maxWidth[x], y)
}
}
t := fmt.Sprintf("%s %s %s %s %s", formatStr(title[0], maxWidth[0]), formatStr(title[1], maxWidth[1]), formatStr(title[2], maxWidth[2]), formatStr(title[3], maxWidth[3]), formatStr(title[4], maxWidth[4]))
fmt.Println(titleStyle.Render(t))
sort.Slice(out, func(i, j int) bool {
return out[i][1] < out[j][1]
})
for k, v := range out {
o := fmt.Sprintf("%s %s %s %s %s", formatStr(v[0], maxWidth[0]),
formatStr(v[1], maxWidth[1]),
formatStr(v[2], maxWidth[2]),
formatStr(v[3], maxWidth[3]),
formatStr(v[4], maxWidth[4]))
switch k % 3 {
case 0:
fmt.Println(s1.Render(o))
case 1:
fmt.Println(s2.Render(o))
case 2:
fmt.Println(s3.Render(o))
}
}
}
func formatStr(raw string, l int) string {
lstr := len(raw)
if l >= lstr {
return fmt.Sprintf("%s%s", raw, strings.Repeat(" ", l-lstr))
}
return raw[:l]
}
func RemoveDuplicates[T comparable](s []T) []T {
m := make(map[T]struct{})
index := 0
for _, v := range s {
if _, have := m[v]; !have {
m[v] = struct{}{}
s[index] = v
index++
}
}
return s[:index]
}
...@@ -193,6 +193,7 @@ var ( ...@@ -193,6 +193,7 @@ var (
reDockerApi = regexp.MustCompile(`(?i)^\s+api\s+version:\s+([0-9.]+).*$`) reDockerApi = regexp.MustCompile(`(?i)^\s+api\s+version:\s+([0-9.]+).*$`)
) )
// GetDockerAPIVersion 获取docker服务端API版本
func GetDockerAPIVersion() (string, error) { func GetDockerAPIVersion() (string, error) {
output, err := exec.Command("docker", "version").Output() output, err := exec.Command("docker", "version").Output()
if err != nil { if err != nil {
...@@ -208,6 +209,7 @@ func GetDockerAPIVersion() (string, error) { ...@@ -208,6 +209,7 @@ func GetDockerAPIVersion() (string, error) {
return "", nil return "", nil
} }
// GetDockerClient 获取docker客户端
func GetDockerClient() (*client.Client, error) { func GetDockerClient() (*client.Client, error) {
ver, err := GetDockerAPIVersion() ver, err := GetDockerAPIVersion()
if err != nil { if err != nil {
......
...@@ -120,3 +120,24 @@ func TestGetDockerAPIVersion(t *testing.T) { ...@@ -120,3 +120,24 @@ func TestGetDockerAPIVersion(t *testing.T) {
} }
t.Log(v) t.Log(v)
} }
func TestTop(t *testing.T) {
cli, err := GetDockerClient()
if err != nil {
t.Error(err)
}
cs, err := cli.ContainerList(context.Background(), client.ContainerListOptions{})
if err != nil {
cli.Close()
t.Error(err)
}
for _, v := range cs {
top, err := cli.ContainerTop(context.Background(), v.ID, nil)
if err != nil {
continue
}
t.Logf("%v", top)
}
cli.Close()
}
...@@ -9,6 +9,8 @@ require ( ...@@ -9,6 +9,8 @@ require (
github.com/lrstanley/bubblezone v0.0.0-20240914071701-b48c55a5e78e github.com/lrstanley/bubblezone v0.0.0-20240914071701-b48c55a5e78e
github.com/moby/moby/api v1.52.0-beta.2 github.com/moby/moby/api v1.52.0-beta.2
github.com/moby/moby/client v0.1.0-beta.2 github.com/moby/moby/client v0.1.0-beta.2
github.com/pquerna/otp v1.5.0
github.com/samber/mo v1.16.0
github.com/shirou/gopsutil/v3 v3.24.5 github.com/shirou/gopsutil/v3 v3.24.5
github.com/shirou/gopsutil/v4 v4.25.9 github.com/shirou/gopsutil/v4 v4.25.9
github.com/spf13/viper v1.21.0 github.com/spf13/viper v1.21.0
...@@ -48,11 +50,9 @@ require ( ...@@ -48,11 +50,9 @@ require (
github.com/muesli/clusters v0.0.0-20200529215643-2700303c1762 // indirect github.com/muesli/clusters v0.0.0-20200529215643-2700303c1762 // indirect
github.com/muesli/kmeans v0.3.1 // indirect github.com/muesli/kmeans v0.3.1 // indirect
github.com/pelletier/go-toml/v2 v2.2.4 // indirect github.com/pelletier/go-toml/v2 v2.2.4 // indirect
github.com/pquerna/otp v1.5.0 // indirect
github.com/quic-go/qpack v0.5.1 // indirect github.com/quic-go/qpack v0.5.1 // indirect
github.com/quic-go/quic-go v0.54.0 // indirect github.com/quic-go/quic-go v0.54.0 // indirect
github.com/sagikazarmark/locafero v0.11.0 // indirect github.com/sagikazarmark/locafero v0.11.0 // indirect
github.com/samber/mo v1.16.0 // indirect
github.com/shoenig/go-m1cpu v0.1.6 // indirect github.com/shoenig/go-m1cpu v0.1.6 // indirect
github.com/sourcegraph/conc v0.3.1-0.20240121214520-5f936abd7ae8 // indirect github.com/sourcegraph/conc v0.3.1-0.20240121214520-5f936abd7ae8 // indirect
github.com/spf13/afero v1.15.0 // indirect github.com/spf13/afero v1.15.0 // indirect
......
...@@ -3,6 +3,6 @@ package utils ...@@ -3,6 +3,6 @@ package utils
import "testing" import "testing"
func TestDetectCmd(t *testing.T) { func TestDetectCmd(t *testing.T) {
a, b := DetectCmd("ps") a, b := DetectCmd("hy-smi")
t.Logf("%v,%v", a, b) t.Logf("%v,%v", a, b)
} }
...@@ -2,6 +2,7 @@ package utils ...@@ -2,6 +2,7 @@ package utils
import ( import (
"fmt" "fmt"
"regexp"
"slices" "slices"
"strings" "strings"
"sync/atomic" "sync/atomic"
...@@ -215,5 +216,19 @@ func TestCHar(t *testing.T) { ...@@ -215,5 +216,19 @@ func TestCHar(t *testing.T) {
} }
func TestSplitN(t *testing.T) { func TestSplitN(t *testing.T) {
t.Logf("%v",strings.SplitN("a=b=c","=",2)) t.Logf("%v", strings.SplitN("a=b=c", "=", 2))
} }
\ No newline at end of file
func TestReg(t *testing.T) {
var RegUser = regexp.MustCompile(`^(?i)/public[0-9]*/([0-9a-z]+)(?:|/.*)$`)
test := []string{"/public2/liming6", "/Public/Liming6/hello"}
for _, v := range test {
if RegUser.MatchString(v) {
f := RegUser.FindStringSubmatch(v)
t.Logf("match %s, %v", v, f)
} else {
t.Logf("not match %s", v)
}
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment