mirror of
https://github.com/1Panel-dev/1Panel.git
synced 2026-09-23 08:00:55 +00:00
116 lines
2.5 KiB
Go
116 lines
2.5 KiB
Go
package gpu
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"strings"
|
||
"sync"
|
||
|
||
"github.com/1Panel-dev/1Panel/agent/utils/cmd"
|
||
)
|
||
|
||
type provider interface {
|
||
LoadInfo(context.Context) (*Info, error)
|
||
}
|
||
|
||
type Client struct {
|
||
providers []provider
|
||
}
|
||
|
||
type providerResult struct {
|
||
info *Info
|
||
err error
|
||
}
|
||
|
||
func New() (bool, Client) {
|
||
client := Client{}
|
||
if cmd.Which(nvidiaSMICommand) {
|
||
client.providers = append(client.providers, nvidiaSMI{})
|
||
}
|
||
if command, ok := findAMDSMI(); ok {
|
||
client.providers = append(client.providers, amdSMI{command: command})
|
||
}
|
||
return len(client.providers) > 0, client
|
||
}
|
||
|
||
func (c Client) LoadInfo() (*Info, error) {
|
||
return c.LoadInfoContext(context.Background())
|
||
}
|
||
|
||
func (c Client) LoadInfoContext(ctx context.Context) (*Info, error) {
|
||
results := make([]providerResult, len(c.providers))
|
||
var wg sync.WaitGroup
|
||
for index, item := range c.providers {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
results[index].info, results[index].err = item.LoadInfo(ctx)
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
|
||
merged := &Info{}
|
||
var (
|
||
errs []error
|
||
types []string
|
||
driverVersions []string
|
||
)
|
||
for _, result := range results {
|
||
if result.err != nil {
|
||
errs = append(errs, result.err)
|
||
continue
|
||
}
|
||
if result.info == nil {
|
||
continue
|
||
}
|
||
if result.info.Type != "" {
|
||
types = append(types, result.info.Type)
|
||
}
|
||
if result.info.DriverVersion != "" {
|
||
driverVersions = append(driverVersions, formatDriverVersion(result.info.Type, result.info.DriverVersion))
|
||
}
|
||
if result.info.CudaVersion != "" {
|
||
merged.CudaVersion = result.info.CudaVersion
|
||
}
|
||
for _, device := range result.info.Devices {
|
||
if device.Type == "" {
|
||
device.Type = result.info.Type
|
||
}
|
||
merged.Devices = append(merged.Devices, device)
|
||
}
|
||
}
|
||
|
||
merged.Type = mergeTypes(types)
|
||
merged.DriverVersion = mergeDriverVersions(driverVersions)
|
||
if len(merged.Devices) == 0 && len(errs) > 0 {
|
||
return nil, fmt.Errorf("calling GPU monitoring tools failed: %w", errors.Join(errs...))
|
||
}
|
||
return merged, nil
|
||
}
|
||
|
||
func formatDriverVersion(deviceType, version string) string {
|
||
if deviceType == "" {
|
||
return version
|
||
}
|
||
return fmt.Sprintf("%s: %s", strings.ToUpper(deviceType), version)
|
||
}
|
||
|
||
func mergeTypes(types []string) string {
|
||
if len(types) == 1 {
|
||
return types[0]
|
||
}
|
||
if len(types) > 1 {
|
||
return "mixed"
|
||
}
|
||
return ""
|
||
}
|
||
|
||
func mergeDriverVersions(versions []string) string {
|
||
if len(versions) == 1 {
|
||
parts := strings.SplitN(versions[0], ": ", 2)
|
||
return parts[len(parts)-1]
|
||
}
|
||
return strings.Join(versions, "；")
|
||
}
|