llmgw/internal/agent/provision/recipe.go

// Package provision installs the engines a node is asked to have.
//
// The control plane says which engines at which versions; everything else is
// worked out here, on the node, because only the node knows it: which build
// fits its processor and driver, what its operating system is missing, and
// whether the result actually runs. Nothing already on the machine is taken
// on trust. An engine counts as installed when this package installed it,
// checked every file against the checksum its publisher gave, and watched it
// start.
package provision

import (
	"fmt"
	"slices"
	"strings"

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

// Build is one published build of an engine, chosen for one machine.
type Build struct {
	// Variant names the build: "cpu", "cuda-12.8".
	Variant string
	// Assets are the release files to fetch, unpacked in order into one
	// directory.
	Assets []string
	// Packages are operating system packages the build links against and a
	// minimal install may not carry, by packaging family ("deb", "rpm"): the
	// same library is libgomp1 on one and libgomp on the other.
	Packages map[string][]string
	// Binary is the executable's file name, found anywhere in the unpacked
	// tree.
	Binary string
	// CPUFlags are extensions the build is slow without. Missing ones are
	// reported, not refused: slow is still running.
	CPUFlags []string

	// Python, when set, makes the build a Python environment rather than an
	// archive: that interpreter version, with Requirements installed from
	// PyPI by uv. Binary is then looked for in the environment's bin.
	Python       string
	Requirements []string
	// VersionArgs is what Binary is run with to prove the build works, when
	// --version is not it.
	VersionArgs []string
	// GPUBackend, when set, is the shared library the build computes on GPUs
	// with, loaded at run time rather than linked: "libggml-cuda.so". The
	// build must list a device with it, or it is not installed — a llama.cpp
	// whose CUDA backend cannot load runs on the CPU without complaint.
	GPUBackend string
	// Wheel, when set, is the engine itself as a wheel from a GitHub
	// release rather than from PyPI, checked against the digest GitHub
	// publishes for it. ExtraIndex is a package index besides PyPI, for
	// IndexPackages only: the CUDA 12.9 torch the wheel was built on. Every
	// other package comes from PyPI alone.
	Wheel         *Wheel
	ExtraIndex    string
	IndexPackages []string
	// GPUCheck, on a node with an NVIDIA GPU, has the installed environment
	// compute something on one before the build counts as installed. A build
	// with no kernels for the card — a GPU generation its torch dropped —
	// imports and prints its version, and fails on its first model.
	GPUCheck bool
	// Version is the engine version this build is, when it is not the one
	// asked for: the newest built for this machine's GPUs, where the one
	// asked for has no kernels for them. Empty is the version asked for.
	Version string
	// Note says why this build, when it is not the plain one: shown beside
	// it wherever the installation is.
	Note string
	// FreeGB is the disk space the installation needs before it starts: a
	// Python engine is several gigabytes, and one that fills the disk halfway
	// leaves the node worse off than one that never began.
	FreeGB int
}

// Wheel is a package published as a GitHub release asset: the asset whose
// name starts with Prefix and ends with Suffix.
type Wheel struct {
	Project, Repo, Tag, Prefix, Suffix string
}

// Recipe knows how to install one engine.
type Recipe struct {
	// Repo is the GitHub repository whose releases carry the builds; empty
	// for an engine installed from PyPI.
	Repo string
	// PyPI is the project an engine installed from PyPI is published as.
	PyPI string
	// Candidates are the builds that may run on a machine, best first, or
	// why none can. The installer takes the first that installs and is seen
	// to work — computing on the GPU, where there is one — so a build this
	// table is wrong about costs a download, not a node. variant is a build
	// chosen for the node; empty leaves the choice to the machine.
	Candidates func(version, variant string, h discovery.Host) ([]Build, error)
}

// Recipes are the engines this agent can install.
//
// vLLM and SGLang are Python environments of several gigabytes, installed
// from PyPI by uv, which is itself installed like llama.cpp: a GitHub release
// checked against its published checksum.
var Recipes = map[string]Recipe{
	"llamacpp": {Repo: "ggml-org/llama.cpp", Candidates: llamaCppCandidates},
	"vllm":     {PyPI: "vllm", Candidates: vllmCandidates},
	"sglang":   {PyPI: "sglang", Candidates: sglangCandidates},
	"uv":       {Repo: "astral-sh/uv", Candidates: uvCandidates},
}

// Installable are the engines this agent can install, as it tells the
// control plane: every recipe but uv, which is a tool for the others.
func Installable() []string {
	var out []string
	for name := range Recipes {
		if name != "uv" {
			out = append(out, name)
		}
	}
	slices.Sort(out)
	return out
}

// Builds are every build of an engine version published for an
// architecture: what an offline bundle carries, for nodes it has not seen.
// Read from the same table the nodes choose from, asked as a machine with
// no GPU and as one with any NVIDIA GPU, so a bundle and the nodes that
// install from it cannot disagree about which files a build is. Archive
// engines only: a Python engine is a PyPI environment, which no bundle
// carries.
func Builds(engine, version, arch string) ([]Build, error) {
	r, ok := Recipes[engine]
	switch {
	case !ok:
		return nil, fmt.Errorf("this agent does not know how to install %s", engine)
	case r.Repo == "":
		return nil, fmt.Errorf("%s is installed from PyPI, which an offline bundle does not carry", engine)
	case !validVersion.MatchString(version):
		return nil, fmt.Errorf("%q is not a version of %s", version, engine)
	}
	var out []Build
	seen := map[string]bool{}
	var errs []string
	for _, h := range []discovery.Host{
		{Arch: arch},
		// Any card and the newest driver: every GPU build is a candidate.
		{Arch: arch, NVIDIADriver: "9999"},
	} {
		bs, err := r.Candidates(version, "", h)
		if err != nil {
			errs = append(errs, err.Error())
			continue
		}
		for _, b := range bs {
			if !seen[b.Variant] {
				seen[b.Variant] = true
				out = append(out, b)
			}
		}
	}
	if len(out) == 0 {
		return nil, fmt.Errorf("no build of %s %s for %s: %s", engine, version, arch, strings.Join(errs, "; "))
	}
	return out, nil
}

// UVVersion is the uv that installs the Python engines.
const UVVersion = "0.12.20"

// pythonForEngines is the interpreter the Python engines run on. uv fetches
// it, so the node's own Python — or its absence — does not matter.
const pythonForEngines = "3.12"

// NotHere is the error a recipe gives for a node the engine is not meant
// for, as opposed to one it should run on and cannot: vLLM on a machine
// with no GPU is not a fault, and is not reported as one.
type NotHere struct{ Reason string }

func (e NotHere) Error() string { return e.Reason }

func uvCandidates(version, _ string, h discovery.Host) ([]Build, error) {
	b, err := chooseUV(version, h)
	if err != nil {
		return nil, err
	}
	return []Build{b}, nil
}

func chooseUV(_ string, h discovery.Host) (Build, error) {
	var arch string
	switch h.Arch {
	case "amd64":
		arch = "x86_64"
	case "arm64":
		arch = "aarch64"
	default:
		return Build{}, fmt.Errorf("uv publishes no Linux build for %s", orUnknown(h.Arch))
	}
	return Build{
		Variant: arch,
		Assets:  []string{"uv-" + arch + "-unknown-linux-gnu.tar.gz"},
		Binary:  "uv",
	}, nil
}

// What builds need of the NVIDIA side of a machine, read from what each
// project publishes and how it was built:
//
//   - CUDA 13 needs driver 580; CUDA 12 builds run on driver 525 and newer,
//     by CUDA's minor version compatibility — except code shipped only as
//     PTX, which the driver compiles and so needs a driver as new as the
//     CUDA it was compiled with.
//   - Which GPUs a build has kernels for is the build's own list, not what
//     its CUDA supports.
const (
	minDriverForCUDA13 = 580
	minDriverForCUDA12 = 525
)

// gpuNeeds is what one build needs of a machine's GPUs and driver.
type gpuNeeds struct {
	driver int
	cc     float64
}

// fits says why a machine cannot run a build, or nil.
func (n gpuNeeds) fits(what string, h discovery.Host) error {
	if major := leadingInt(h.NVIDIADriver); n.driver > 0 && major > 0 && major < n.driver {
		return fmt.Errorf("%s needs NVIDIA driver %d or newer; this node has %s", what, n.driver, h.NVIDIADriver)
	}
	if cc := h.MinComputeCapability; n.cc > 0 && cc > 0 && cc < n.cc {
		return fmt.Errorf("%s has kernels for compute capability %.1f and newer; this node has %.1f", what, n.cc, cc)
	}
	return nil
}

// candidate is a build and what it needs.
type candidate struct {
	b     Build
	needs gpuNeeds
	what  string
}

// pick keeps the candidates a machine can run, in order, and those a
// variant chosen for the node names; with none left, it says why each went.
func pick(engine, variant string, cands []candidate, h discovery.Host) ([]Build, error) {
	var out []Build
	var why []string
	for _, c := range cands {
		if variant != "" && c.b.Variant != variant {
			continue
		}
		if err := c.needs.fits(c.what, h); err != nil {
			why = append(why, err.Error())
			continue
		}
		out = append(out, c.b)
	}
	switch {
	case len(out) > 0:
		return out, nil
	case len(why) == 0 && variant != "":
		var names []string
		for _, c := range cands {
			names = append(names, c.b.Variant)
		}
		return nil, fmt.Errorf("%s has no build %q here; it has %s", engine, variant, strings.Join(names, ", "))
	}
	return nil, fmt.Errorf("no build of %s runs on this node: %s", engine, strings.Join(why, "; "))
}

// --- vLLM ---

// VLLMVariants are the builds of vLLM a node can be given by choice: CUDA
// 13 from PyPI, and CUDA 12.9 from the project's GitHub release — the same
// version, built for drivers CUDA 13 does not run on.
var VLLMVariants = []string{"cu13", "cu129"}

// pytorchCU129 is where the CUDA 12.9 build's torch is published. PyPI's
// torch is built for CUDA 13; the cu129 vLLM wheel needs this one.
const pytorchCU129 = "https://download.pytorch.org/whl/cu129"

// pytorchPackages are what is taken from PyTorch's index: the builds for one
// CUDA version, which PyPI does not have.
var pytorchPackages = []string{"torch", "torchvision", "torchaudio"}

// vllmCUDA13Since is the first vLLM whose PyPI build is CUDA 13, read from
// its releases: 0.19.1 publishes "+cu130" beside a CUDA 12 default, 0.20.0
// "+cu129" beside a CUDA 13 one.
const vllmCUDA13Since = "0.20"

// vllmMinCompute is the oldest GPU a vLLM version's builds have kernels for,
// from its release pipeline: on x86, 7.0 (Volta) up to 0.19.1 — the
// Dockerfile's default "7.0 7.5 8.0 8.9 9.0 10.0 12.0" — and 7.5 from 0.20.0,
// CUDA 12.9 build included; on arm64, built for Grace machines, 8.7 before
// 0.20 and 8.0 since.
func vllmMinCompute(version, arch string) float64 {
	old := versionBefore(version, vllmCUDA13Since)
	switch {
	case arch == "arm64" && old:
		return 8.7
	case arch == "arm64":
		return 8.0
	case old:
		return 7.0
	}
	return 7.5
}

// vllmNewestFor is the newest vLLM built for GPUs of compute capability cc
// on an architecture, for GPUs older than every build of the version asked
// for: a V100 gets 0.19.1, where 0.20 and later have nothing for it. Nothing
// for a GPU no vLLM was built for.
var vllmNewestFor = []struct {
	arch    string
	from    float64 // the oldest GPU it covers
	version string
}{
	{"amd64", 7.0, "0.19.1"},
}

// vllmCandidates are the vLLM builds for a machine.
//
// The version asked for, in the build the machine runs: CUDA 13 where the
// driver allows, the same version's CUDA 12.9 build where it does not. One
// version across the cluster wherever the GPUs allow it, because the flags
// llmgw builds are checked against it. Where they do not — GPUs older than
// every build of that version — the newest version built for them, said so
// beside the installation: vLLM on a V100 at all, rather than not.
func vllmCandidates(version, variant string, h discovery.Host) ([]Build, error) {
	arch, err := gpuArch("vLLM", h)
	if err != nil {
		return nil, err
	}
	v, note := version, ""
	if cc := h.MinComputeCapability; cc > 0 && cc < vllmMinCompute(version, h.Arch) {
		older := ""
		for _, n := range vllmNewestFor {
			if n.arch == h.Arch && cc >= n.from {
				older = n.version
			}
		}
		if older == "" {
			return nil, fmt.Errorf("no vLLM is built for GPUs of compute capability %.1f on %s; "+
				"the node serves llama.cpp models", cc, arch)
		}
		v, note = older, fmt.Sprintf("vLLM %s, the newest built for compute capability %.1f, instead of %s", older, cc, version)
	}
	pypi := func(variant string) Build {
		b := Build{
			Variant:      variant,
			Python:       pythonForEngines,
			Requirements: []string{"vllm==" + v},
			Binary:       "vllm",
			Packages:     compilerPackages,
			FreeGB:       25,
			GPUCheck:     true,
			Note:         note,
		}
		if v != version {
			b.Version = v
		}
		return b
	}
	minCC := vllmMinCompute(v, h.Arch)
	if versionBefore(v, vllmCUDA13Since) {
		// PyPI's own build is CUDA 12, and is the one there is; asked for
		// cu129, this is it.
		if variant == "cu129" {
			variant = "cu12"
		}
		return pick("vLLM "+v, variant, []candidate{
			{pypi("cu12"), gpuNeeds{minDriverForCUDA12, minCC}, "vLLM " + v},
		}, h)
	}
	cu129 := pypi("cu129")
	cu129.Requirements = nil
	// Found in the release's assets rather than named: the manylinux tag
	// in the file's name changes between versions (2_28, 2_31, 2_34).
	cu129.Wheel = &Wheel{
		Project: "vllm", Repo: "vllm-project/vllm", Tag: "v" + v,
		Prefix: "vllm-" + v + "+cu129-", Suffix: "_" + map[string]string{"amd64": "x86_64", "arm64": "aarch64"}[h.Arch] + ".whl",
	}
	cu129.ExtraIndex = pytorchCU129
	cu129.IndexPackages = pytorchPackages
	return pick("vLLM "+v, variant, []candidate{
		{pypi("cu13"), gpuNeeds{minDriverForCUDA13, minCC}, "vLLM " + v + " built for CUDA 13"},
		{cu129, gpuNeeds{minDriverForCUDA12, minCC}, "vLLM " + v + " built for CUDA 12.9"},
	}, h)
}

// gpuArch says whether a Python GPU engine is published for this machine
// and whether it has an NVIDIA GPU at all.
func gpuArch(engine string, h discovery.Host) (string, error) {
	switch h.Arch {
	case "amd64", "arm64":
	default:
		return "", fmt.Errorf("%s publishes no Linux build for %s", engine, orUnknown(h.Arch))
	}
	if h.NVIDIADriver == "" {
		return "", NotHere{engine + " runs on NVIDIA GPUs and this node has none; it serves llama.cpp models"}
	}
	return h.Arch, nil
}

// --- SGLang ---

// sglangCUDA13Since is the first SGLang built for CUDA 13.
const sglangCUDA13Since = "0.5"

func sglangCandidates(version, variant string, h discovery.Host) ([]Build, error) {
	if _, err := gpuArch("SGLang", h); err != nil {
		return nil, err
	}
	b := Build{
		Variant:      "cu13",
		Python:       pythonForEngines,
		Requirements: []string{"sglang==" + version},
		// SGLang is launched as a module, so what the agent is given is the
		// environment's interpreter, and the proof is importing it.
		Binary:      "python",
		VersionArgs: []string{"-c", "import sglang; print('sglang', sglang.__version__)"},
		Packages:    compilerPackages,
		FreeGB:      25,
		GPUCheck:    true,
	}
	needs := gpuNeeds{minDriverForCUDA13, 7.5}
	if versionBefore(version, sglangCUDA13Since) {
		b.Variant, needs = "cu12", gpuNeeds{minDriverForCUDA12, 7.0}
	}
	return pick("SGLang "+version, variant, []candidate{{b, needs, "SGLang " + version}}, h)
}

// Python engines need a C compiler at run time: Triton builds a small
// launcher for its kernels the first time a model loads, and without gcc
// that load fails long after the installation looked finished.
var compilerPackages = map[string][]string{"deb": {"gcc"}, "rpm": {"gcc"}}

// versionBefore reports whether version a is older than b, comparing the
// dotted numbers in turn ("0.25.1" before "0.26"). A part that is not a
// number counts as zero.
func versionBefore(a, b string) bool {
	if b == "" {
		return false
	}
	pa, pb := strings.Split(strings.TrimPrefix(a, "v"), "."), strings.Split(b, ".")
	for i := 0; i < max(len(pa), len(pb)); i++ {
		x, y := 0, 0
		if i < len(pa) {
			x = leadingInt(pa[i])
		}
		if i < len(pb) {
			y = leadingInt(pb[i])
		}
		if x != y {
			return x < y
		}
	}
	return false
}

// leadingInt is the number a version string starts with: 580 of "580.65.06".
func leadingInt(s string) int {
	n := 0
	for _, c := range strings.TrimSpace(s) {
		if c < '0' || c > '9' {
			break
		}
		n = n*10 + int(c-'0')
	}
	return n
}

// --- llama.cpp ---

// llamaCUDA12Real are the GPUs llama.cpp's CUDA 12.8 build carries machine
// code for. The rest of what it covers — 5.0, 6.1, 7.0, 7.5, 8.0, 9.0 — it
// carries as PTX, which a driver older than the CUDA it came from cannot
// compile.
var llamaCUDA12Real = []float64{8.6, 8.9, 12.0}

// llamaCppCandidates are the llama.cpp builds for a machine.
//
// On an NVIDIA node, the GPU builds best first: CUDA 13.3 (driver 580,
// Turing and newer), CUDA 12.8 on x64 (Maxwell to Blackwell; driver 570 for
// the GPUs it ships as PTX, 525 for those it ships as machine code), and
// Vulkan, which runs on any card whose driver has Vulkan — the generations
// CUDA builds no longer cover. Never the CPU build on an NVIDIA node: that is
// a GPU placement served on the CPU, which is what installing checks for.
// The CPU build where there is no NVIDIA GPU.
//
// The published Linux builds are built on Ubuntu, but nothing about them is
// Ubuntu's: they need glibc 2.34 and the libstdc++ of GCC 12, which Debian 12
// and RHEL 10 have and RHEL 9 does not. So the distribution's name decides
// nothing here. The build is tried everywhere, and ldd says whether this
// system can load it — in the words of the missing symbol.
func llamaCppCandidates(version, variant string, h discovery.Host) ([]Build, error) {
	var arch string
	switch h.Arch {
	case "amd64":
		arch = "x64"
	case "arm64":
		arch = "arm64"
	default:
		return nil, fmt.Errorf("llama.cpp publishes no Linux build for %s", orUnknown(h.Arch))
	}
	base := func(variant string, assets ...string) Build {
		b := Build{
			Variant: variant,
			Assets:  assets,
			// Everything else the build links is in the archive or in any
			// current glibc system; OpenMP's runtime is the one a cloud image
			// lacks, and the server does not start without it.
			Packages: map[string][]string{"deb": {"libgomp1"}, "rpm": {"libgomp"}},
			Binary:   "llama-server",
		}
		if arch == "x64" {
			// The archive carries a backend per CPU generation and loads the
			// best one, down to SSE4.2. Below AVX2 it runs, several times
			// slower.
			b.CPUFlags = []string{"avx2"}
		}
		return b
	}
	if h.NVIDIADriver == "" {
		return pick("llama.cpp", variant, []candidate{
			{base("cpu", fmt.Sprintf("llama-%s-bin-ubuntu-%s.tar.gz", version, arch)), gpuNeeds{}, "llama.cpp"},
		}, h)
	}
	cuda := func(cv string) Build {
		name := fmt.Sprintf("bin-ubuntu-cuda-%s-%s.tar.gz", cv, arch)
		// The release carries the CUDA runtime beside the build, so the node
		// needs a driver and nothing else — no toolkit.
		b := base("cuda-"+cv, fmt.Sprintf("llama-%s-%s", version, name), fmt.Sprintf("cudart-llama-%s-%s", version, name))
		b.GPUBackend = "libggml-cuda.so"
		return b
	}
	vulkan := base("vulkan", fmt.Sprintf("llama-%s-bin-ubuntu-vulkan-%s.tar.gz", version, arch))
	vulkan.GPUBackend = "libggml-vulkan.so"
	vulkan.Packages = map[string][]string{"deb": {"libgomp1", "libvulkan1"}, "rpm": {"libgomp", "vulkan-loader"}}

	cuda12Driver := 570
	if slices.Contains(llamaCUDA12Real, h.MinComputeCapability) {
		cuda12Driver = minDriverForCUDA12
	}
	cands := []candidate{{cuda("13.3"), gpuNeeds{minDriverForCUDA13, 7.5}, "llama.cpp built for CUDA 13.3"}}
	if arch == "x64" {
		cands = append(cands, candidate{cuda("12.8"), gpuNeeds{cuda12Driver, 5.0}, "llama.cpp built for CUDA 12.8"})
	}
	cands = append(cands, candidate{vulkan, gpuNeeds{}, "llama.cpp built for Vulkan"})
	return pick("llama.cpp", variant, cands, h)
}

func orUnknown(s string) string {
	if strings.TrimSpace(s) == "" {
		return "an unrecognised system"
	}
	return s
}