初始提交：边缘AI算力机统一AI通讯层

2026-08-03 07:44:05 +08:00
commit 93a469061d
51 changed files with 11565 additions and 0 deletions
@@ -0,0 +1,197 @@
+package resource
+
+import (
+	"context"
+	"fmt"
+	"os/exec"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+)
+
+// GPUMetrics represents GPU utilization data from nvidia-smi.
+type GPUMetrics struct {
+	Index         int
+	Name           string
+	TemperatureC   int
+	UtilizationGPU int // percentage 0-100
+	MemoryUsedMB   int
+	MemoryTotalMB  int
+	MemoryUtilPct  float64
+	PowerDrawW     float64
+	PowerLimitW    float64
+	Timestamp      time.Time
+}
+
+// GPUCollector collects GPU metrics via nvidia-smi.
+type GPUCollector struct {
+	mu      sync.RWMutex
+	metrics []GPUMetrics
+	enabled bool
+}
+
+// NewGPUCollector creates a new GPU collector.
+func NewGPUCollector() *GPUCollector {
+	return &GPUCollector{enabled: true}
+}
+
+// Collect runs nvidia-smi and parses the output.
+func (c *GPUCollector) Collect(ctx context.Context) ([]GPUMetrics, error) {
+	if !c.enabled {
+		return nil, nil
+	}
+
+	// Use nvidia-smi with CSV format for structured output
+	cmd := exec.CommandContext(ctx, "nvidia-smi",
+		"--query-gpu=index,name,temperature.gpu,utilization.gpu,memory.used,memory.total,memory.utilization,power.draw,power.limit",
+		"--format=csv,noheader,nounits",
+	)
+
+	output, err := cmd.Output()
+	if err != nil {
+		// If nvidia-smi is not available, disable collector
+		c.mu.Lock()
+		c.enabled = false
+		c.mu.Unlock()
+		return nil, fmt.Errorf("nvidia-smi not available: %w", err)
+	}
+
+	metrics := parseNvidiaSMI(string(output))
+
+	c.mu.Lock()
+	c.metrics = metrics
+	c.mu.Unlock()
+
+	return metrics, nil
+}
+
+func parseNvidiaSMI(output string) []GPUMetrics {
+	lines := strings.Split(strings.TrimSpace(output), "\n")
+	metrics := make([]GPUMetrics, 0, len(lines))
+
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if line == "" {
+			continue
+		}
+
+		fields := strings.Split(line, ",")
+		if len(fields) < 9 {
+			continue
+		}
+
+		m := GPUMetrics{Timestamp: time.Now()}
+		m.Index = parseIntSafe(fields[0])
+		m.Name = strings.TrimSpace(fields[1])
+		m.TemperatureC = parseIntSafe(fields[2])
+		m.UtilizationGPU = parseIntSafe(fields[3])
+		m.MemoryUsedMB = parseIntSafe(fields[4])
+		m.MemoryTotalMB = parseIntSafe(fields[5])
+		m.MemoryUtilPct = parseFloatSafe(fields[6])
+		m.PowerDrawW = parseFloatSafe(fields[7])
+		m.PowerLimitW = parseFloatSafe(fields[8])
+
+		metrics = append(metrics, m)
+	}
+
+	return metrics
+}
+
+func parseIntSafe(s string) int {
+	s = strings.TrimSpace(s)
+	v, err := strconv.Atoi(s)
+	if err != nil {
+		return 0
+	}
+	return v
+}
+
+func parseFloatSafe(s string) float64 {
+	s = strings.TrimSpace(s)
+	v, err := strconv.ParseFloat(s, 64)
+	if err != nil {
+		return 0
+	}
+	return v
+}
+
+// GetMetrics returns the last collected metrics (thread-safe).
+func (c *GPUCollector) GetMetrics() []GPUMetrics {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.metrics
+}
+
+// IsEnabled returns whether GPU collection is enabled.
+func (c *GPUCollector) IsEnabled() bool {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	return c.enabled
+}
+
+// StartPeriodicCollection starts a background goroutine that collects GPU metrics at regular intervals.
+func (c *GPUCollector) StartPeriodicCollection(ctx context.Context, interval time.Duration) {
+	go func() {
+		ticker := time.NewTicker(interval)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				c.Collect(ctx)
+			}
+		}
+	}()
+}
+
+// TotalMemoryUsedMB returns total GPU memory used across all GPUs.
+func (c *GPUCollector) TotalMemoryUsedMB() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	total := 0
+	for _, m := range c.metrics {
+		total += m.MemoryUsedMB
+	}
+	return total
+}
+
+// TotalMemoryTotalMB returns total GPU memory capacity across all GPUs.
+func (c *GPUCollector) TotalMemoryTotalMB() int {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	total := 0
+	for _, m := range c.metrics {
+		total += m.MemoryTotalMB
+	}
+	return total
+}
+
+// AverageUtilization returns average GPU utilization percentage.
+func (c *GPUCollector) AverageUtilization() float64 {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+
+	if len(c.metrics) == 0 {
+		return 0
+	}
+
+	total := 0
+	for _, m := range c.metrics {
+		total += m.UtilizationGPU
+	}
+	return float64(total) / float64(len(c.metrics))
+}
+
+// MemoryUtilizationRatio returns memory used / memory total (0.0-1.0).
+func (c *GPUCollector) MemoryUtilizationRatio() float64 {
+	total := c.TotalMemoryTotalMB()
+	if total == 0 {
+		return 0
+	}
+	return float64(c.TotalMemoryUsedMB()) / float64(total)
+}
@@ -0,0 +1,130 @@
+package resource
+
+import (
+	"testing"
+)
+
+func TestParseNvidiaSMI(t *testing.T) {
+	output := `0, NVIDIA GeForce RTX 4090, 45, 30, 4096, 24576, 16.67, 150.5, 450.0
+1, NVIDIA GeForce RTX 4090, 52, 75, 8192, 24576, 33.33, 320.0, 450.0`
+
+	metrics := parseNvidiaSMI(output)
+	if len(metrics) != 2 {
+		t.Fatalf("expected 2 GPUs, got %d", len(metrics))
+	}
+
+	if metrics[0].Index != 0 {
+		t.Errorf("expected index 0, got %d", metrics[0].Index)
+	}
+	if metrics[0].Name != "NVIDIA GeForce RTX 4090" {
+		t.Errorf("unexpected name: %s", metrics[0].Name)
+	}
+	if metrics[0].TemperatureC != 45 {
+		t.Errorf("expected temp 45, got %d", metrics[0].TemperatureC)
+	}
+	if metrics[0].UtilizationGPU != 30 {
+		t.Errorf("expected util 30, got %d", metrics[0].UtilizationGPU)
+	}
+	if metrics[0].MemoryUsedMB != 4096 {
+		t.Errorf("expected mem used 4096, got %d", metrics[0].MemoryUsedMB)
+	}
+	if metrics[0].MemoryTotalMB != 24576 {
+		t.Errorf("expected mem total 24576, got %d", metrics[0].MemoryTotalMB)
+	}
+	if metrics[0].PowerDrawW != 150.5 {
+		t.Errorf("expected power 150.5, got %f", metrics[0].PowerDrawW)
+	}
+
+	if metrics[1].Index != 1 {
+		t.Errorf("expected index 1, got %d", metrics[1].Index)
+	}
+	if metrics[1].UtilizationGPU != 75 {
+		t.Errorf("expected util 75, got %d", metrics[1].UtilizationGPU)
+	}
+}
+
+func TestParseNvidiaSMIEmpty(t *testing.T) {
+	metrics := parseNvidiaSMI("")
+	if len(metrics) != 0 {
+		t.Errorf("expected 0 metrics for empty input, got %d", len(metrics))
+	}
+}
+
+func TestParseNvidiaSMIInvalidLines(t *testing.T) {
+	output := `invalid line
+0, GPU0, 40, 50, 1024, 8192, 12.5, 100.0, 300.0
+, , , , , , , , `
+
+	metrics := parseNvidiaSMI(output)
+	// Both lines with 9 fields parse; the empty-name one has Name=""
+	validCount := 0
+	for _, m := range metrics {
+		if m.Name != "" {
+			validCount++
+		}
+	}
+	if validCount != 1 {
+		t.Errorf("expected 1 valid metric with name, got %d", validCount)
+	}
+}
+
+func TestGPUCollectorTotals(t *testing.T) {
+	c := &GPUCollector{
+		metrics: []GPUMetrics{
+			{MemoryUsedMB: 4096, MemoryTotalMB: 24576, UtilizationGPU: 30},
+			{MemoryUsedMB: 8192, MemoryTotalMB: 24576, UtilizationGPU: 75},
+		},
+	}
+
+	if c.TotalMemoryUsedMB() != 12288 {
+		t.Errorf("expected 12288, got %d", c.TotalMemoryUsedMB())
+	}
+	if c.TotalMemoryTotalMB() != 49152 {
+		t.Errorf("expected 49152, got %d", c.TotalMemoryTotalMB())
+	}
+
+	avg := c.AverageUtilization()
+	if avg != 52.5 {
+		t.Errorf("expected 52.5, got %f", avg)
+	}
+
+	ratio := c.MemoryUtilizationRatio()
+	expectedRatio := 12288.0 / 49152.0
+	if ratio != expectedRatio {
+		t.Errorf("expected %f, got %f", expectedRatio, ratio)
+	}
+}
+
+func TestGPUCollectorEmpty(t *testing.T) {
+	c := &GPUCollector{}
+	if c.TotalMemoryUsedMB() != 0 {
+		t.Error("expected 0 for empty collector")
+	}
+	if c.AverageUtilization() != 0 {
+		t.Error("expected 0 for empty collector")
+	}
+	if c.MemoryUtilizationRatio() != 0 {
+		t.Error("expected 0 for empty collector")
+	}
+}
+
+func TestParseIntSafe(t *testing.T) {
+	if parseIntSafe("42") != 42 {
+		t.Error("expected 42")
+	}
+	if parseIntSafe("invalid") != 0 {
+		t.Error("expected 0 for invalid")
+	}
+	if parseIntSafe("  100  ") != 100 {
+		t.Error("expected 100 with whitespace")
+	}
+}
+
+func TestParseFloatSafe(t *testing.T) {
+	if parseFloatSafe("3.14") != 3.14 {
+		t.Error("expected 3.14")
+	}
+	if parseFloatSafe("invalid") != 0 {
+		t.Error("expected 0 for invalid")
+	}
+}