llmgw/internal/agent/engines/llamacpp.go

package engines

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
)

// LlamaCppConfig locates the llama.cpp installation.
type LlamaCppConfig struct {
	// Bin is the llama-server binary. It is the server, not the CLI: the
	// gateway talks OpenAI over HTTP and needs the one that serves it.
	// Empty means the agent installs llama.cpp itself, at the version the
	// control plane asks for. Setting it opts out of that.
	Bin string `koanf:"bin"`
	// Parallel is how many requests a server handles at once, when a profile
	// does not say. llama.cpp defaults to one, which would make every second
	// request wait behind the first no matter what the router decided.
	Parallel int `koanf:"parallel"`
	// Threads is how many CPU threads a shard with no GPU computes on. Zero
	// takes half this machine's CPUs.
	//
	// Not all of them, which is llama.cpp's own default: two busy servers
	// that each spin on every CPU fight for them, and on an 8-vCPU node
	// that made 7B generation fall from 10 tokens a second to under 1 each.
	// At half, one server alone is about a third slower and two together
	// keep 5.5 each. A node that only ever runs one CPU replica does better
	// with all of them; say so here.
	Threads int `koanf:"threads"`
}

// numCPU is the CPUs this process may run on; a variable for tests.
var numCPU = runtime.NumCPU

// LlamaCpp builds llama.cpp launch commands.
//
// It is the engine that makes a cluster without GPUs useful: llama.cpp runs a
// small model on CPU at a speed that is fine for testing the whole path, and it
// serves the same OpenAI API as the others, so nothing downstream can tell the
// difference.
type LlamaCpp struct {
	cfg       LlamaCppConfig
	installed atomic.Pointer[string]
}

// Name identifies the engine in the catalog.
func (*LlamaCpp) Name() string { return "llamacpp" }

// HealthPath is llama.cpp's readiness endpoint. It returns 503 while the model
// is still loading, which is exactly what the supervisor needs.
func (*LlamaCpp) HealthPath() string { return "/health" }

// MetricsPath is llama.cpp's Prometheus endpoint, which only exists when the
// server was started with --metrics. Commands always passes it.
func (*LlamaCpp) MetricsPath() string { return "/metrics" }

// ServesHTTP reports whether this rank answers the API. Every llama.cpp shard
// does, because a llama.cpp replica is always exactly one shard.
func (*LlamaCpp) ServesHTTP(ShardSpec) bool { return true }

// Validate checks the profile against the GPUs it would run on.
//
// llama.cpp spreads a model over a node's GPUs by layer (--split-mode layer),
// so gpus_per_node is how many it gets. It has no tensor parallelism and no
// group spanning machines, and accepting tp or nodes would produce a replica
// that serves from fewer cards, or one machine, while the catalog believes
// otherwise.
func (*LlamaCpp) Validate(p Profile, gpus []GPUInfo) error {
	if p.Nodes > 1 {
		return fmt.Errorf("profile %q: llamacpp does not span nodes; use vllm or sglang for a %d-node replica",
			p.Name, p.Nodes)
	}
	if p.TP > 1 || p.PP > 1 {
		return fmt.Errorf("profile %q: llamacpp splits a model across a node's GPUs by layer, not by tp or pp; "+
			"write gpus_per_node: %d instead", p.Name, max(p.TP, 1)*max(p.PP, 1))
	}
	// The common checks count a replica's GPUs as tp; for llama.cpp the
	// count is gpus_per_node.
	q := p
	q.TP = max(p.GPUsPerNode, 1)
	return validateCommon(q, gpus)
}

