AIRouter/internal/resource/gpu.go

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)
}