llmgw/internal/web/gpumap.go

package web

import (
	"fmt"
	"html/template"
	"sort"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
)

// The GPU map is every card in the fleet, a square each, a row per node,
// coloured by one figure: how busy, how hot, or how full its memory is. A
// hundred cards in a table are a hundred rows to read; on the map the one
// that is not like the others is seen before anything is read.

// gpuMapModes are what the map can be coloured by.
var gpuMapModes = []struct{ Key, Label string }{
	{"util", "busy"}, {"memory", "memory"}, {"temp", "temperature"},
}

type gpuMap struct {
	By    string
	Modes []struct{ Key, Label string }
	Rows  []gpuMapRow
	// Legend says what the colours mean in this mode.
	Legend string
}

type gpuMapRow struct {
	Node  string
	Cells []gpuCell
}

type gpuCell struct {
	Index int
	// Model and Replica hold the card; empty is a free card.
	Model, Replica string
	// Style is the cell's colour, built here from numbers alone, which is
	// what makes it safe to put in a style attribute; Title says everything
	// known about the card.
	Style template.CSS
	Title string
	// Label is the figure written in the cell.
	Label string
	// Faulty is a card whose wear needs seeing to.
	Faulty bool
}

// idleHolding is a replica holding GPUs that are doing nothing.
type idleHolding struct {
	Model, Replica, Node string
	GPUs                 []int
	// MaxUtil is the busiest of its cards, now.
	MaxUtil float64
	// Idle is how long since a request was routed to it, or "" when not
	// known.
	Idle string
}

// buildGPUMap draws the map from what the nodes last reported and where
// replicas are. Nil for a fleet without GPUs.
func buildGPUMap(nodes []adminapi.NodeView, reps []adminapi.ReplicaView, by, prefix string) *gpuMap {
	held := heldBy(reps)
	m := &gpuMap{By: by, Modes: gpuMapModes}
	if m.By == "" {
		m.By = "util"
	}
	switch m.By {
	case "temp":
		m.Legend = "Blue is cool, red is 90 °C and over."
	case "memory":
		m.Legend = "Pale is empty, dark is full. Engines reserve most of a card's memory when they start."
	default:
		m.By = "util"
		m.Legend = "Pale is idle, dark green is fully busy. A dashed border is a card no replica holds."
	}
	for _, n := range nodes {
		if n.GPUTotal == 0 {
			continue
		}
		row := gpuMapRow{Node: strings.TrimPrefix(n.ID, prefix)}
		for i := range n.GPUTotal {
			c := gpuCell{Index: i}
			if h, ok := held[gpuKey(n.ID, i)]; ok {
				c.Model, c.Replica = h.model, h.replica
			}
			var title []string
			title = append(title, fmt.Sprintf("%s GPU %d", n.ID, i))
			if c.Model != "" {
				title = append(title, c.Model+" ("+c.Replica+")")
			} else {
				title = append(title, "free")
			}
			s, sampled := sensorOf(n, i)
			if sampled {
				if s.MemoryTotalMB > 0 && s.UtilPercent >= 0 {
					title = append(title, fmt.Sprintf("%.0f%% busy", s.UtilPercent))
				}
				if s.MemoryTotalMB > 0 {
					title = append(title, fmt.Sprintf("%.1f of %.0f GB", s.MemoryUsedMB/1024, s.MemoryTotalMB/1024))
				}
				if s.TempC > 0 {
					title = append(title, fmt.Sprintf("%.0f °C", s.TempC))
				}
				if s.PowerW > 0 {
					title = append(title, fmt.Sprintf("%.0f W", s.PowerW))
				}
				if len(s.Throttle) > 0 {
					title = append(title, "held back by "+strings.Join(s.Throttle, ", "))
				}
				title = append(title, s.Problems...)
			}
			c.Title = strings.Join(title, " · ")
			style, label := cellColour(m.By, s, sampled)
			if c.Model == "" {
				style += ";border-style:dashed"
			}
			// A card wearing out is outlined whatever the map is coloured
			// by: it is the one thing on the map that needs doing.
			if len(s.Problems) > 0 {
				style += ";border:2px solid #c0392b"
				c.Faulty = true
			}
			c.Style, c.Label = template.CSS(style), label
			row.Cells = append(row.Cells, c)
		}
		m.Rows = append(m.Rows, row)
	}
	if len(m.Rows) == 0 {
		return nil
	}
	return m
}