// Commands returns the processes for one shard, which is always one process.
func (l *LlamaCpp) Commands(s ShardSpec) ([]Command, error) {
	bin := l.bin()
	p := s.Profile

	model, projector, err := ggufIn(p.ModelPath)
	if err != nil {
		return nil, err
	}
	args := []string{
		"--model", model,
		"--host", s.Host(),
		"--port", strconv.Itoa(s.HTTPPort),
		// Without this there is no /metrics, and an endpoint the router cannot
		// measure is one it has to route to blind.
		"--metrics",
	}
	// The model's own chat template rather than a generic one, which is what
	// renders tool definitions in the form the model was trained to read.
	//
	// Pinned rather than relied on: build 10991 has it on by default, older
	// builds do not, and an engine upgrade should not be able to change how
	// prompts are built underneath a running fleet. Passing it explicitly was
	// a no-op on the build here — it was added while chasing tool calls that
	// came back as text, and did not fix that, which turned out to be the
	// model rather than the template.
	if !hasFlag(p.Args, "--jinja", "--no-jinja") {
		args = append(args, "--jinja")
	}
	if p.ServedModelName != "" {
		args = append(args, "--alias", p.ServedModelName)
	}
	// A vision model's projector, fetched beside it: what turns an image in
	// a request into something the language model reads.
	if projector != "" && !hasFlag(p.Args, "--mmproj", "-mm") {
		args = append(args, "--mmproj", projector)
	}
	if !hasFlag(p.Args, "--parallel", "-np") {
		n := l.cfg.Parallel
		if n <= 0 {
			n = 4
		}
		args = append(args, "--parallel", strconv.Itoa(n))
	}
	// With GPUs assigned, offload everything; llama.cpp keeps the model on the
	// CPU otherwise, which is not what a GPU placement asked for. Without them
	// the flag is omitted entirely rather than set to zero, so a build with no
	// GPU support does not see an argument it rejects.
	if len(s.GPUIndices) > 0 && !hasFlag(p.Args, "--n-gpu-layers", "-ngl") {
		args = append(args, "--n-gpu-layers", "999")
	}
	// Several GPUs: the layers spread across all of them, which is the only
	// split llama.cpp does without a matched interconnect. In proportion to
	// each card's free memory, which for the whole cards placement gives is
	// an even split.
	if len(s.GPUIndices) > 1 && !hasFlag(p.Args, "--split-mode", "-sm") {
		args = append(args, "--split-mode", "layer")
	}
	if len(s.GPUIndices) == 0 && !hasFlag(p.Args, "--threads", "-t") {
		// What placement reserved first: it is what the node's other
		// replicas were placed around.
		n := s.Threads
		if n <= 0 {
			n = l.cfg.Threads
		}
		if n <= 0 {
			n = max(numCPU()/2, 1)
		}
		args = append(args, "--threads", strconv.Itoa(n))
	}
	// Loaded at start, applied per request: a request that names an adapter
	// carries "lora": [{"id": i, "scale": s}], and one that names none is the
	// base model. Each is loaded at scale zero, not with --lora and
	// --lora-init-without-apply: that leaves an adapter's scale at 1, and
	// llama-server applies it at that scale to every request with no "lora"
	// of its own — every request for the model itself.
	for _, a := range p.Adapters {
		args = append(args, "--lora-scaled", a.Path+":0")
	}
	args = append(args, p.Args...)

	env := baseEnv(s)
	if s.APIKey != "" {
		env["LLAMA_API_KEY"] = s.APIKey
	}
	// The Vulkan build — what a GPU older than the CUDA builds gets — does
	// not read CUDA_VISIBLE_DEVICES; this is its own. Vulkan numbers devices
	// as the driver lists them, which on NVIDIA follows the same bus order
	// as nvidia-smi; were it not to, the shard would not hold memory on its
	// own GPUs, and starting it checks exactly that.
	//
	// Not on a MIG instance, which Vulkan does not see at all: only CUDA
	// opens one, by the UUID CUDA_VISIBLE_DEVICES carries, and an index here
	// would name some other card.
	if len(s.GPUIndices) > 0 && !s.OnMIG() {
		env["GGML_VK_VISIBLE_DEVICES"] = joinInts(s.GPUIndices)
	}
	return []Command{{
		Name: "llama-server", Path: bin, Args: args, Env: materializeEnv(env),
	}}, nil
}

// hasFlag reports whether the profile already sets one of these flags, in
// either "--flag value" or "--flag=value" form.
func hasFlag(args []string, names ...string) bool {
	for _, a := range args {
		for _, n := range names {
			if a == n || strings.HasPrefix(a, n+"=") {
				return true
			}
		}
	}
	return false
}

