llmgw/internal/modelplan/strata.go

package modelplan

import (
	"slices"

	"gitea.mynodes.xyz/ilker/llmgw/internal/strata"
)

// Strata runs a few mixture-of-experts models with their experts in the
// node's RAM and one card for the rest: a model that fits no card, or only
// many, served from one. Offered for the files of a family it runs, on nodes
// that have it installed, a card of at least strataMinCardGB and the RAM.

// strataMinCardGB is the smallest card Strata is worth running on: what its
// guide recommends.
const strataMinCardGB = 12

// StrataOptions are the ways a family's files could run on Strata here.
func StrataOptions(family string, arts []Artifact, nodes []PoolNode) []Option {
	var out []Option
	for _, a := range arts {
		size, ok := strata.SizeForFile(family, a.Name)
		if !ok {
			continue
		}
		need := float64(size.MemoryGB())
		fam, _ := strata.FamilyNamed(family)
		byClass := map[string]*Option{}
		var order []string
		for _, n := range nodes {
			if !slices.Contains(n.Engines, "strata") || n.HostMemGB <= 0 {
				continue
			}
			// A family whose engine needs a newer card: refused by Strata on
			// this one, after the node was chosen.
			if fam.MinCC > 0 && n.CC > 0 && n.CC < fam.MinCC {
				continue
			}
			have := float64(n.HostMemGB) * hostShare
			if need > have {
				continue
			}
			for _, k := range n.kinds() {
				if k.Class == "" || isMIG(k.Class) || k.GPUMemGB < strataMinCardGB || k.GPUs < 1 {
					continue
				}
				o := byClass[k.Class]
				if o == nil {
					o = &Option{Artifact: a.Name, Format: a.Format, Engine: "strata", Class: k.Class,
						GPUsPerNode: 1, Nodes: 1, NeedGB: need, StrataFamily: family, StrataSize: size.Name}
					byClass[k.Class] = o
					order = append(order, k.Class)
				}
				o.Hosts++
				o.HaveGB = max(o.HaveGB, have)
				if k.FreeGPUs > 0 && float64(n.HostFreeGB) >= need {
					o.FreeNow = true
				}
			}
		}
		for _, c := range order {
			out = append(out, *byClass[c])
		}
	}
	return out
}