llmgw/internal/control/catalog/catalog.go

// Package catalog holds the desired state: which models exist, how each one can
// be run, and how many copies of it there should be (plan §6.1).
package catalog

import (
	"fmt"
	"slices"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/placement"
)

// Scaling is how many replicas a model should have.
type Scaling struct {
	Min int `koanf:"min"`
	Max int `koanf:"max"`
}

// SLO is the latency target a model is expected to meet. `ctl bench` measures
// capacity against it (plan §6.2.1).
type SLO struct {
	TTFTp99 time.Duration `koanf:"ttft_p99"`
	TPOTp99 time.Duration `koanf:"tpot_p99"`
}

// Profile is how a model runs on one GPU class.
type Profile struct {
	Name     string `koanf:"name"`
	Engine   string `koanf:"engine"`
	GPUClass string `koanf:"gpu_class"`

	TP int `koanf:"tp"`
	PP int `koanf:"pp"`
	// GPUsPerNode defaults to TP for a single-node profile.
	GPUsPerNode int `koanf:"gpus_per_node"`
	// Nodes above 1 makes this a multi-node profile, which requires a fabric.
	Nodes int `koanf:"nodes"`

	Args []string          `koanf:"args"`
	Env  map[string]string `koanf:"env"`

	// Weight is the router's capacity weight. It should come from ctl bench
	// rather than a guess (plan §6.2.1).
	Weight float64 `koanf:"weight"`
	// Priority orders profiles during placement; highest that fits wins.
	Priority int `koanf:"priority"`
	// StartupTimeout bounds model loading. A 70B model across two nodes takes
	// tens of minutes, so this is per profile rather than global.
	StartupTimeout time.Duration `koanf:"startup_timeout"`
}

// CPUOnly reports whether this profile needs no GPUs.
//
// Naming no GPU class is how a profile says so, and only engines that can
// actually serve without a card are allowed to leave it out — otherwise an
// omitted class would quietly become a replica placed on a machine with no
// hardware to run it.
func (p Profile) CPUOnly() bool { return p.GPUClass == "" && cpuCapable(p.Engine) }

// cpuCapable reports whether an engine can serve a model without a GPU.
//
// llama.cpp is the one that can, and it is what makes a cluster with no cards
// testable end to end.
func cpuCapable(engine string) bool { return engine == "llamacpp" }

// Model is a catalog entry as clients see it.
type Model struct {
	Name    string   `koanf:"name"`
	Aliases []string `koanf:"aliases"`
	// Source is the weights path or a HuggingFace id.
	Source    string    `koanf:"source"`
	Profiles  []Profile `koanf:"profiles"`
	Scaling   Scaling   `koanf:"scaling"`
	Fallbacks []string  `koanf:"fallbacks"`
	// RoutingProfile selects how the router scores endpoints for this model.
	RoutingProfile string `koanf:"routing_profile"`
	SLO            SLO    `koanf:"slo"`
}

// Catalog is the set of models the cluster should serve.
type Catalog struct {
	models  map[string]Model
	aliases map[string]string
	order   []string
}

// New builds a catalog, rejecting definitions that would fail confusingly later.
func New(models []Model) (*Catalog, error) {
	c := &Catalog{
		models:  make(map[string]Model, len(models)),
		aliases: map[string]string{},
	}

	for i := range models {
		m := models[i].normalize()
		if err := m.validate(); err != nil {
			return nil, err
		}
		if _, dup := c.models[m.Name]; dup {
			return nil, fmt.Errorf("model %q is declared twice", m.Name)
		}
		c.models[m.Name] = m
		c.order = append(c.order, m.Name)
	}

	// Aliases share a namespace with names. A collision would silently shadow a
	// model at request time, so it has to fail at load.
	for _, name := range c.order {
		for _, a := range c.models[name].Aliases {
			if _, clash := c.models[a]; clash {
				return nil, fmt.Errorf("model %q: alias %q collides with another model name", name, a)
			}
			if prev, clash := c.aliases[a]; clash {
				return nil, fmt.Errorf("alias %q is claimed by both %q and %q", a, prev, name)
			}
			c.aliases[a] = name
		}
	}

	// Fallbacks are resolved now rather than at request time: a typo in a
	// fallback should stop a deployment, not surface during an outage when the
	// fallback is finally needed.
	for _, name := range c.order {
		for _, f := range c.models[name].Fallbacks {
			if _, ok := c.Resolve(f); !ok {
				return nil, fmt.Errorf("model %q: fallback %q is not in the catalog", name, f)
			}
		}
	}
	return c, nil
}

func (m Model) normalize() Model {
	if m.Scaling.Min < 0 {
		m.Scaling.Min = 0
	}
	if m.Scaling.Max == 0 {
		m.Scaling.Max = max(m.Scaling.Min, 1)
	}
	for i := range m.Profiles {
		p := &m.Profiles[i]
		if p.Engine == "" {
			p.Engine = "vllm"
		}
		if p.TP <= 0 {
			p.TP = 1
		}
		if p.Nodes <= 0 {
			p.Nodes = 1
		}
		switch {
		case p.CPUOnly():
			// Nothing to divide up: this profile asks for no cards at all.
			p.GPUsPerNode = 0
		case p.GPUsPerNode <= 0:
			// The common case: a single-node profile uses TP GPUs on one node.
			p.GPUsPerNode = p.TP / p.Nodes
			if p.GPUsPerNode <= 0 {
				p.GPUsPerNode = 1
			}
		}
		if p.StartupTimeout <= 0 {
			p.StartupTimeout = 20 * time.Minute
		}
		if p.Weight <= 0 {
			p.Weight = 1
		}
	}
	return m
}

