llmgw/internal/agent/provision.go

package agent

import (
	"context"
	"fmt"
	"log/slog"
	"maps"
	"slices"
	"strings"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
)

// engines is what this node can run: what the agent installed and saw run,
// and what the configuration names and answers. Nothing merely found on the
// machine — a llama-server on PATH is somebody's, not necessarily one that
// works, and placing on it is how a replica fails on a node nobody watches.
func (a *Agent) engines(ctx context.Context) map[string]string {
	a.engineOnce.Do(func() {
		vers := a.reg.ProbeConfigured(ctx)
		a.mu.Lock()
		a.configuredVers, a.lastEngineProbe = vers, time.Now()
		a.mu.Unlock()
		if len(vers) > 0 {
			a.log.Info("engines named in the configuration",
				"engines", engines.ProbedNames(vers))
		}
	})
	a.mu.Lock()
	defer a.mu.Unlock()
	vers := a.engineReportLocked()
	// Whatever registration sends is what the control plane now knows, so a
	// heartbeat repeats it only once it differs.
	a.enginesSent = enginesFingerprint(vers)
	return vers
}

// engineReportLocked merges the two sources. Installed wins: it is the one
// that was seen to run. Called with mu held.
func (a *Agent) engineReportLocked() map[string]string {
	out := maps.Clone(a.configuredVers)
	if out == nil {
		out = map[string]string{}
	}
	maps.Copy(out, a.installedVers)
	maps.DeleteFunc(out, func(e, _ string) bool { return a.disabled(e) })
	return out
}

// byHand reports whether agent.yaml names this engine's binary, which opts
// it out of being installed.
func (a *Agent) byHand(engine string) bool {
	if a.reg == nil {
		return false
	}
	e, err := a.reg.Get(engine)
	if err != nil {
		return false
	}
	c, ok := e.(engines.Configured)
	return ok && c.Configured()
}

// disabled reports whether agent.yaml says this node never runs an engine:
// not installed when asked, and not advertised when found.
func (a *Agent) disabled(engine string) bool {
	return slices.Contains(a.cfg.Engines.Disabled, engine)
}

// provisionEvery is how often installed engines are checked again. A library
// removed by an unrelated upgrade is found here rather than by the next
// replica to start.
const provisionEvery = 10 * time.Minute

// setProvision takes the control plane's list of wanted engines.
func (a *Agent) setProvision(p *llmgwv1.Provision) {
	wants := []provision.Want{}
	for _, e := range p.GetEngines() {
		if a.disabled(e.GetEngine()) {
			continue // a control plane older than the node's say in it
		}
		if a.byHand(e.GetEngine()) {
			// The operator named a binary for it: a local build, a CUDA
			// runtime of their own. Installing one beside it would win —
			// the installed build is used first — and quietly replace theirs.
			continue
		}
		w := provision.Want{Engine: e.GetEngine(), Version: e.GetVersion(), Variant: e.GetVariant()}
		if base := e.GetOfflineBase(); base != "" {
			// An offline cluster: from its mirror, checked against what the
			// control plane holds, never the publisher.
			w.Offline = &provision.Offline{Base: base, SHA256: maps.Clone(e.GetOfflineSha256())}
		}
		wants = append(wants, w)
	}
	a.mu.Lock()
	a.provWants = wants
	a.mu.Unlock()
	select {
	case a.provKick <- struct{}{}:
	default:
	}
}

// provisionLoop installs and re-checks the wanted engines, one pass at a time.
func (a *Agent) provisionLoop(ctx context.Context) {
	// After a failure the next pass comes sooner, and backs off: a lock held
	// by a new machine's own updates clears in a minute or two, while ten
	// minutes' wait left a node idle for no reason.
	next := provisionEvery
	retry := time.Minute
	for {
		select {
		case <-ctx.Done():
			return
		case <-a.provKick:
		case <-time.After(next):
		}
		a.mu.Lock()
		wants := slices.Clone(a.provWants)
		managed := a.provWants != nil
		a.mu.Unlock()
		if !managed {
			continue
		}

		statuses := a.prov.Converge(ctx, wants)
		next = provisionEvery
		for _, st := range statuses {
			if st.State == provision.Failed {
				next, retry = retry, min(retry*2, provisionEvery)
				break
			}
		}
		if next == provisionEvery {
			retry = time.Minute
		}
		installed := map[string]string{}
		keep := map[string]bool{}
		for _, st := range statuses {
			keep[st.Engine] = true
			if st.State == provision.Ready {
				a.reg.UseInstalled(st.Engine, st.Binary)
				installed[st.Engine] = st.Version
			}
		}
		a.mu.Lock()
		if enginesFingerprint(installed) != enginesFingerprint(a.installedVers) {
			a.log.Info("engines installed and running", "engines", engines.ProbedNames(installed))
		}
		a.installedVers = installed
		// An engine no longer asked for is no longer reported on.
		for e := range a.provStatus {
			if !keep[e] {
				delete(a.provStatus, e)
			}
		}
		a.mu.Unlock()
	}
}

// provisionStatus records a change in one engine's installation.
func (a *Agent) provisionStatus(st provision.Status) {
	a.mu.Lock()
	prev, had := a.provStatus[st.Engine]
	a.provStatus[st.Engine] = st
	a.mu.Unlock()
	if !had || prev.State != st.State || prev.Detail != st.Detail {
		lvl := slog.LevelInfo
		if st.State == provision.Failed || st.State == provision.Unsupported {
			lvl = slog.LevelWarn
		}
		a.log.Log(context.Background(), lvl, "engine installation",
			"engine", st.Engine, "version", st.Version, "variant", st.Variant,
			"state", st.State, "detail", st.Detail)
	}
}

func provisionFingerprint(m map[string]provision.Status) string {
	var b strings.Builder
	for _, k := range slices.Sorted(maps.Keys(m)) {
		st := m[k]
		fmt.Fprintf(&b, "%s|%s|%s|%s|%s\n", st.Engine, st.Version, st.Variant, st.State, st.Detail)
	}
	return b.String()
}

// engineReprobe is how often the installed engines are looked for again.
//
// Minutes, not seconds: an engine is installed by somebody at a keyboard, so a
// minute's delay in noticing is invisible, while SGLang's probe starts a
// Python interpreter that imports it.
const engineReprobe = time.Minute

// reprobeEngines looks for installed engines again when it is time, off the
// heartbeat goroutine so a slow probe never delays a heartbeat.
func (a *Agent) reprobeEngines() {
	a.mu.Lock()
	due := !a.probingEngines && !a.lastEngineProbe.IsZero() &&
		time.Since(a.lastEngineProbe) >= engineReprobe
	if due {
		a.probingEngines = true
	}
	a.mu.Unlock()
	if !due {
		return
	}
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		vers := a.reg.ProbeConfigured(ctx)

		a.mu.Lock()
		changed := enginesFingerprint(vers) != enginesFingerprint(a.configuredVers)
		a.configuredVers, a.lastEngineProbe, a.probingEngines = vers, time.Now(), false
		a.mu.Unlock()
		if changed {
			a.log.Info("engines named in the configuration changed",
				"engines", engines.ProbedNames(vers))
		}
	}()
}

func enginesFingerprint(vers map[string]string) string {
	var b strings.Builder
	for _, k := range slices.Sorted(maps.Keys(vers)) {
		fmt.Fprintf(&b, "%s=%s\n", k, vers[k])
	}
	return b.String()
}