llmgw/internal/agent/engines/probe.go

package engines

import (
	"context"
	"errors"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"time"
)

// Unknown is the version reported for an engine that is installed and would
// not say which version it is.
//
// A value, not an absent key, because the two mean different things and
// placement reads the keys: absent is "this node cannot run it", Unknown is
// "it is here and the probe could not get a number out of it". Collapsing them
// would make a slow engine look like a missing one and take a node out of
// scheduling for being slow to answer --version.
const Unknown = "unknown"

// Prober is an engine that can say whether it is installed on this machine.
//
// Optional, and asserted for rather than part of Engine: an adapter that has no
// way to check should report nothing rather than be forced to invent an answer.
type Prober interface {
	// Probe returns the engine's version, or an error when it is not installed.
	Probe(ctx context.Context) (string, error)
}

// Probe asks every registered engine whether it is here.
//
// The result is what the node reports at registration, and its key set is what
// placement filters on. Until this existed nobody knew: the registry registers
// every adapter whether or not its binary exists, so a node with no vLLM
// accepted a vLLM shard, registered healthy, and failed the shard minutes later
// with "executable file not found" — in the engine's own words, on a machine
// nobody was watching.
//
// Bounded, because it runs on the path to registering and an agent that cannot
// reach the control plane is worse than one that reports an incomplete list.
func (r *Registry) Probe(ctx context.Context) map[string]string {
	out := map[string]string{}
	for name, e := range r.engines {
		p, ok := e.(Prober)
		if !ok {
			continue
		}
		v, err := p.Probe(ctx)
		if err != nil {
			continue
		}
		if v == "" {
			v = Unknown
		}
		out[name] = v
	}
	return out
}

// Names of the engines in a probe result, sorted.
func ProbedNames(m map[string]string) []string {
	out := make([]string, 0, len(m))
	for n := range m {
		out = append(out, n)
	}
	sort.Strings(out)
	return out
}

// probeTimeout bounds one engine's version check.
//
// Generous, because vLLM's --version imports torch and that is seconds on a
// cold page cache, and mean, because this is startup: an engine that has not
// answered in this long is reported as present-with-unknown-version rather
// than waited for.
const probeTimeout = 8 * time.Second

// run executes a version command, saying whether the program was even there.
//
// The two failures are different answers and the caller has to tell them
// apart. A binary that is absent is an engine this node cannot run. A binary
// that is present and exits non-zero is an engine it probably can — unless the
// command was a test of whether a module imports, where a non-zero exit is the
// whole answer.
func run(ctx context.Context, bin string, args ...string) (out string, found bool, ok bool) {
	path, err := exec.LookPath(bin)
	if err != nil {
		return "", false, false
	}
	ctx, cancel := context.WithTimeout(ctx, probeTimeout)
	defer cancel()

	// Combined, not just stdout: engines print their version to either, and one
	// that chose stderr would otherwise read as having no version at all.
	b, err := exec.CommandContext(ctx, path, args...).CombinedOutput()
	return string(b), true, err == nil
}

// installedBinary reports an engine that is its own executable.
//
// It resolves the path and runs nothing. The first version of this ran the
// engine with --version, which is fine for an engine and not fine for the
// wrapper scripts people put in front of one: a script that shifts past flags
// it does not recognise and then execs the server will start an inference
// server, at agent startup, on every node. The agent's own test suite installs
// exactly such a script, and caught it.
//
// So the version is not read here. Presence is what placement needs, and what
// the version is for — telling whether an engine still accepts the flags llmgw
// builds — is `llmgw ctl engine-check`, which is run deliberately by somebody
// who meant to run the engine.
func installedBinary(bin string) (string, error) {
	if _, err := exec.LookPath(bin); err != nil {
		return "", errNotInstalled
	}
	return Unknown, nil
}

// versionOfModule probes an engine that is a module inside an interpreter.
//
// This one does execute, and safely: the command is llmgw's own one-line
// import, not the engine's launcher, so there is no wrapper to shift past
// flags and start a server. There is also no alternative — a module has no
// path to look up.
//
// A non-zero exit is the answer: python3 exists on almost every machine, so
// presence of the interpreter says nothing. Treating a failed import as
// "installed" reported SGLang on every node in a fleet that had none, which
// would take nodes into scheduling for an engine they cannot run — the exact
// failure this whole change exists to stop.
func versionOfModule(ctx context.Context, python string, args ...string) (string, error) {
	out, found, ok := run(ctx, python, args...)
	if !found || !ok {
		return "", errNotInstalled
	}
	return firstVersion(out), nil
}

var errNotInstalled = errors.New("not installed on this node")

var versionPattern = regexp.MustCompile(`\b\d+\.\d+(\.\d+)?\S*`)

// firstVersion is the first thing in the output that looks like a version.
//
// Pattern rather than parsing each engine's format: they print "vllm 0.6.3",
// "version: b4123 (…)", and a paragraph with the number in the middle, and all
// of those change between releases. What this is for is telling one build from
// another, which the first number does.
func firstVersion(s string) string {
	for _, line := range strings.Split(s, "\n") {
		if m := versionPattern.FindString(line); m != "" {
			return strings.TrimRight(m, ",;)")
		}
	}
	return Unknown
}

// Configured is an engine whose location the agent's configuration names.
type Configured interface {
	Configured() bool
}

// ProbeConfigured probes only the engines the configuration names.
//
// Probe looks everywhere — PATH, the default interpreter — and reports what
// it finds, which is exactly the trust in whatever happens to be on a machine
// that installing engines replaces. A node the agent provisions reports what
// it installed, plus what the operator named on purpose; nothing else.
func (r *Registry) ProbeConfigured(ctx context.Context) map[string]string {
	out := map[string]string{}
	for name, e := range r.engines {
		c, ok := e.(Configured)
		if !ok || !c.Configured() {
			continue
		}
		p, ok := e.(Prober)
		if !ok {
			continue
		}
		if v, err := p.Probe(ctx); err == nil {
			if v == "" {
				v = Unknown
			}
			out[name] = v
		}
	}
	return out
}

// Probe reports whether the llama.cpp server is installed.
func (l *LlamaCpp) Probe(context.Context) (string, error) {
	return installedBinary(l.bin())
}

// Configured reports whether the configuration names vLLM.
func (v *VLLM) Configured() bool { return v.cfg.Bin != "" }

// Configured reports whether the configuration names SGLang's interpreter.
func (g *SGLang) Configured() bool { return g.cfg.Python != "" }

// Probe reports whether vLLM is installed.
func (v *VLLM) Probe(context.Context) (string, error) {
	bin := v.cfg.Bin
	if bin == "" {
		bin = "vllm"
	}
	return installedBinary(bin)
}

// Probe reports the SGLang version.
//
// Through the interpreter, because SGLang is a module rather than a binary:
// asking the shell for "sglang" finds nothing on a machine that can run it
// perfectly well.
func (g *SGLang) Probe(ctx context.Context) (string, error) {
	python := g.cfg.Python
	if python == "" {
		python = "python3"
	}
	return versionOfModule(ctx, python, "-c",
		"import sglang, sys; sys.stdout.write(getattr(sglang, '__version__', ''))")
}