func (m Model) validate() error {
	if m.Name == "" {
		return fmt.Errorf("a model has no name")
	}
	if m.Source == "" {
		return fmt.Errorf("model %q: source is required", m.Name)
	}
	if len(m.Profiles) == 0 {
		return fmt.Errorf("model %q: at least one profile is required", m.Name)
	}
	if m.Scaling.Max < m.Scaling.Min {
		return fmt.Errorf("model %q: scaling max %d is below min %d", m.Name, m.Scaling.Max, m.Scaling.Min)
	}

	seen := map[string]struct{}{}
	for _, p := range m.Profiles {
		if p.Name == "" {
			return fmt.Errorf("model %q: a profile has no name", m.Name)
		}
		if _, dup := seen[p.Name]; dup {
			return fmt.Errorf("model %q: profile %q is declared twice", m.Name, p.Name)
		}
		seen[p.Name] = struct{}{}

		if p.GPUClass == "" && !cpuCapable(p.Engine) {
			return fmt.Errorf("model %q profile %q: gpu_class is required for engine %q",
				m.Name, p.Name, p.Engine)
		}
		if p.CPUOnly() {
			if p.Nodes > 1 {
				return fmt.Errorf("model %q profile %q: a CPU profile cannot span %d nodes",
					m.Name, p.Name, p.Nodes)
			}
			continue
		}
		// The arithmetic has to add up, or the engine is handed a shard spec it
		// cannot satisfy and fails minutes later during model loading.
		if want := p.GPUsPerNode * p.Nodes; want != p.TP*max(p.PP, 1) {
			return fmt.Errorf("model %q profile %q: %d nodes x %d GPUs does not match TP %d x PP %d",
				m.Name, p.Name, p.Nodes, p.GPUsPerNode, p.TP, max(p.PP, 1))
		}
	}
	return nil
}

// Resolve maps a client-supplied name or alias onto a model.
func (c *Catalog) Resolve(name string) (Model, bool) {
	if m, ok := c.models[name]; ok {
		return m, true
	}
	if canonical, ok := c.aliases[name]; ok {
		return c.models[canonical], true
	}
	return Model{}, false
}

// Models returns every model, in declaration order.
func (c *Catalog) Models() []Model {
	out := make([]Model, 0, len(c.order))
	for _, n := range c.order {
		out = append(out, c.models[n])
	}
	return out
}

// Names returns the canonical model names, sorted.
func (c *Catalog) Names() []string {
	out := slices.Clone(c.order)
	slices.Sort(out)
	return out
}

// AliasesOf returns the alternative names for a model.
func (c *Catalog) AliasesOf(name string) []string {
	m, ok := c.Resolve(name)
	if !ok {
		return nil
	}
	return slices.Clone(m.Aliases)
}

// PlacementProfiles converts a model's profiles into placement's view.
func (m Model) PlacementProfiles() []placement.Profile {
	out := make([]placement.Profile, 0, len(m.Profiles))
	for _, p := range m.Profiles {
		out = append(out, placement.Profile{
			Name:        p.Name,
			Engine:      p.Engine,
			GPUClass:    p.GPUClass,
			GPUsPerNode: p.GPUsPerNode,
			Nodes:       p.Nodes,
			Priority:    p.Priority,
		})
	}
	return out
}

// EngineProfile converts a catalog profile into the engine adapter's view.
func (m Model) EngineProfile(name string) (engines.Profile, bool) {
	for _, p := range m.Profiles {
		if p.Name != name {
			continue
		}
		return engines.Profile{
			Name:            p.Name,
			Engine:          p.Engine,
			GPUClass:        p.GPUClass,
			TP:              p.TP,
			PP:              p.PP,
			GPUsPerNode:     p.GPUsPerNode,
			Nodes:           p.Nodes,
			Args:            slices.Clone(p.Args),
			Env:             p.Env,
			ModelPath:       m.Source,
			ServedModelName: m.Name,
		}, true
	}
	return engines.Profile{}, false
}

// ProfileByName returns one of a model's profiles.
func (m Model) ProfileByName(name string) (Profile, bool) {
	for _, p := range m.Profiles {
		if p.Name == name {
			return p, true
		}
	}
	return Profile{}, false
}

// String renders the catalog for `ctl models`.
func (c *Catalog) String() string {
	var b strings.Builder
	for _, m := range c.Models() {
		fmt.Fprintf(&b, "%s (min %d, max %d)\n", m.Name, m.Scaling.Min, m.Scaling.Max)
		for _, p := range m.Profiles {
			fmt.Fprintf(&b, "  %-20s %-10s %-12s tp=%d nodes=%d priority=%d\n",
				p.Name, p.Engine, p.GPUClass, p.TP, p.Nodes, p.Priority)
		}
	}
	return b.String()
}