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; each is named by the key
// dash.gpumap.mode.<mode>.
var gpuMapModes = []string{"util", "memory", "temp"}

func init() {
	for _, m := range gpuMapModes {
		registerKeys("dash.gpumap.mode." + m)
	}
}

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(l *Lang, nodes []adminapi.NodeView, reps []adminapi.ReplicaView, by, prefix string) *gpuMap {
	held := heldBy(reps)
	m := &gpuMap{By: by}
	for _, k := range gpuMapModes {
		m.Modes = append(m.Modes, struct{ Key, Label string }{k, l.T("dash.gpumap.mode." + k)})
	}
	if m.By == "" {
		m.By = "util"
	}
	switch m.By {
	case "temp":
		m.Legend = l.T("dash.gpumap.legend.temp")
	case "memory":
		m.Legend = l.T("dash.gpumap.legend.memory")
	default:
		m.By = "util"
		m.Legend = l.T("dash.gpumap.legend.util")
	}
	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, l.T("dash.gpumap.free"))
			}
			s, sampled := sensorOf(n, i)
			if sampled {
				if s.MemoryTotalMB > 0 && s.UtilPercent >= 0 {
					title = append(title, l.T("dash.gpumap.busy", l.Num(s.UtilPercent, 0)))
				}
				if s.MemoryTotalMB > 0 {
					title = append(title, l.T("dash.gpumap.memory", l.Num(s.MemoryUsedMB/1024, 1), l.Num(s.MemoryTotalMB/1024, 0)))
				}
				if s.TempC > 0 {
					title = append(title, l.Num(s.TempC, 0)+" °C")
				}
				if s.PowerW > 0 {
					title = append(title, l.Num(s.PowerW, 0)+" W")
				}
				if len(s.Throttle) > 0 {
					// The driver's own names for the reasons, which are
					// the same in every language.
					title = append(title, l.T("dash.gpumap.held_back", strings.Join(s.Throttle, ", ")))
				}
				// What wears a card out is said by the inventory, in English.
				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(l *Lang, 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 = l.clock(time.Duration(r.IdleSec) * time.Second)
			}
			out = append(out, h)
		}
	}
	sort.SliceStable(out, func(i, j int) bool { return len(out[i].GPUs) > len(out[j].GPUs) })
	return out
}

// clock is a duration the way Go writes one, 30m0s, with the units in this
// language's letters: 30dk0sn. Exact to the second, unlike Age, because
// how long a replica has sat unused is read against the idle threshold.
func (l *Lang) clock(d time.Duration) string {
	s := d.String()
	if l == english {
		return s
	}
	// Runs of digits are kept, with the decimal point this language
	// writes; each run of letters between them is a unit.
	var b strings.Builder
	unit := func(u string) {
		switch u {
		case "h", "m", "s", "ms":
			b.WriteString(l.T("dash.clock." + u))
		default:
			b.WriteString(u)
		}
	}
	start := 0
	for i, r := range s {
		if (r >= '0' && r <= '9') || r == '.' || r == '-' {
			if start < i {
				unit(s[start:i])
			}
			start = i + 1
			if r == '.' {
				b.WriteString(l.decimal)
			} else {
				b.WriteRune(r)
			}
		}
	}
	unit(s[start:])
	return b.String()
}

func init() {
	registerKeys("dash.clock.h", "dash.clock.m", "dash.clock.s", "dash.clock.ms")
}

// clockText is clock for a duration the control plane has already written
// out, such as a node's last-seen age. What is not a duration is left as it
// is.
func (l *Lang) clockText(s string) string {
	d, err := time.ParseDuration(s)
	if err != nil {
		return s
	}
	return l.clock(d)
}

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
}