功能：完善云管助手与GPU故障恢复核对

2026-11-13 19:08:40 +08:00
1 parent b6026f9759
commit 9f7b01e4db
@@ -0,0 +1,141 @@
+package server
+
+import (
+	"context"
+	"crypto/subtle"
+	"strings"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"vdi-cloud/platform/pkg/proxmox"
+)
+
+// gpuAgentSample is the small, vendor-neutral contract implemented by a node
+// collector. The collector may use nvidia-smi, rocm-smi, intel_gpu_top or a
+// vendor API, while the platform keeps the same health model for every GPU.
+type gpuAgentSample struct {
+	Address            string   `json:"address" binding:"required"`
+	ObservedAt         string   `json:"observedAt"`
+	TemperatureC       *float64 `json:"temperatureC"`
+	UtilizationPercent *float64 `json:"utilizationPercent"`
+	MemoryUsedBytes    *int64   `json:"memoryUsedBytes"`
+	MemoryTotalBytes   *int64   `json:"memoryTotalBytes"`
+	Driver             string   `json:"driver"`
+	DriverBinding      string   `json:"driverBinding"`
+}
+
+type gpuAgentReport struct {
+	Node    string           `json:"node" binding:"required,max=100"`
+	Samples []gpuAgentSample `json:"samples" binding:"required,min=1,max=256"`
+}
+
+func (s *Server) gpuAgentAuth(c *gin.Context) {
+	if strings.TrimSpace(s.Config.GPUAgentToken) == "" {
+		s.fail(c, 503, "gpu_agent_unconfigured", "尚未配置 GPU 节点采集令牌")
+		return
+	}
+	provided := c.GetHeader("X-GPU-Agent-Token")
+	if len(provided) != len(s.Config.GPUAgentToken) || subtle.ConstantTimeCompare([]byte(provided), []byte(s.Config.GPUAgentToken)) != 1 {
+		s.fail(c, 401, "gpu_agent_unauthorized", "GPU 节点采集令牌无效")
+		return
+	}
+	c.Next()
+}
+
+func validGPUHealth(sample gpuAgentSample) bool {
+	if sample.TemperatureC != nil && (*sample.TemperatureC < -40 || *sample.TemperatureC > 150) {
+		return false
+	}
+	if sample.UtilizationPercent != nil && (*sample.UtilizationPercent < 0 || *sample.UtilizationPercent > 100) {
+		return false
+	}
+	if sample.MemoryUsedBytes != nil && *sample.MemoryUsedBytes < 0 {
+		return false
+	}
+	if sample.MemoryTotalBytes != nil && *sample.MemoryTotalBytes < 0 {
+		return false
+	}
+	if sample.MemoryUsedBytes != nil && sample.MemoryTotalBytes != nil && *sample.MemoryUsedBytes > *sample.MemoryTotalBytes {
+		return false
+	}
+	_, ok := proxmox.NormalizePCIAddress(sample.Address)
+	return ok
+}
+
+func (s *Server) ingestGPUHealth(c *gin.Context) {
+	var in gpuAgentReport
+	if !s.decode(c, &in) {
+		return
+	}
+	in.Node = strings.TrimSpace(in.Node)
+	if in.Node == "" || strings.ContainsAny(in.Node, "/\\?# ") {
+		s.fail(c, 400, "invalid_node", "节点名称无效")
+		return
+	}
+	ctx := c.Request.Context()
+	tx, err := s.DB.BeginTxx(ctx, nil)
+	if err != nil {
+		s.dbError(c, err)
+		return
+	}
+	defer tx.Rollback()
+	for _, sample := range in.Samples {
+		if !validGPUHealth(sample) {
+			s.fail(c, 400, "invalid_gpu_health", "GPU 健康数据超出范围或 PCI 地址无效")
+			return
+		}
+		address, _ := proxmox.NormalizePCIAddress(sample.Address)
+		observedAt := time.Now().UTC()
+		if strings.TrimSpace(sample.ObservedAt) != "" {
+			observedAt, err = time.Parse(time.RFC3339Nano, sample.ObservedAt)
+			if err != nil || observedAt.After(time.Now().Add(2*time.Minute)) || observedAt.Before(time.Now().Add(-10*time.Minute)) {
+				s.fail(c, 400, "invalid_gpu_health", "采集时间必须是最近 10 分钟内的 RFC3339 时间")
+				return
+			}
+			observedAt = observedAt.UTC()
+		}
+		_, err = tx.ExecContext(ctx, `INSERT INTO gpu_health_samples(node,address,observed_at,temperature_c,utilization_percent,memory_used_bytes,memory_total_bytes,driver,driver_binding,source)
+VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,'node-agent') ON CONFLICT DO NOTHING`, in.Node, address, observedAt, sample.TemperatureC, sample.UtilizationPercent, sample.MemoryUsedBytes, sample.MemoryTotalBytes, strings.TrimSpace(sample.Driver), strings.TrimSpace(sample.DriverBinding))
+		if err != nil {
+			s.dbError(c, err)
+			return
+		}
+	}
+	if _, err = tx.ExecContext(ctx, "DELETE FROM gpu_health_samples WHERE observed_at < now() - interval '7 days'"); err != nil {
+		s.dbError(c, err)
+		return
+	}
+	if err = tx.Commit(); err != nil {
+		s.dbError(c, err)
+		return
+	}
+	c.JSON(202, gin.H{"accepted": len(in.Samples), "node": in.Node})
+}
+
+type latestGPUHealth struct {
+	Address            string    `db:"address"`
+	ObservedAt         time.Time `db:"observed_at"`
+	TemperatureC       *float64  `db:"temperature_c"`
+	UtilizationPercent *float64  `db:"utilization_percent"`
+	MemoryUsedBytes    *int64    `db:"memory_used_bytes"`
+	MemoryTotalBytes   *int64    `db:"memory_total_bytes"`
+	Driver             string    `db:"driver"`
+	DriverBinding      string    `db:"driver_binding"`
+	Source             string    `db:"source"`
+}
+
+func (s *Server) latestGPUHealth(ctx context.Context, node string) (map[string]latestGPUHealth, error) {
+	// Kept separate from inventory assembly so a missing collector is an
+	// ordinary empty result and remains visibly "not_collected" in the API.
+	items := []latestGPUHealth{}
+	query := `SELECT DISTINCT ON(address) address,observed_at,temperature_c,utilization_percent,memory_used_bytes,memory_total_bytes,driver,driver_binding,source
+FROM gpu_health_samples WHERE node=$1 AND observed_at > now()-interval '10 minutes' ORDER BY address,observed_at DESC`
+	if err := s.DB.SelectContext(ctx, &items, query, node); err != nil {
+		return nil, err
+	}
+	result := make(map[string]latestGPUHealth, len(items))
+	for _, item := range items {
+		result[item.Address] = item
+	}
+	return result, nil
+}