// cellColour is a cell's background, and the figure written in it.
func cellColour(by string, s sensorView, sampled bool) (string, string) {
	const unknown = "background:var(--panel-2,#eee)"
	switch by {
	case "temp":
		if !sampled || s.TempC <= 0 {
			return unknown, ""
		}
		// 30 °C blue to 90 °C red.
		f := clamp01((s.TempC - 30) / 60)
		return fmt.Sprintf("background:hsl(%.0f,70%%,55%%)", 220-220*f), fmt.Sprintf("%.0f", s.TempC)
	case "memory":
		if !sampled || s.MemoryTotalMB <= 0 {
			return unknown, ""
		}
		f := clamp01(s.MemoryUsedMB / s.MemoryTotalMB)
		return fmt.Sprintf("background:hsl(265,45%%,%.0f%%)", 90-55*f), fmt.Sprintf("%.0f", 100*f)
	}
	if !sampled || s.MemoryTotalMB <= 0 || s.UtilPercent < 0 {
		return unknown, ""
	}
	f := clamp01(s.UtilPercent / 100)
	return fmt.Sprintf("background:hsl(140,%.0f%%,%.0f%%)", 15+50*f, 88-50*f), fmt.Sprintf("%.0f", s.UtilPercent)
}

// idleHoldings are the replicas whose GPUs all read under 5% busy now:
// capacity held and not asked for, the first place to find room.
func idleHoldings(nodes []adminapi.NodeView, reps []adminapi.ReplicaView, prefix string) []idleHolding {
	byID := map[string]adminapi.NodeView{}
	for _, n := range nodes {
		byID[n.ID] = n
	}
	var out []idleHolding
	for _, r := range reps {
		if r.State != "ready" {
			continue
		}
		for _, sh := range r.Shards {
			if len(sh.GPUs) == 0 {
				continue
			}
			h := idleHolding{Model: r.Model, Replica: r.ID, Node: strings.TrimPrefix(sh.Node, prefix), GPUs: sh.GPUs}
			idle, known := true, true
			for _, g := range sh.GPUs {
				s, ok := sensorOf(byID[sh.Node], g)
				if !ok || s.MemoryTotalMB <= 0 || s.UtilPercent < 0 {
					known = false
					break
				}
				h.MaxUtil = max(h.MaxUtil, s.UtilPercent)
				if s.UtilPercent >= 5 {
					idle = false
				}
			}
			if !known || !idle {
				continue
			}
			if r.IdleSec >= 0 {
				h.Idle = (time.Duration(r.IdleSec) * time.Second).String()
			}
			out = append(out, h)
		}
	}
	sort.SliceStable(out, func(i, j int) bool { return len(out[i].GPUs) > len(out[j].GPUs) })
	return out
}

type holder struct{ model, replica string }

func heldBy(reps []adminapi.ReplicaView) map[string]holder {
	out := map[string]holder{}
	for _, r := range reps {
		switch r.State {
		case "stopped", "failed":
			continue
		}
		for _, sh := range r.Shards {
			for _, g := range sh.GPUs {
				out[gpuKey(sh.Node, g)] = holder{r.Model, r.ID}
			}
		}
	}
	return out
}

func gpuKey(node string, gpu int) string { return fmt.Sprintf("%s/%d", node, gpu) }

// sensorView is a card's last sample.
type sensorView = struct {
	TempC, PowerW, UtilPercent, MemoryUsedMB, MemoryTotalMB float64
	Throttle                                                []string
	Problems                                                []string
}

func sensorOf(n adminapi.NodeView, gpu int) (sensorView, bool) {
	if n.Load == nil {
		return sensorView{}, false
	}
	for _, s := range n.Load.GPUs {
		if s.Index == gpu {
			return sensorView{TempC: s.TempC, PowerW: s.PowerW, UtilPercent: s.UtilPercent,
				MemoryUsedMB: s.MemoryUsedMB, MemoryTotalMB: s.MemoryTotalMB, Throttle: s.Throttle,
				Problems: s.Wear.Problems()}, true
		}
	}
	return sensorView{}, false
}