llmgw/internal/control/catalog/hardware.go

package catalog

import (
	"fmt"
	"slices"
	"strconv"
	"strings"

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

// Profiles that fit the hardware.
//
// A profile names its GPU class — 6000-24g, a100-80g — and placement matches
// it exactly. A catalog entry meant for many sites cannot know their class
// names: the PoC has one Turing card, a customer an A100, another an L40S. So
// a profile may say what card it needs instead — a compute capability, a
// memory — and stands for every class in the cluster that meets it. What
// differs between cards goes on the profile too: its own weights (Source),
// its own workflow and nodes (ComfyUI), its own engine arguments. The model
// then has, say, an fp8 profile for compute 8.9 and up at a higher priority
// and an fp16 one for anything with 16 GB, and each site gets the best its
// cards can run.

// GPUNeed is what card a profile needs.
type GPUNeed struct {
	// MinComputeCapability is the oldest architecture that runs it: 7.5 is
	// Turing, 8.0 Ampere (bf16, flash-attention 2), 8.9 Ada (fp8), 9.0 Hopper.
	MinComputeCapability float64 `json:"min_compute_capability,omitempty" koanf:"min_compute_capability"`
	// MinMemoryGB is the card's own memory it needs.
	MinMemoryGB int `json:"min_memory_gb,omitempty" koanf:"min_memory_gb"`
}

// ProfileComfyUI is what a ComfyUI profile runs differently from the model:
// each field set replaces the model's.
type ProfileComfyUI struct {
	Workflow   string       `json:"workflow,omitempty" koanf:"workflow"`
	Checkpoint string       `json:"checkpoint,omitempty" koanf:"checkpoint"`
	Nodes      []CustomNode `json:"nodes,omitempty" koanf:"nodes"`
}

// GPUKind is one kind of card in the cluster.
type GPUKind struct {
	Class             string
	ComputeCapability float64
	MemoryGB          int
}

// GPUKinds are the kinds of card among gpus, one each, by class.
func GPUKinds(gpus []discovery.GPU) []GPUKind {
	seen := map[string]bool{}
	var out []GPUKind
	for _, g := range gpus {
		c := g.Class()
		if seen[c] {
			continue
		}
		seen[c] = true
		cc, _ := strconv.ParseFloat(strings.TrimSpace(g.ComputeCapability), 64)
		// A MIG instance's memory is the instance's: what a process sees.
		out = append(out, GPUKind{Class: c, ComputeCapability: cc, MemoryGB: (g.MemoryTotalMB + 512) / 1024})
	}
	slices.SortFunc(out, func(a, b GPUKind) int { return strings.Compare(a.Class, b.Class) })
	return out
}

// Fits reports whether a kind of card meets the need.
func (n GPUNeed) Fits(k GPUKind) bool {
	// A card that does not say its capability is not taken for one that has
	// it: an fp8 workflow on a card of unknown age fails at its first step.
	if n.MinComputeCapability > 0 && k.ComputeCapability+1e-9 < n.MinComputeCapability {
		return false
	}
	return n.MinMemoryGB <= 0 || k.MemoryGB >= n.MinMemoryGB
}

func (n GPUNeed) String() string {
	var parts []string
	if n.MinComputeCapability > 0 {
		parts = append(parts, fmt.Sprintf("compute %.1f+", n.MinComputeCapability))
	}
	if n.MinMemoryGB > 0 {
		parts = append(parts, fmt.Sprintf("%d GB+", n.MinMemoryGB))
	}
	if len(parts) == 0 {
		return "any GPU"
	}
	return strings.Join(parts, ", ")
}

// forClass is a profile that names its GPU by need, as the one it stands for
// on a class.
func (p Profile) forClass(class string) Profile {
	p.Name = p.Name + "@" + class
	p.GPUClass = class
	p.GPU = nil
	return p
}

func (p Profile) validateHardware(m Model) error {
	if strings.Contains(p.Name, "@") {
		return fmt.Errorf("model %q profile %q: a profile's name has no '@': it is how one standing for a class is named",
			m.Name, p.Name)
	}
	if p.GPU != nil {
		if p.GPUClass != "" {
			return fmt.Errorf("model %q profile %q: gpu_class and gpu are one or the other: a class, or what card it needs",
				m.Name, p.Name)
		}
		if p.GPU.MinComputeCapability < 0 || p.GPU.MinComputeCapability > 20 || p.GPU.MinMemoryGB < 0 {
			return fmt.Errorf("model %q profile %q: gpu %+v is not a compute capability and a memory", m.Name, p.Name, *p.GPU)
		}
	}
	if p.ComfyUI != nil {
		if p.Engine != EngineComfyUI {
			return fmt.Errorf("model %q profile %q: comfyui: is for a profile of engine comfyui", m.Name, p.Name)
		}
		if err := checkWorkflow(m.Name, p.ComfyUI.Workflow); err != nil {
			return err
		}
		for _, n := range p.ComfyUI.Nodes {
			if err := n.validate(m.Name); err != nil {
				return err
			}
		}
	}
	return nil
}

// PlacementProfilesFor is a model's profiles in placement's view, each one
// that names its GPU by need standing for every kind of card among kinds
// that meets it, at its own priority. One that no card meets stands for
// none, and is not tried.
func (m Model) PlacementProfilesFor(kinds []GPUKind) []placement.Profile {
	var expanded []Profile
	for _, p := range m.Profiles {
		if p.GPU == nil {
			expanded = append(expanded, p)
			continue
		}
		for _, k := range kinds {
			if p.GPU.Fits(k) {
				expanded = append(expanded, p.forClass(k.Class))
			}
		}
	}
	mm := m
	mm.Profiles = expanded
	return mm.PlacementProfiles()
}

// ComfyWorkflowFor is what a profile's replicas are sent: the profile's own
// workflow and checkpoint where it has them, else the model's.
func (m Model) ComfyWorkflowFor(p Profile) (workflow, checkpoint string) {
	workflow, checkpoint = m.ComfyWorkflow()
	if p.ComfyUI != nil {
		if p.ComfyUI.Workflow != "" {
			workflow = p.ComfyUI.Workflow
		}
		if p.ComfyUI.Checkpoint != "" {
			checkpoint = p.ComfyUI.Checkpoint
		}
	}
	if p.Source != "" && (p.ComfyUI == nil || p.ComfyUI.Checkpoint == "") && !strings.HasSuffix(p.Source, "/") {
		// A profile with weights of its own loads them by their own name.
		checkpoint = lastPathPart(strings.TrimPrefix(p.Source, "hf:"))
	}
	return workflow, checkpoint
}

// ComfyNodesFor are the custom nodes a profile's replicas carry.
func (m Model) ComfyNodesFor(p Profile) []CustomNode {
	if p.ComfyUI != nil && p.ComfyUI.Nodes != nil {
		return p.ComfyUI.Nodes
	}
	if m.ComfyUI != nil {
		return m.ComfyUI.Nodes
	}
	return nil
}

func lastPathPart(s string) string {
	if i := strings.LastIndex(s, "/"); i >= 0 {
		return s[i+1:]
	}
	return s
}