// ggufIn resolves a model path that may be a directory.
//
// A multi-file version arrives as the directory it was downloaded into, because
// that is the only shape a set can take. llama.cpp wants a file: given the
// first part of a split GGUF it finds the rest beside it, so the first part is
// what it gets.
//
// A path that is already a file is returned untouched, which is every
// single-file model and every hand-staged one.
//
// A vision projector in the directory — an mmproj file — is not the model:
// it is returned second, to be loaded beside it.
func ggufIn(path string) (model, projector string, err error) {
	st, err := os.Stat(path)
	if err != nil || !st.IsDir() {
		// Not a directory, or not there at all — the second is the engine's to
		// report, with the message it would have given anyway.
		return path, "", nil //nolint:nilerr // a missing path is the engine's error to give
	}
	entries, err := os.ReadDir(path)
	if err != nil {
		return "", "", err
	}
	var names, projectors []string
	for _, e := range entries {
		lower := strings.ToLower(e.Name())
		if e.IsDir() || !strings.HasSuffix(lower, ".gguf") {
			continue
		}
		if strings.Contains(lower, "mmproj") {
			projectors = append(projectors, e.Name())
		} else {
			names = append(names, e.Name())
		}
	}
	if len(names) == 0 {
		if len(projectors) > 0 {
			return "", "", fmt.Errorf("%s holds a vision projector and no model", path)
		}
		return "", "", fmt.Errorf("%s holds no .gguf file", path)
	}
	// Sorted, so a split model's first part comes first: the parts are named
	// "-00001-of-00004" and lexical order is shard order.
	sort.Strings(names)
	if len(projectors) > 0 {
		sort.Strings(projectors)
		projector = filepath.Join(path, projectors[0])
	}
	return filepath.Join(path, names[0]), projector, nil
}

// UseInstalled points the engine at a build the agent installed and checked.
// It takes precedence over everything else, because it is the only one that
// was seen to run.
func (l *LlamaCpp) UseInstalled(path string) { l.installed.Store(&path) }

// Configured reports whether the agent's configuration names a binary. A
// named binary is the operator's deliberate choice; one merely found on PATH
// is not, and is not reported as installed.
func (l *LlamaCpp) Configured() bool { return l.cfg.Bin != "" }

// bin is the llama-server to run: the installed build, else the configured
// one, else whatever PATH has — which is only ever reached by a node not
// provisioned by the agent and not told where its engine is.
func (l *LlamaCpp) bin() string {
	if p := l.installed.Load(); p != nil {
		return *p
	}
	if l.cfg.Bin != "" {
		return l.cfg.Bin
	}
	return "llama-server"
}

// FallsBackToCPU says that llama.cpp keeps running when its GPU backend does
// not load: --n-gpu-layers asks for layers on a GPU it then has none of, and
// the server comes up and answers on the CPU, at a tenth of the speed, with no
// error anywhere. Its readiness on a GPU placement needs checking separately.
func (*LlamaCpp) FallsBackToCPU() bool { return true }

// Devices lists the GPUs this build can compute on, as llama-server's
// --list-devices names them ("CUDA0: NVIDIA L4 (22994 MiB, 22713 MiB free)").
// Empty, on a machine with GPUs, is a build whose GPU backend did not load:
// the CPU build, or a CUDA build missing a library its backend links.
func (l *LlamaCpp) Devices(ctx context.Context) ([]string, error) {
	out, err := exec.CommandContext(ctx, l.bin(), "--list-devices").CombinedOutput() //nolint:gosec // the engine's own binary
	if err != nil {
		return nil, fmt.Errorf("%s --list-devices: %v: %s", l.bin(), err, strings.TrimSpace(string(out)))
	}
	return parseDevices(string(out)), nil
}

// parseDevices reads --list-devices: the lines after "Available devices:",
// which says "(none)" when there are none. Log lines before it are skipped.
func parseDevices(out string) []string {
	var devs []string
	listing := false
	for _, line := range strings.Split(out, "\n") {
		line = strings.TrimSpace(line)
		switch {
		case strings.HasPrefix(line, "Available devices"):
			listing = true
		case listing && line != "" && line != "(none)":
			devs = append(devs, line)
		}
	}
	return devs
}