llmgw/internal/web/explore_fit.go

package web

import (
	"slices"
	"strconv"
	"strings"

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

// The explorer lists what this cluster can run, judged by the nodes it has
// now rather than by a size picked from a list: a site with H100s sees the
// 70B models, one with a T4 and CPUs sees what fits there. From the
// registry's listing alone — a parameter count, safetensors by dtype, tags
// — so no repository is opened to decide; the repository's page sizes each
// file exactly.

// clusterFit is the hardware a listing is judged against.
type clusterFit struct {
	pool []modelplan.PoolNode
	// fp8 are the nodes with a card vLLM runs FP8 weights on, compute
	// capability 8.9 or newer, as the agent requires.
	fp8 []modelplan.PoolNode
	// card is the largest card's memory in GB, cc its compute capability;
	// host the most memory one node has.
	card int
	cc   string
	host int
	// cpus are the nodes whose memory's speed their agent measured, for
	// how fast a model would answer on one.
	cpus []cpuHost
	// classBW is each card class's memory bandwidth in GB/s, from its
	// cards' names; nvlink how many cards of a class one node has joined
	// by NVLink, at most; fabrics what each fabric is, by name.
	classBW map[string]float64
	nvlink  map[string]int
	fabrics map[string]modelplan.Links
}

// cpuHost is a node as a model on its CPU sees it.
type cpuHost struct {
	id    string
	ramGB int
	bwGBs float64
}

// fp8Capability is the oldest card vLLM runs an FP8 checkpoint on: Ampere,
// by Marlin's kernels, which read FP8 weights on cards without FP8 maths.
// Ada and Hopper compute in it; the agent's 8.9 is for --quantization fp8,
// which makes FP8 of other weights.
const fp8Capability = 8.0

// fitOf is the cluster's hardware now, nil when there is none to judge by:
// no fleet to ask, or no node ready.
func (p *AccountPages) fitOf() *clusterFit {
	if p.fleet == nil {
		return nil
	}
	nodes := p.fleet.Nodes()
	c := &clusterFit{pool: adminapi.PoolNodesOf(nodes)}
	if len(c.pool) == 0 {
		return nil
	}
	c.classBW, c.nvlink, c.fabrics = map[string]float64{}, map[string]int{}, map[string]modelplan.Links{}
	newest := map[string]float64{}
	for _, n := range nodes {
		if n.State != "ready" {
			continue
		}
		byClass := map[string][]adminapi.GPUView{}
		for _, g := range n.GPUs {
			bw := modelplan.GPUBandwidthGBs(g.Name)
			// A MIG instance has its share of the card's memory, and of the
			// bandwidth to it.
			if parent, _, ok := strings.Cut(g.Class, "-mig-"); ok {
				if whole := modelplan.MemFromClass(parent); whole > 0 {
					bw *= float64(modelplan.MemFromClass(g.Class)) / float64(whole)
				}
			}
			if bw > c.classBW[g.Class] {
				c.classBW[g.Class] = bw
			}
			byClass[g.Class] = append(byClass[g.Class], g)
		}
		for class, gs := range byClass {
			c.nvlink[class] = max(c.nvlink[class], nvlinkClique(gs))
		}
		for _, p := range n.IBPorts {
			if p.Fabric != "" && p.Active {
				l := c.fabrics[p.Fabric]
				l.InfiniBand = p.LinkLayer == "InfiniBand"
				if l.RateGbps == 0 || p.RateGbps < l.RateGbps {
					l.RateGbps = p.RateGbps
				}
				c.fabrics[p.Fabric] = l
			}
		}
		if n.State == "ready" && n.MemBandwidthGBs > 0 {
			c.cpus = append(c.cpus, cpuHost{id: n.ID, ramGB: n.MemoryGB, bwGBs: n.MemBandwidthGBs})
		}
		for _, s := range n.ComputeCapabilities {
			if v, err := strconv.ParseFloat(s, 64); err == nil && v > newest[n.ID] {
				newest[n.ID] = v
			}
		}
	}
	for _, n := range c.pool {
		c.host = max(c.host, n.HostMemGB)
		if n.GPUs > 0 && n.GPUMemGB >= c.card {
			if n.GPUMemGB > c.card || newest[n.ID] > atof(c.cc) {
				c.cc = strconv.FormatFloat(newest[n.ID], 'f', 1, 64)
			}
			c.card = n.GPUMemGB
		}
		if newest[n.ID] >= fp8Capability {
			c.fp8 = append(c.fp8, n)
		}
	}
	return c
}

func atof(s string) float64 { v, _ := strconv.ParseFloat(s, 64); return v }

// gpus reports whether any node has a card.
func (c *clusterFit) gpus() bool { return c.card > 0 }

// pictureTasks are a pipeline ComfyUI runs, on a card, rather than an
// engine for text.
var pictureTasks = []string{"text-to-image", "image-to-image", "image-to-3d", "text-to-3d",
	"text-to-video", "image-to-video", "unconditional-image-generation"}

// ggufBitsPerWeight is the quantisation a GGUF repository is judged by:
// Q4_K_M, the smallest of those published that answers as well as most
// people want. A repository of Q8 only is judged kindly; its page says.
const ggufBitsPerWeight = 4.8

// runs reports whether some node here could hold the repository's model.
func (c *clusterFit) runs(m hf.Model) bool {
	params := m.Params()
	if slices.Contains(pictureTasks, m.PipelineTag) || m.HasTag("diffusers") {
		b := m.WeightBytes()
		if b == 0 {
			b = params * 2
		}
		// ComfyUI holds the model on one card, with room for its working.
		return c.gpus() && float64(b+2<<30) <= float64(c.card)*0.9*(1<<30)
	}
	if m.GGUF != nil || m.HasTag("gguf") && m.Safetensors == nil {
		if params == 0 {
			return true
		}
		b := int64(float64(params) * ggufBitsPerWeight / 8)
		if anyFits(modelplan.Options([]modelplan.Artifact{{Name: "m-Q4_K_M.gguf", Format: "gguf", Bytes: b,
			Quant: "Q4_K_M"}}, modelplan.Shape{Params: params}, c.pool, 0)) {
			return true
		}
	}
	// MLX safetensors are Apple's, which no engine here reads; a repository
	// tagged so may still carry a GGUF, judged above.
	if (m.Safetensors != nil || m.HasTag("safetensors")) && !m.HasTag("mlx") {
		// vLLM, which runs on cards only.
		pool := c.pool
		if m.FP8() {
			pool = c.fp8
		}
		if !c.gpus() || len(pool) == 0 {
			return false
		}
		b := m.WeightBytes()
		if b == 0 {
			if params == 0 {
				return true
			}
			b = params * 2
		}
		return anyFits(modelplan.Options([]modelplan.Artifact{{Name: "m.safetensors", Format: "safetensors",
			Bytes: b}}, modelplan.Shape{Params: params}, pool, 0))
	}
	return false
}

func anyFits(opts []modelplan.Option) bool {
	return slices.ContainsFunc(opts, modelplan.Option.Fits)
}

// note is what the listing was judged by, for the page to say.
func (c *clusterFit) note(l *Lang) string {
	if !c.gpus() {
		return l.T("explore.fit_cpu", c.host)
	}
	return l.T("explore.fit", c.card, c.cc, c.host)
}

// holds judges one file set of an open repository, whose size is known: a
// verdict that it cannot be served here, and false, when no node holds it.
func (c *clusterFit) holds(l *Lang, name string, bytes, params int64) (gguf.Verdict, bool) {
	format := artifactFormat(name)
	if format == "" || bytes <= 0 {
		return gguf.Verdict{}, true
	}
	opts := modelplan.Options([]modelplan.Artifact{{Name: name, Format: format, Bytes: bytes,
		Quant: modelplan.QuantOf(name)}}, modelplan.Shape{Params: params}, c.pool, 0)
	if anyFits(opts) {
		return gguf.Verdict{}, true
	}
	if format == "safetensors" && !c.gpus() {
		return gguf.Verdict{Checked: true, Reason: l.T("card.no_gpu")}, false
	}
	if len(opts) == 0 {
		// No width of any node's cards holds it, so none is offered: what
		// the weights take against all the cards one node has.
		have := 0
		for _, n := range c.pool {
			have = max(have, n.GPUs*n.GPUMemGB)
		}
		return gguf.Verdict{Checked: true, Reason: l.T("card.too_big",
			float64(bytes)/(1<<30), float64(have)*0.9)}, false
	}
	best := opts[0]
	for _, o := range opts[1:] {
		if o.HaveGB-o.NeedGB > best.HaveGB-best.NeedGB {
			best = o
		}
	}
	return gguf.Verdict{Checked: true, Reason: l.T("card.too_big", best.NeedGB, best.HaveGB)}, false
}

// cpuSpeed is about how many tokens a second one reply gets from a GGUF of
// bytes on the fastest node whose RAM holds it, and that node; share is the
// part a token reads (ActiveShare). 0 when no measured node holds it.
func (c *clusterFit) cpuSpeed(bytes int64, share float64) (float64, string) {
	var best cpuHost
	for _, n := range c.cpus {
		need := float64(bytes) + float64(1<<30) + float64(bytes)/20
		if need <= float64(n.ramGB)*0.85*(1<<30) && n.bwGBs > best.bwGBs {
			best = n
		}
	}
	if best.id == "" {
		return 0, ""
	}
	return modelplan.CPUTokensPerSecond(int64(float64(bytes)*share), best.bwGBs), best.id
}

// speedLabel is a rate as the page shows it: whole tokens from ten up, one
// decimal below, and under one said so.
func speedLabel(tps float64) string {
	switch {
	case tps <= 0:
		return ""
	case tps < 1:
		return "<1"
	case tps < 10:
		return "~" + strconv.FormatFloat(tps, 'f', 1, 64)
	}
	return "~" + strconv.FormatFloat(tps, 'f', 0, 64)
}

// rowSpeed is a listed repository's speed on the CPU, judged at Q4_K_M:
// GGUF only, which is what runs there.
func (c *clusterFit) rowSpeed(m hf.Model) (string, string) {
	if m.GGUF == nil && !m.HasTag("gguf") || m.Params() <= 0 {
		return "", ""
	}
	b := int64(float64(m.Params()) * ggufBitsPerWeight / 8)
	tps, node := c.cpuSpeed(b, modelplan.ActiveShare(m.ID, m.Params()))
	return speedLabel(tps), node
}

// gpuSpeed is about how many tokens a second one reply gets from art on
// the cards, laid out as the planner would place it — the fewest cards
// that hold it, for each class — the fastest of those, and how it is laid
// out. 0 when nothing here holds it or no card's speed is known.
func (c *clusterFit) gpuSpeed(l *Lang, art modelplan.Artifact, params int64, share float64, fp8 bool) (float64, string) {
	pool := c.pool
	if fp8 {
		pool = c.fp8
	}
	best, layout := 0.0, ""
	for _, o := range modelplan.Options([]modelplan.Artifact{art}, modelplan.Shape{Params: params}, pool, 0) {
		if o.CPU || !o.Fits() {
			continue
		}
		links := modelplan.Links{NVLink: c.nvlink[o.Class] >= o.GPUsPerNode && o.GPUsPerNode > 1}
		if o.Nodes > 1 {
			f := c.fabrics[o.Fabric]
			links.InfiniBand, links.RateGbps = f.InfiniBand, f.RateGbps
		}
		tps := modelplan.GPUTokensPerSecond(o, art.Bytes, share, params, c.classBW[o.Class], links)
		if tps > best {
			best, layout = tps, layoutLabel(l, o, links)
		}
	}
	return best, layout
}

// layoutLabel says how a replica would be laid out: "1× h100-80g",
// "4× h100-80g, NVLink, vLLM", "2 nodes × 8× h100-80g over ib0 (InfiniBand
// 400 Gb/s), vLLM".
func layoutLabel(l *Lang, o modelplan.Option, links modelplan.Links) string {
	s := strconv.Itoa(max(o.GPUsPerNode, 1)) + "× " + o.Class
	if o.GPUsPerNode > 1 {
		if links.NVLink {
			s += ", NVLink"
		} else {
			s += ", PCIe"
		}
	}
	if o.Nodes > 1 {
		kind := "RoCE"
		if links.InfiniBand {
			kind = "InfiniBand"
		}
		if links.RateGbps > 0 {
			kind += " " + strconv.Itoa(links.RateGbps) + " Gb/s"
		}
		s = l.T("explore.layout.nodes", o.Nodes, s, o.Fabric, kind)
	}
	engine := "llama.cpp"
	if o.Engine == "vllm" {
		engine = "vLLM"
	}
	return s + ", " + engine
}

// rowGPU is a listed repository's speed on the cards: its GGUF at Q4_K_M
// on llama.cpp, its safetensors on vLLM, whichever is faster here.
func (c *clusterFit) rowGPU(l *Lang, m hf.Model) (string, string) {
	if !c.gpus() || m.Params() <= 0 {
		return "", ""
	}
	params, share := m.Params(), modelplan.ActiveShare(m.ID, m.Params())
	best, layout := 0.0, ""
	if m.GGUF != nil || m.HasTag("gguf") {
		b := int64(float64(params) * ggufBitsPerWeight / 8)
		best, layout = c.gpuSpeed(l, modelplan.Artifact{Name: "m-Q4_K_M.gguf", Format: "gguf", Bytes: b, Quant: "Q4_K_M"},
			params, share, false)
	}
	if (m.Safetensors != nil || m.HasTag("safetensors")) && !m.HasTag("mlx") {
		b := m.WeightBytes()
		if b == 0 {
			b = params * 2
		}
		if tps, lay := c.gpuSpeed(l, modelplan.Artifact{Name: "m.safetensors", Format: "safetensors", Bytes: b},
			params, share, m.FP8()); tps > best {
			best, layout = tps, lay
		}
	}
	return speedLabel(best), layout
}

// artifactFormat is "gguf" or "safetensors" by a file set's name, or "".
func artifactFormat(name string) string {
	switch lower := strings.ToLower(name); {
	case strings.HasSuffix(lower, ".gguf"):
		return "gguf"
	case strings.HasSuffix(lower, ".safetensors"):
		return "safetensors"
	}
	return ""
}

// nvlinkClique is the most cards of these that are each joined by NVLink to
// every other: tensor parallel across four needs all four on one fabric, not
// two bridged pairs.
func nvlinkClique(gs []adminapi.GPUView) int {
	if len(gs) > 16 {
		gs = gs[:16]
	}
	peer := map[int]map[int]bool{}
	for _, g := range gs {
		peer[g.Index] = map[int]bool{}
		for _, p := range g.NVLinkPeers {
			peer[g.Index][p] = true
		}
	}
	best := 0
	for set := 1; set < 1<<len(gs); set++ {
		var in []int
		for i := range gs {
			if set&(1<<i) != 0 {
				in = append(in, gs[i].Index)
			}
		}
		if len(in) <= best {
			continue
		}
		ok := true
		for _, a := range in {
			for _, b := range in {
				if a != b && !peer[a][b] {
					ok = false
				}
			}
		}
		if ok && len(in) > 1 {
			best = len(in)
		}
	}
	return best
}