llmgw/internal/agent/agent.go

// Package agent is the node side of the control plane stream (plan §5.6).
package agent

import (
	"context"
	"crypto/sha256"
	"crypto/tls"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"maps"
	"math/rand/v2"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/http2"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// Config configures the agent.
type Config struct {
	// NodeID defaults to the hostname. It must equal the common name of the
	// agent's certificate, so set it only where the two differ.
	NodeID string `koanf:"node_id"`
	// Server is the control plane, and Servers is the same thing when there is
	// more than one. Both are accepted: every configuration written before
	// there could be a second control plane names one address, and none of
	// them should have to change.
	//
	// The order is tried in order, and a refusal moves to the next. Only the
	// leader accepts an agent, so a standby is skipped in the time it takes to
	// be told so — which is what makes failover a reconnect rather than an
	// edit to every node's configuration.
	Server  string   `koanf:"server"`
	Servers []string `koanf:"servers"`
	// Hostname defaults to the machine's.
	Hostname string `koanf:"hostname"`
	// MgmtAddr is the address engines are reached on. It defaults to this
	// machine's address on the route to the control plane, which is the one
	// the control plane and the router can reach.
	MgmtAddr string `koanf:"mgmt_addr"`
	// Labels are free-form facts about a node for a profile's node_selector.
	// Nothing needs them: memory, GPUs, engines and disk are measured, and a
	// label restating one of those goes stale the day the machine changes.
	Labels map[string]string `koanf:"labels"`

	// MgmtInterface is the ethernet interface NCCL bootstraps over. It
	// defaults to the interface holding MgmtAddr, then eth0.
	MgmtInterface string `koanf:"mgmt_interface"`
	// Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's
	// default subnet prefix is the same on every fabric, so nothing can be
	// inferred from the wire (plan §5.2).
	Fabrics []inventory.FabricAttachment `koanf:"fabrics"`

	HeartbeatInterval time.Duration `koanf:"heartbeat_interval"`
	// ReconnectBase is the first reconnect delay; it backs off from there.
	ReconnectBase time.Duration `koanf:"reconnect_base"`
	ReconnectMax  time.Duration `koanf:"reconnect_max"`

	// AdoptReadyTimeout bounds how long an adopted engine is given to answer a
	// health check before the agent stops waiting on it.
	//
	// Shorter than a startup timeout on purpose: this engine is already
	// running, so it is not loading a model, it is either serving or it is not.
	AdoptReadyTimeout time.Duration `koanf:"adopt_ready_timeout"`

	// TLS authenticates both ends of the control-plane stream. Leaving it unset
	// keeps the plaintext behaviour, which only belongs on a trusted network.
	TLS tlsconf.Config `koanf:"tls"`

	// ModelDirs are the directories this node keeps model weights in, beyond
	// WeightsCache, which is always looked in. Storage differs at every installation,
	// so llmgw observes it rather than owning it: a Lustre mount, an image
	// exported over the fabric, local NVMe and a copy made by hand all look
	// identical from here, and placement only needs to know whether the file
	// is present.
	ModelDirs []string `koanf:"model_dirs"`

	// EnginesDir is where the agent installs engines, one directory per
	// build. Defaults to /opt/llmgw/engines.
	EnginesDir string `koanf:"engines_dir"`

	// WeightsCache is where a remote source is downloaded to. Defaults to the
	// first entry of ModelDirs, because a node that says where its models live
	// has already answered this question, and otherwise to
	// /var/lib/llmgw-weights.
	//
	// Fetching is the simple arrangement: every node keeps its own copy. The
	// exotic ones need no code — weights on a shared filer or an image exported
	// over the fabric are simply present, and present means nothing is fetched.
	WeightsCache string `koanf:"weights_cache"`

	Discovery discovery.Options `koanf:"discovery"`
	Engines   engines.Config    `koanf:"engines"`
	Superv    supervisor.Config `koanf:"supervisor"`

	Version string
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	// No MgmtInterface: normalize finds the one holding the management
	// address, and falls back to eth0 only when it cannot. Defaulted here, it
	// was never found — every node said eth0 unless told otherwise, and a node
	// on ens3 or bond0 had NCCL bootstrap over an interface it does not have.
	return Config{
		HeartbeatInterval: 5 * time.Second,
		ReconnectBase:     time.Second,
		ReconnectMax:      30 * time.Second,
		AdoptReadyTimeout: 60 * time.Second,
		Superv:            supervisor.DefaultConfig(),
		// NVML gives NVLink topology, so prefer it when this build has it.
		Discovery: discovery.Options{PreferNVML: true},
	}
}

func (c *Config) normalize() error {
	d := DefaultConfig()
	if err := c.TLS.Validate(); err != nil {
		return err
	}
	if err := c.normalizeServerURL(); err != nil {
		return err
	}
	// Everything below is found rather than written. A file that has to say
	// which machine it is on, and at what address, is a file that differs on
	// every node — and a copy taken from one node to the next registers the
	// second as the first.
	host, _ := os.Hostname()
	if c.NodeID == "" {
		c.NodeID = host
	}
	if c.NodeID == "" {
		return errors.New("node_id is not set and the hostname could not be read")
	}
	if c.Hostname == "" {
		c.Hostname = host
	}
	if c.MgmtAddr == "" {
		c.MgmtAddr = addrToward(c.Server)
	}
	if c.MgmtInterface == "" {
		c.MgmtInterface = interfaceOf(c.MgmtAddr)
	}
	if c.MgmtInterface == "" {
		c.MgmtInterface = "eth0"
	}
	// The directory downloads land in is always one that is looked in. When
	// they were two settings, a node could fill its disk with fetched weights
	// that placement never saw — which is what every node here was doing.
	if c.WeightsCache == "" && len(c.ModelDirs) > 0 {
		c.WeightsCache = c.ModelDirs[0]
	}
	if c.WeightsCache == "" {
		c.WeightsCache = DefaultWeightsCache
	}
	if !slices.Contains(c.ModelDirs, c.WeightsCache) {
		c.ModelDirs = append(c.ModelDirs, c.WeightsCache)
	}
	if c.HeartbeatInterval <= 0 {
		c.HeartbeatInterval = d.HeartbeatInterval
	}
	if c.ReconnectBase <= 0 {
		c.ReconnectBase = d.ReconnectBase
	}
	if c.AdoptReadyTimeout <= 0 {
		c.AdoptReadyTimeout = d.AdoptReadyTimeout
	}
	if c.ReconnectMax < c.ReconnectBase {
		c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax)
	}
	return nil
}

// normalizeServerURL makes the configured address a URL connect can use, and
// refuses the combination that would quietly undo the certificates: a scheme
// that says plaintext while TLS is configured. Silently connecting anyway would
// mean a cluster believing it has mutual authentication when it has none.
func (c *Config) normalizeServerURL() error {
	// One list, however it was written. Server first when both are given,
	// because an operator adding a standby to an existing file expects the
	// one already there to stay preferred.
	all := make([]string, 0, 1+len(c.Servers))
	if c.Server != "" {
		all = append(all, c.Server)
	}
	all = append(all, c.Servers...)

	seen := map[string]struct{}{}
	out := make([]string, 0, len(all))
	for _, s := range all {
		fixed, err := normalizeOne(s, c.TLS.Enabled())
		if err != nil {
			return err
		}
		// A duplicate would be tried twice in a row on failover, which reads
		// as the rotation not working.
		if _, dup := seen[fixed]; dup {
			continue
		}
		seen[fixed] = struct{}{}
		out = append(out, fixed)
	}
	if len(out) == 0 {
		return errors.New("server is required")
	}
	c.Servers = out
	c.Server = out[0]
	return nil
}

// DefaultWeightsCache is where fetched weights go when nothing says otherwise.
const DefaultWeightsCache = "/var/lib/llmgw-weights"

// addrToward is this machine's address on the route to the control plane.
//
// The address engines are reached on is whichever one the control plane can
// reach, and the kernel already knows which that is: a UDP "connection" picks
// the source address for the route without sending a packet. Asking the
// interfaces instead would have to guess among docker bridges, VPNs and
// link-local addresses.
func addrToward(server string) string {
	u, err := url.Parse(server)
	if err != nil || u.Host == "" {
		return ""
	}
	host := u.Hostname()
	port := u.Port()
	if port == "" {
		port = "443"
	}
	conn, err := net.Dial("udp", net.JoinHostPort(host, port))
	if err != nil {
		return ""
	}
	defer conn.Close()
	if a, ok := conn.LocalAddr().(*net.UDPAddr); ok {
		return a.IP.String()
	}
	return ""
}

// interfaceOf names the interface that holds an address, or "" when none does.
func interfaceOf(addr string) string {
	ip := net.ParseIP(addr)
	if ip == nil {
		return ""
	}
	ifaces, err := net.Interfaces()
	if err != nil {
		return ""
	}
	for _, ifc := range ifaces {
		addrs, err := ifc.Addrs()
		if err != nil {
			continue
		}
		for _, a := range addrs {
			if n, ok := a.(*net.IPNet); ok && n.IP.Equal(ip) {
				return ifc.Name
			}
		}
	}
	return ""
}

// normalizeOne applies the scheme rules to a single address.
func normalizeOne(addr string, secure bool) (string, error) {
	switch {
	case strings.HasPrefix(addr, "https://"):
		if !secure {
			return "", fmt.Errorf("%s is an https:// address but no tls certificates are configured", addr)
		}
	case strings.HasPrefix(addr, "http://"):
		if secure {
			return "", fmt.Errorf("tls is configured but %s is an http:// address; use https://", addr)
		}
	case strings.Contains(addr, "://"):
		return "", fmt.Errorf("server %q must be an http:// or https:// address", addr)
	case secure:
		return "https://" + addr, nil
	default:
		return "http://" + addr, nil
	}
	return addr, nil
}

// Agent runs on a GPU node.
type Agent struct {
	cfg  Config
	log  *slog.Logger
	sup  *supervisor.Supervisor
	reg  *engines.Registry
	http connectHTTPClient

	// engineOnce probes the configured engines once before registering.
	//
	// Once, because probing SGLang imports it in a Python interpreter — seconds
	// on a cold cache — and a reconnect is not a reason to pay that again.
	// Later changes are found by the slower re-probe the heartbeat starts.
	engineOnce sync.Once
	// configuredVers are the engines the configuration names and that answer.
	// Guarded by mu once the first probe has run.
	configuredVers map[string]string
	// installedVers are the engines this agent installed and saw run.
	// Guarded by mu.
	installedVers map[string]string

	// prov installs the engines the control plane asks for. provWants is what
	// it asked for, nil until it has; provStatus is where each stands and
	// provSent the fingerprint of what the control plane last heard. All
	// guarded by mu. provKick starts a pass without waiting for the timer.
	logRing *logRing

	prov       *provision.Installer
	provWants  []provision.Want
	provStatus map[string]provision.Status
	provSent   string
	provKick   chan struct{}
	// enginesSent is the fingerprint of the engines the control plane last
	// heard about, from registration or a heartbeat. Guarded by mu.
	enginesSent string
	// lastEngineProbe and probingEngines pace the re-probe. Guarded by mu.
	lastEngineProbe time.Time
	probingEngines  bool

	mu  sync.Mutex
	inv discovery.Inventory
	// shardReplica remembers which replica a shard belongs to, so a status
	// report can name both. The control plane keys on the replica.
	shardReplica map[string]string
	fetcher      *weights.Fetcher

	// lastWeights is the fingerprint of the file list most recently sent, so a
	// heartbeat carries the list only when it changed. Guarded by mu: during a
	// failover the old stream's heartbeat loop can still be running when the
	// new one's starts.
	lastWeights string
	// fabricsSent is what the control plane was last told about fabrics.
	fabricsSent string
	// sm caches which subnet manager each InfiniBand port sees, asked once
	// a minute: a management datagram per port per heartbeat is more than
	// an answer that changes on a failover is worth.
	sm      map[string]*discovery.SMInfo
	smErr   map[string]string
	smAsked time.Time
	// nvlAsked is when the GPUs' NVLink fabric state was last read.
	nvlAsked time.Time
	// sms is every subnet manager on a fabric, asked through one port per
	// master, keyed by that port; smsErr says why a port could not ask.
	sms    map[string][]discovery.SMRecord
	smsErr map[string]string
	// load measures the machine between heartbeats.
	load *discovery.LoadSampler

	// fetchWatchers are the shards waiting on each source, so progress on one
	// download reaches every shard that asked for it. A fetch is single-flight
	// — eight shards starting together produce one download — and all eight
	// are waiting, so all eight should say so.
	fetchWatchers map[string][]func(weights.Progress)
}

// connectHTTPClient is the HTTP client connect uses; an interface so tests can
// hand in an httptest server's client.
type connectHTTPClient interface {
	Do(*http.Request) (*http.Response, error)
}

// New builds an agent.
func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, error) {
	if err := cfg.normalize(); err != nil {
		return nil, err
	}
	if log == nil {
		log = slog.Default()
	}
	// Everything the agent logs also goes to a ring the admin page can read,
	// so seeing what a node did does not mean ssh and a journal only root may
	// open.
	ring := newLogRing(agentLogLines)
	log = slog.New(teeHandler{log.Handler(), slog.NewTextHandler(ring, nil)})
	if httpClient == nil {
		httpClient = http.DefaultClient
	}

	health := supervisor.NewHTTPHealth(3 * time.Second)
	sup, err := supervisor.New(cfg.Superv, health, log)
	if err != nil {
		return nil, err
	}
	a := &Agent{
		cfg: cfg, log: log, sup: sup,
		reg:          engines.NewRegistry(cfg.Engines),
		http:         httpClient,
		shardReplica: map[string]string{},
		provStatus:   map[string]provision.Status{},
		logRing:      ring,
		provKick:     make(chan struct{}, 1),
		load:         discovery.NewLoadSampler(),
	}
	a.prov = &provision.Installer{Root: cfg.EnginesDir, Log: log, OnStatus: a.provisionStatus}
	return a, nil
}

// Supervisor exposes the process supervisor, for tests and `ctl logs`.
func (a *Agent) Supervisor() *supervisor.Supervisor { return a.sup }

// Run connects to the control plane and stays connected until ctx ends.
//
// Reconnecting forever is the right behaviour for a node agent: a control plane
// restart or a network partition must not require touching every GPU node, and
// the engines keep serving traffic throughout because the router talks to them
// directly.
func (a *Agent) Run(ctx context.Context) error {
	if err := a.scan(ctx); err != nil {
		// Not fatal: a node whose GPUs cannot be read should still register, so
		// an operator can see it and why it is useless.
		a.log.Error("hardware discovery failed", "err", err)
	}
	a.recoverOrphans()
	// Made now rather than by the first download, so a new node reports its
	// disk from its first heartbeat instead of only once it holds a model.
	if err := os.MkdirAll(a.cfg.WeightsCache, 0o755); err != nil {
		a.log.Warn("cannot create the weights directory", "dir", a.cfg.WeightsCache, "err", err)
	}
	go a.provisionLoop(ctx)
	go a.ensureDiags(ctx)

	attempt := 0
	next := 0
	for ctx.Err() == nil {
		server := a.cfg.Servers[next%len(a.cfg.Servers)]
		err := a.session(ctx, server)
		if ctx.Err() != nil {
			break
		}
		// On to the next address whatever went wrong. A standby refuses, a
		// stopped control plane refuses the connection, and a crashed one
		// never answers — from here they are the same thing, and trying
		// somewhere else is the answer to all three.
		next++
		attempt++
		delay := a.reconnectDelay(attempt)
		// No wait when there is somewhere else to try: the delay exists to
		// stop hammering one control plane, and the next address has not been
		// hammered. With a single address this is the old behaviour exactly.
		if len(a.cfg.Servers) > 1 && next%len(a.cfg.Servers) != 0 {
			delay = 0
		}
		a.log.Warn("disconnected from the control plane, retrying",
			"err", err, "server", server,
			"next", a.cfg.Servers[next%len(a.cfg.Servers)],
			"in", delay.Round(time.Millisecond), "attempt", attempt)

		select {
		case <-ctx.Done():
		case <-time.After(delay):
		}
		if err == nil {
			attempt = 0
		}
	}

	// Engines keep running when the agent stops on purpose: an agent restart
	// must not interrupt served traffic. Orphan adoption picks them up again.
	return ctx.Err()
}

func (a *Agent) reconnectDelay(attempt int) time.Duration {
	d := a.cfg.ReconnectBase << min(attempt-1, 16)
	if d > a.cfg.ReconnectMax || d <= 0 {
		d = a.cfg.ReconnectMax
	}
	// Jitter, so a fleet does not reconnect in lockstep after a control plane
	// restart and knock it over again.
	return d/2 + time.Duration(rand.Int64N(int64(d/2))) //nolint:gosec // jitter
}

// recoverOrphans takes back engines a previous agent left running.
//
// Engines deliberately outlive their agent — the router talks to them directly,
// so an agent restart must not interrupt served traffic. The cost is that they
// have to be found again, and anything that cannot be identified has to be
// killed: a process holding GPUs that nothing can describe is worse than no
// process at all, because placement will hand those GPUs out again.
func (a *Agent) recoverOrphans() {
	orphans, err := a.sup.FindOrphans()
	if err != nil {
		a.log.Warn("scanning for leftover engines failed", "err", err)
		return
	}

	for _, o := range orphans {
		switch {
		case !o.Alive:
			// The process is gone; only its state file remains.
			if err := a.sup.KillOrphan(o); err != nil {
				a.log.Warn("cleaning up a dead shard's state failed", "shard", o.ShardID, "err", err)
			}
		case o.ReplicaID == "":
			// Alive, but its state file could not be read, so it cannot be
			// described to the control plane or matched to a replica.
			a.log.Warn("killing an unidentifiable engine", "shard", o.ShardID, "pid", o.PID)
			if err := a.sup.KillOrphan(o); err != nil {
				a.log.Warn("killing it failed", "shard", o.ShardID, "err", err)
			}
		default:
			sh, err := a.sup.Adopt(o)
			if err != nil {
				a.log.Warn("adopting a shard failed", "shard", o.ShardID, "err", err)
				continue
			}
			a.mu.Lock()
			a.shardReplica[o.ShardID] = o.ReplicaID
			a.mu.Unlock()
			a.log.Info("adopted a running engine",
				"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
			go a.settleAdopted(sh, o)
		}
	}
}

// settleAdopted brings an adopted engine to Ready, or lets it fail.
//
// Adoption stopped at "the process is alive", which left the shard in Starting
// with nothing that would ever move it. The control plane then waited out the
// profile's startup timeout and killed a perfectly healthy engine so it could
// start an identical one — so an agent restart, which exists precisely so that
// engines are not interrupted, cost every model on the node a full reload.
//
// Measured on node05: the engine answered /health the whole time, and the
// replica sat in Starting for three minutes until the timeout reaped it.
//
// Run in the background because recovery happens before the agent connects, and
// a slow health check must not delay registration — the control plane needs to
// hear about these shards promptly whatever state they are in.
func (a *Agent) settleAdopted(sh *supervisor.Shard, o supervisor.Orphan) {
	ctx, cancel := context.WithTimeout(context.Background(), a.cfg.AdoptReadyTimeout)
	defer cancel()

	eng, err := a.reg.Get(o.Spec.Profile.Engine)
	if err != nil {
		a.log.Warn("cannot judge an adopted engine; leaving it to the control plane",
			"shard", o.ShardID, "err", err)
		return
	}

	// The same two paths a freshly started shard takes: a rank that serves HTTP
	// is health-checked, a worker is judged by staying alive.
	if eng.ServesHTTP(o.Spec) {
		base := fmt.Sprintf("http://127.0.0.1:%d", o.Spec.HTTPPort)
		err = a.sup.WaitReady(ctx, sh, base, o.Spec.Profile.ServedModelName, a.cfg.AdoptReadyTimeout)
	} else {
		err = a.sup.WaitAlive(ctx, sh)
	}
	if err != nil {
		// Not killed here. An adopted engine that will not answer is the
		// control plane's decision to replace, and it has the timeout for it;
		// killing it locally would race that.
		a.log.Warn("an adopted engine did not become ready",
			"shard", o.ShardID, "replica", o.ReplicaID, "err", err)
		return
	}
	a.log.Info("adopted engine is serving",
		"shard", o.ShardID, "replica", o.ReplicaID)
}

// usableFabrics reports the fabrics this node can actually carry traffic on.
//
// Loud about what it drops. A fabric silently missing from a registration looks
// identical to one never configured, and the operator who typed the adapter
// name is the only person who can fix it — so the reason is logged here, on the
// machine where both halves are visible, rather than inferred later from a
// replica that will not converge.
func (a *Agent) usableFabrics(inv discovery.Inventory) []inventory.FabricAttachment {
	out := a.checkFabrics(inv.HCAs, true)
	a.mu.Lock()
	a.fabricsSent = fabricsFingerprint(out)
	a.mu.Unlock()
	return out
}

// checkFabrics is the declared fabrics this node can carry traffic on, from
// the adapters as they are now. Loud logs every problem; the heartbeat checks
// quietly and speaks only when the answer changes.
func (a *Agent) checkFabrics(hcas []discovery.HCA, loud bool) []inventory.FabricAttachment {
	declared := map[string][]string{}
	for _, f := range a.cfg.Fabrics {
		declared[f.Fabric] = f.HCAs
	}

	var out []inventory.FabricAttachment
	for _, c := range discovery.CheckFabrics(declared, hcas) {
		if loud {
			for _, p := range c.Problems {
				a.log.Error("a declared fabric adapter cannot carry traffic",
					"fabric", c.Fabric, "problem", p)
			}
		}
		if !c.OK() {
			if loud {
				a.log.Error("fabric not advertised; no declared adapter on it can carry traffic",
					"fabric", c.Fabric,
					"consequence", "no multi-node replica will be placed here")
			}
			continue
		}
		out = append(out, inventory.FabricAttachment{
			Fabric: c.Fabric, HCAs: c.Usable, RateGbps: c.RateGbps,
		})
	}
	return out
}

// fabricsFingerprint is enough to notice a fabric, an adapter or a rate
// changing.
func fabricsFingerprint(fs []inventory.FabricAttachment) string {
	var b strings.Builder
	for _, f := range fs {
		fmt.Fprintf(&b, "%s:%s:%d;", f.Fabric, strings.Join(f.HCAs, ","), f.RateGbps)
	}
	return b.String()
}

// sampleIB re-reads the adapters for the heartbeat: every port's state and
// counters, and the fabrics again when they no longer match what the control
// plane was told.
//
// Every heartbeat, because it is a few small files per port, and because a
// port that went down used to be invisible until the agent restarted — the
// node went on being offered multi-node work that would hang on its first
// collective.
func (a *Agent) sampleIB(hb *llmgwv1.Heartbeat) {
	a.mu.Lock()
	had := len(a.inv.HCAs) > 0
	a.mu.Unlock()
	if !had && len(a.cfg.Fabrics) == 0 {
		return
	}
	hcas, err := discovery.DiscoverIB(a.cfg.Discovery.IBRoot)
	if err != nil {
		return
	}
	a.attachSM(hcas)
	for _, h := range hcas {
		for _, p := range h.Ports {
			ip := &llmgwv1.IBPort{
				Hca: h.Name, Port: int32(p.Num), State: p.State, PhysState: p.PhysState,
				RateGbps: int32(p.RateGbps), LinkLayer: string(p.LinkLayer), Counters: p.Counters,
				NodeGuid: h.NodeGUID, PortGuid: p.PortGUID, Network: p.Network, GidPrefix: p.GIDPrefix,
			}
			a.mu.Lock()
			if p.SM != nil {
				ip.SmGuid, ip.SmLid, ip.SmPriority, ip.SmState = p.SM.GUID, int32(p.SM.LID), int32(p.SM.Priority), p.SM.State
			} else {
				ip.SmError = a.smErr[smKey(h.Name, p.Num)]
			}
			for _, r := range a.sms[smKey(h.Name, p.Num)] {
				ip.SubnetManagers = append(ip.SubnetManagers, &llmgwv1.SubnetManager{
					Guid: r.GUID, Lid: int32(r.LID), Priority: int32(r.Priority), State: r.State})
			}
			ip.SubnetManagersError = a.smsErr[smKey(h.Name, p.Num)]
			a.mu.Unlock()
			hb.IbPorts = append(hb.IbPorts, ip)
		}
	}
	a.mu.Lock()
	a.inv.HCAs = hcas // a reconnect registers what is true now
	a.mu.Unlock()
	if len(a.cfg.Fabrics) == 0 {
		return
	}
	now := a.checkFabrics(hcas, false)
	fp := fabricsFingerprint(now)
	a.mu.Lock()
	changed := fp != a.fabricsSent
	a.fabricsSent = fp
	a.mu.Unlock()
	if !changed {
		return
	}
	// Said once, with the reasons, the moment it happens.
	a.checkFabrics(hcas, true)
	a.log.Warn("the fabrics this node can carry traffic on changed", "now", fp)
	hb.FabricsChanged = true
	for _, f := range now {
		hb.Fabrics = append(hb.Fabrics, &llmgwv1.FabricAttachment{
			Fabric: f.Fabric, Hcas: f.HCAs, RateGbps: int32(f.RateGbps),
		})
	}
}

// weights builds the fetcher lazily, so an agent that never starts a shard
// never needs a cache directory.
func (a *Agent) weights() *weights.Fetcher {
	a.mu.Lock()
	defer a.mu.Unlock()
	if a.fetcher != nil {
		return a.fetcher
	}
	cache := a.cfg.WeightsCache
	if cache == "" && len(a.cfg.ModelDirs) > 0 {
		cache = a.cfg.ModelDirs[0]
	}
	a.fetcher = &weights.Fetcher{
		CacheDir: cache,
		Log:      a.log,
		OnProgress: func(p weights.Progress) {
			a.log.Info("fetching weights",
				"source", p.Source, "percent", p.Percent(),
				"bytes", p.BytesDone, "total", p.BytesTotal)
			// And upward, to whoever is watching from the control plane. The
			// log line above is on the node, which is where nobody is looking
			// when a replica has been starting for ten minutes.
			a.mu.Lock()
			watchers := slices.Clone(a.fetchWatchers[p.Source])
			a.mu.Unlock()
			for _, w := range watchers {
				if w != nil {
					w(p)
				}
			}
		},
	}
	return a.fetcher
}

// runningShards describes what this agent is supervising, for registration.
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
	var out []*llmgwv1.RunningShard
	for _, st := range a.sup.Statuses() {
		if st.State.Terminal() {
			continue
		}
		sh, err := a.sup.Get(st.ID)
		if err != nil {
			continue
		}
		spec := sh.Spec
		out = append(out, &llmgwv1.RunningShard{
			ReplicaId:  spec.ReplicaID,
			ShardId:    st.ID,
			Rank:       int32(spec.Rank),
			Nnodes:     int32(spec.NNodes),
			GpuIndices: int32Slice(spec.GPUIndices),
			Hcas:       spec.HCAs,
			HttpPort:   int32(spec.HTTPPort),
			State:      agentsvc.StateToWire(st.State),
			Model:      spec.Profile.ServedModelName,
			Profile:    spec.Profile.Name,
			Pool:       spec.Pool,
			Fabric:     spec.Fabric,
			ProfileFp:  spec.ProfileFP,
			LeaderAddr: a.cfg.MgmtAddr,
			LeaderPort: int32(spec.LeaderPort),
		})
	}
	return out
}

func int32Slice(v []int) []int32 {
	if len(v) == 0 {
		return nil
	}
	out := make([]int32, len(v))
	for i, x := range v {
		out[i] = int32(x)
	}
	return out
}

// scan refreshes the node's hardware inventory.
func (a *Agent) scan(ctx context.Context) error {
	// Carried from the agent's own configuration rather than duplicated into
	// the discovery block, so an operator names the directories once.
	opts := a.cfg.Discovery
	opts.ModelDirs = a.cfg.ModelDirs
	inv, err := discovery.Discover(ctx, opts, a.log)
	host := discovery.DiscoverHost()
	host.MinComputeCapability = discovery.MinComputeCapability(inv.GPUs)
	a.mu.Lock()
	a.inv = inv
	a.prov.Host = host
	a.mu.Unlock()
	a.log.Info("this node", "os", host.OSID+" "+host.OSVersion, "arch", host.Arch,
		"cpu_flags", host.CPUFlags, "nvidia_driver", host.NVIDIADriver)
	return err
}

// session runs one connection until it breaks.
func (a *Agent) session(ctx context.Context, server string) error {
	client := llmgwv1connect.NewAgentServiceClient(a.http, server)

	// parent is kept because the derived context is cancelled from inside this
	// function; asking the derived one whether we are shutting down would
	// always say yes, and every stream error would read as a clean close.
	parent := ctx
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	stream := client.Connect(ctx)
	defer func() { _ = stream.CloseRequest() }()

	a.mu.Lock()
	inv := a.inv
	a.mu.Unlock()

	// Declared fabric membership against what the node actually has. A fabric
	// whose adapters are absent or have no active port is not advertised at
	// all, so the control plane will not place a multi-node replica on it —
	// rather than placing one that starts, passes health checks and hangs on
	// its first all-reduce.
	reg := agentsvc.RegisterFromInventory(
		a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
		a.cfg.Labels, inv, a.usableFabrics(inv))
	// Every registration reports what is already running, not just the first
	// one: a reconnect may be a control plane restart, and it has to recover
	// the same way (plan §6.5).
	// What this node can actually run, so the control plane stops placing
	// engines that are not here. The registry holds an adapter for every
	// engine whether or not its binary exists, so without this a node with no
	// vLLM accepted a vLLM shard and failed it minutes later.
	reg.EngineVersions = a.engines(parent)
	reg.EnginesProbed = true
	reg.InstallableEngines = provision.Installable()
	reg.Features = []string{FeatureFabricTest}
	reg.RunningShards = a.runningShards()
	a.mu.Lock()
	h := a.prov.Host
	// A new stream is a control plane that knows nothing of where installs
	// stand, so the next heartbeat tells it again.
	a.provSent = ""
	a.mu.Unlock()
	reg.Host = &llmgwv1.Host{
		OsId: h.OSID, OsVersion: h.OSVersion, OsCodename: h.OSCodename,
		Arch: h.Arch, CpuFlags: h.CPUFlags, NvidiaDriver: h.NVIDIADriver,
	}

	if err := stream.Send(&llmgwv1.AgentMessage{
		Msg: &llmgwv1.AgentMessage_Register{Register: reg},
	}); err != nil {
		return fmt.Errorf("registering: %w", err)
	}

	// Closing the response is what actually unblocks Receive below. Cancelling
	// the context is not enough: HTTP/2 watches it while the request body is
	// being written, and a bidirectional stream spends its life blocked on the
	// read side instead. Without this an agent told to stop keeps reading until
	// something kills it — which for systemd means the full stop timeout on
	// every restart.
	watchDone := make(chan struct{})
	defer close(watchDone)
	go func() {
		select {
		case <-ctx.Done():
			_ = stream.CloseResponse()
		case <-watchDone:
		}
	}()

	// Sends are serialized through this channel: connect streams are not safe
	// for concurrent writes, and heartbeats race with status reports.
	out := make(chan *llmgwv1.AgentMessage, 32)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		a.sendLoop(ctx, stream, out)
	}()
	defer wg.Wait()

	for {
		msg, err := stream.Receive()
		if err != nil {
			cancel()
			switch {
			case parent.Err() != nil:
				// The agent is stopping; the stream ending is the consequence.
				return nil
			case errors.Is(err, io.EOF):
				// The control plane closed the stream, which a restart does.
				// Reconnect immediately rather than backing off.
				return nil
			}
			// Everything else is worth both reporting and backing off from. A
			// refused certificate arrives here, and an agent that cannot say
			// why it is not registering is one an operator cannot fix.
			return err
		}
		a.handle(ctx, msg, out)
	}
}

type sender interface {
	Send(*llmgwv1.AgentMessage) error
}

func (a *Agent) sendLoop(ctx context.Context, s sender, out <-chan *llmgwv1.AgentMessage) {
	t := time.NewTicker(a.cfg.HeartbeatInterval)
	defer t.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case msg := <-out:
			if err := s.Send(msg); err != nil {
				return
			}
		case <-t.C:
			if err := s.Send(a.heartbeat()); err != nil {
				return
			}
		}
	}
}

func (a *Agent) heartbeat() *llmgwv1.AgentMessage {
	a.mu.Lock()
	inv := a.inv
	a.mu.Unlock()

	hb := &llmgwv1.Heartbeat{NodeId: a.cfg.NodeID}

	// Read every time; it is one small file. Memory is added to a running VM,
	// and a node that grew from 7 GB to 32 GB has to be seen to have grown
	// without anybody restarting anything.
	hb.MemoryTotalMb = int64(discovery.MemoryTotalMB())

	a.reprobeEngines()
	a.mu.Lock()
	if vers := a.engineReportLocked(); a.configuredVers != nil &&
		enginesFingerprint(vers) != a.enginesSent {
		a.enginesSent = enginesFingerprint(vers)
		hb.EnginesChanged = true
		hb.EngineVersions = vers
	}
	if fp := provisionFingerprint(a.provStatus); fp != a.provSent {
		a.provSent = fp
		hb.ProvisionChanged = true
		for _, e := range slices.Sorted(maps.Keys(a.provStatus)) {
			st := a.provStatus[e]
			hb.Provision = append(hb.Provision, &llmgwv1.ProvisionStatus{
				Engine: st.Engine, Version: st.Version, Variant: st.Variant,
				State: st.State, Detail: st.Detail,
			})
		}
	}
	a.mu.Unlock()

	for _, g := range inv.GPUs {
		hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB))
	}

	// Rescanned here rather than reused from registration. Disk is not a fact
	// about a machine, it is a fact about a moment: a node that had 200 GB when
	// its agent started can be full an hour later, and it is precisely the node
	// busy filling up whose figure matters.
	for _, st := range discovery.ScanStorage(a.cfg.ModelDirs) {
		hb.Storage = append(hb.Storage, &llmgwv1.Storage{
			Path: st.Path, FreeBytes: st.FreeBytes,
			TotalBytes: st.TotalBytes, Shared: st.Shared,
		})
	}

	// The file list is longer and changes rarely, so it goes only when it has.
	// The flag is what makes an empty list mean "this node now holds nothing"
	// rather than "nothing to report" — without it, the first heartbeat after
	// a model was deleted would be indistinguishable from silence.
	weights := discovery.ScanWeights(a.cfg.ModelDirs)
	fp := weightsFingerprint(weights)
	a.mu.Lock()
	changed := fp != a.lastWeights
	a.lastWeights = fp
	a.mu.Unlock()
	if changed {
		hb.WeightsChanged = true
		for _, w := range weights {
			hb.Weights = append(hb.Weights, &llmgwv1.LocalWeights{
				Path: w.Path, SizeBytes: w.SizeBytes, Shared: w.Shared,
			})
		}
	}
	hb.Load = a.sampleLoad()
	a.sampleIB(hb)
	a.sampleNVLink(hb)
	return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}}
}

// sampleLoad measures the machine and each running engine since the last
// heartbeat. Nil on the first, which has nothing to compare with.
func (a *Agent) sampleLoad() *llmgwv1.NodeLoad {
	if a.load == nil {
		return nil
	}
	pids := map[string]int{}
	for _, st := range a.sup.Statuses() {
		if st.PID > 0 {
			pids[st.ID] = st.PID
		}
	}
	l, ok := a.load.Sample(pids)
	if !ok {
		return nil
	}
	out := &llmgwv1.NodeLoad{
		CpuPercent: l.CPUPercent, CpuCount: int32(l.CPUCount),
		Load1: l.Load1, Load5: l.Load5, Load15: l.Load15,
		MemoryUsedBytes: l.MemUsed, MemoryAvailableBytes: l.MemAvail,
	}
	for _, sh := range l.Shards {
		out.Shards = append(out.Shards, &llmgwv1.ShardLoad{
			ShardId: sh.ID, CpuCores: sh.CPUCores, MemoryRssBytes: sh.RSSBytes,
		})
	}
	return out
}

// weightsFingerprint is enough to notice a model appearing, disappearing or
// finishing its download.
//
// Paths and sizes, because a fetch in progress has the same name as the file
// it will become — under a different name until it is renamed, but a shared
// filesystem can also grow a file in place.
func weightsFingerprint(ws []discovery.LocalWeights) string {
	h := sha256.New()
	for _, w := range ws {
		fmt.Fprintf(h, "%s:%d\n", w.Path, w.SizeBytes)
	}
	return hex.EncodeToString(h.Sum(nil))
}

// removeWeights deletes cached model files the control plane says are no
// longer needed.
//
// Every path is checked here as well as there, and that duplication is the
// point. The control plane computes these from the catalog, but it is the node
// that holds the files, and a node must not be talkable into deleting a model
// an operator spent a week copying onto it — not by a control plane that is
// wrong, not by one that is newer than this agent, not by one that has been
// tampered with. So this deletes only what it can prove llmgw itself wrote:
// inside its own cache directory, under a name this system generates.
//
// Failures are logged and skipped rather than reported upward. A file that is
// already gone is the expected outcome of a retry, and one that will not
// delete is a disk to look at, not a reason to fail a command whose purpose
// was to free space.
func (a *Agent) removeWeights(paths []string) {
	cache := a.cfg.WeightsCache
	if cache == "" && len(a.cfg.ModelDirs) > 0 {
		cache = a.cfg.ModelDirs[0]
	}
	if cache == "" {
		a.log.Warn("asked to remove weights with no cache directory configured")
		return
	}
	cacheAbs, err := filepath.Abs(cache)
	if err != nil {
		a.log.Warn("cannot resolve the cache directory", "dir", cache, "err", err)
		return
	}

	for _, p := range paths {
		abs, err := filepath.Abs(p)
		if err != nil {
			continue
		}
		// Directly inside the cache directory, not merely under it: a nested
		// path would let ".." or a symlinked subdirectory reach elsewhere, and
		// nothing this package writes is nested anyway.
		if filepath.Dir(abs) != cacheAbs {
			a.log.Warn("refused to remove a file outside the weights cache", "path", p)
			continue
		}
		if !weights.Owned(abs) {
			a.log.Warn("refused to remove a file llmgw did not fetch", "path", p)
			continue
		}
		if err := os.Remove(abs); err != nil && !os.IsNotExist(err) {
			a.log.Warn("could not remove cached weights", "path", p, "err", err)
			continue
		}
		a.log.Info("removed cached weights no model asks for", "path", p)
	}
}

func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) {
	switch m := msg.Msg.(type) {
	case *llmgwv1.ControlMessage_Start:
		go a.startShard(ctx, m.Start, out)
	case *llmgwv1.ControlMessage_Stop:
		go a.stopShard(ctx, m.Stop.GetShardId(), out)
	case *llmgwv1.ControlMessage_Drain:
		if err := a.sup.Drain(m.Drain.GetShardId()); err != nil {
			a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
		}
	case *llmgwv1.ControlMessage_RemoveWeights:
		a.removeWeights(m.RemoveWeights.GetPaths())
	case *llmgwv1.ControlMessage_Provision:
		a.setProvision(m.Provision)
	case *llmgwv1.ControlMessage_Logs:
		reply := a.logsReply(m.Logs)
		// Off the receive loop: a full send queue must not stop commands being
		// read.
		go func() {
			select {
			case out <- &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Logs{Logs: reply}}:
			case <-ctx.Done():
			}
		}()
	case *llmgwv1.ControlMessage_FabricScan:
		// A scan takes seconds to minutes: off the receive loop entirely.
		go func(req *llmgwv1.FabricScanRequest) {
			start := time.Now()
			reply := &llmgwv1.FabricScanReply{RequestId: req.GetRequestId()}
			text, err := scanFabric(ctx, req.GetHca(), int(req.GetPort()))
			if err != nil {
				reply.Error = err.Error()
			} else {
				reply.Output = text
			}
			reply.TookMs = time.Since(start).Milliseconds()
			select {
			case out <- &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_FabricScan{FabricScan: reply}}:
			case <-ctx.Done():
			}
		}(m.FabricScan)
	case *llmgwv1.ControlMessage_FabricTest:
		// Seconds of traffic, and a package install the first time: off the
		// receive loop too.
		go func(req *llmgwv1.FabricTestRequest) {
			reply := a.fabricTest(ctx, req)
			select {
			case out <- &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_FabricTest{FabricTest: reply}}:
			case <-ctx.Done():
			}
		}(m.FabricTest)
	case *llmgwv1.ControlMessage_Ack:
		if !m.Ack.GetOk() {
			a.log.Warn("control plane rejected a message",
				"id", m.Ack.GetMessageId(), "err", m.Ack.GetError())
		}
	}
}

// startShard launches a shard and reports its progress.
//
// It runs in its own goroutine because model loading takes minutes, and the
// stream has to keep carrying heartbeats and other commands meanwhile.
// watchFetch registers interest in a source's download progress and returns
// the function that stops watching.
//
// Keyed by source rather than by shard because a fetch is single-flight: eight
// shards starting together produce one download, and every one of them is
// waiting on it.
func (a *Agent) watchFetch(source string, fn func(weights.Progress)) (stop func()) {
	a.mu.Lock()
	if a.fetchWatchers == nil {
		a.fetchWatchers = map[string][]func(weights.Progress){}
	}
	a.fetchWatchers[source] = append(a.fetchWatchers[source], fn)
	i := len(a.fetchWatchers[source]) - 1
	a.mu.Unlock()

	var once sync.Once
	return func() {
		once.Do(func() {
			a.mu.Lock()
			defer a.mu.Unlock()
			w := a.fetchWatchers[source]
			if i < len(w) {
				// Cleared rather than removed, so the indexes other watchers
				// hold stay valid. The slice is dropped when the last one goes.
				w[i] = nil
			}
			for _, f := range w {
				if f != nil {
					return
				}
			}
			delete(a.fetchWatchers, source)
		})
	}
}

func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<- *llmgwv1.AgentMessage) {
	shardID, replicaID := m.GetShardId(), m.GetReplicaId()

	a.mu.Lock()
	a.shardReplica[shardID] = replicaID
	a.mu.Unlock()

	send := func(s *llmgwv1.ShardStatus) {
		s.NodeId, s.ReplicaId, s.ShardId = a.cfg.NodeID, replicaID, shardID
		select {
		case out <- &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: s}}:
		case <-ctx.Done():
		}
	}
	report := func(st supervisor.State, errMsg string, tail []string) {
		send(&llmgwv1.ShardStatus{
			State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail,
		})
	}

	spec, timeout := agentsvc.ShardSpecFromStart(m)
	a.localizeSpec(&spec)

	eng, err := a.reg.Get(spec.Profile.Engine)
	if err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}

	a.mu.Lock()
	inv := a.inv
	a.mu.Unlock()

	// Validate before launching. An engine that refuses a profile does so after
	// minutes of loading; catching it here costs nothing.
	if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}

	// The weights, before anything is launched. A source that is already a
	// path is checked and used; one that names a remote file is fetched once
	// and shared by every shard that wants it.
	//
	// Checked here rather than left to the engine because the engine's answer
	// arrives minutes into a load, in its own words, on a node nobody is
	// watching — and because eight shards starting together would otherwise
	// each fetch the same file.
	//
	// Progress is reported upward while it runs, so the control plane can tell
	// a download from a hang.
	stop := a.watchFetch(spec.Profile.ModelPath, func(p weights.Progress) {
		send(&llmgwv1.ShardStatus{
			State: agentsvc.StateToWire(supervisor.Starting),
			Fetch: &llmgwv1.FetchProgress{
				Source: p.Source, BytesDone: p.BytesDone, BytesTotal: p.BytesTotal,
			},
		})
	})
	if path, err := a.weights().Ensure(ctx, spec.Profile.ModelPath); err != nil {
		stop()
		report(supervisor.Failed, err.Error(), nil)
		return
	} else {
		stop()
		spec.Profile.ModelPath = path
	}

	cmds, err := eng.Commands(spec)
	if err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}

	sh, err := a.sup.Start(ctx, shardID, spec, cmds)
	if err != nil {
		if errors.Is(err, supervisor.ErrAlreadyRunning) {
			// Commands are idempotent: the control plane retries, and a repeat
			// must not produce a second process (plan §5.6).
			report(sh.Status().State, "", nil)
			return
		}
		report(supervisor.Failed, err.Error(), sh.Logs(10))
		return
	}
	report(supervisor.Starting, "", nil)

	// Only a rank that serves HTTP can be health-checked. A worker is judged by
	// staying alive, and the replica's real readiness is the leader answering,
	// which only happens once every worker has joined (plan §11, phase 4).
	if eng.ServesHTTP(spec) {
		base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort)
		err = a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout)
	} else {
		err = a.sup.WaitAlive(ctx, sh)
	}
	if err != nil {
		report(supervisor.Failed, err.Error(), sh.Logs(20))
		return
	}
	report(supervisor.Ready, "", nil)

	// Watch for the engine dying later. Nothing else would notice: the router
	// would simply see failing requests.
	go a.watch(ctx, sh, replicaID, shardID, out)
}

// watch reports a shard that exits after it was ready.
func (a *Agent) watch(ctx context.Context, sh *supervisor.Shard, replicaID, shardID string, out chan<- *llmgwv1.AgentMessage) {
	t := time.NewTicker(time.Second)
	defer t.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			st := sh.Status()
			if !st.State.Terminal() {
				continue
			}
			msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
				NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
				State: agentsvc.StateToWire(st.State), Error: st.Err,
				ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
			}}}
			select {
			case out <- msg:
			case <-ctx.Done():
			}
			return
		}
	}
}

func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgwv1.AgentMessage) {
	a.mu.Lock()
	replicaID := a.shardReplica[shardID]
	delete(a.shardReplica, shardID)
	a.mu.Unlock()

	if err := a.sup.Stop(ctx, shardID); err != nil {
		a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
	}
	a.sup.Remove(shardID)

	msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
		NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID,
		State: agentsvc.StateToWire(supervisor.Stopped),
	}}}
	select {
	case out <- msg:
	case <-ctx.Done():
	}
}

// gpuInfo narrows the inventory to the GPUs a shard will use.
func gpuInfo(inv discovery.Inventory, indices []int) []engines.GPUInfo {
	want := map[int]struct{}{}
	for _, i := range indices {
		want[i] = struct{}{}
	}
	var out []engines.GPUInfo
	for _, g := range inv.GPUs {
		if _, ok := want[g.Index]; !ok {
			continue
		}
		out = append(out, engines.GPUInfo{
			Index: g.Index, Name: g.Name,
			ComputeCapability: g.ComputeCapability, MemoryTotalMB: g.MemoryTotalMB,
		})
	}
	return out
}

// Client returns the client the agent reaches the control plane with.
//
// Both branches speak HTTP/2, which connect's bidirectional streams require and
// which Go will not do by default in either case: without TLS it has to be
// asked explicitly, and with TLS it has to be negotiated through ALPN.
//
// The plaintext branch exists for local runs against sim-engines. Anything
// using it must stay on a trusted network, because nothing there stops a node
// from claiming to be another one.
func Client(cfg tlsconf.Config, timeout time.Duration) (*http.Client, error) {
	if !cfg.Enabled() {
		return &http.Client{
			Transport: &http2.Transport{
				AllowHTTP: true,
				DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) {
					return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr)
				},
			},
		}, nil
	}

	// ServerName is left to the transport, which takes it from the address
	// being dialled. Hardcoding it here would break the moment the control
	// plane moved.
	tc, err := cfg.ClientConfig("")
	if err != nil {
		return nil, err
	}
	return &http.Client{Transport: &http2.Transport{TLSClientConfig: tc}}, nil
}

// 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)
	return out
}

// 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() {
		wants = append(wants, provision.Want{Engine: e.GetEngine(), Version: e.GetVersion()})
	}
	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()
}

// logsReply answers a request for one of this node's logs.
func (a *Agent) logsReply(req *llmgwv1.LogsRequest) *llmgwv1.LogsReply {
	n := int(req.GetLines())
	if n <= 0 {
		n = 200
	}
	n = min(n, agentLogLines)
	shards := a.sup.IDs()
	slices.Sort(shards)
	out := &llmgwv1.LogsReply{
		RequestId: req.GetRequestId(),
		Sources:   append([]string{"agent"}, shards...),
	}
	switch src := req.GetSource(); src {
	case "", "agent":
		out.Lines = a.logRing.tail(n)
	default:
		sh, err := a.sup.Get(src)
		if err != nil {
			out.Error = "no engine " + src + " runs on this node"
			return out
		}
		out.Lines = sh.Logs(n)
	}
	return out
}

// localizeSpec fills in what only this node knows.
//
// The interface NCCL bootstraps over is this node's, not the cluster's: the
// control plane sends one name for every node, and a node whose management
// network is on ens3 rather than eth0 would have NCCL bind nothing. And on
// RoCE, which GID entry routes is a property of these adapters.
func (a *Agent) localizeSpec(spec *engines.ShardSpec) {
	if a.cfg.MgmtInterface != "" {
		spec.MgmtInterface = a.cfg.MgmtInterface
	}
	if len(spec.HCAs) == 0 {
		return
	}
	a.mu.Lock()
	hcas := a.inv.HCAs
	a.mu.Unlock()
	want := map[string]bool{}
	for _, h := range spec.HCAs {
		want[h] = true
	}
	for _, h := range hcas {
		if !want[h.Name] {
			continue
		}
		for _, p := range h.Ports {
			if p.Active() && p.RoCEv2GID >= 0 {
				idx := p.RoCEv2GID
				spec.GIDIndex = &idx
				return
			}
		}
	}
}

// ResolveManagement is the address and interface this node's agent would use,
// found the way normalize finds them, for a tool that reports them.
func ResolveManagement(c Config) (addr, iface string) {
	addr, iface = c.MgmtAddr, c.MgmtInterface
	server := c.Server
	if server == "" && len(c.Servers) > 0 {
		server = c.Servers[0]
	}
	if addr == "" {
		addr = addrToward(server)
	}
	if iface == "" {
		iface = interfaceOf(addr)
	}
	return addr, iface
}

func smKey(hca string, port int) string { return fmt.Sprintf("%s/%d", hca, port) }

// scanFabric is discovery.ScanFabric, replaced in tests.
var scanFabric = discovery.ScanFabric

// querySMs is discovery.QuerySMs, replaced in tests.
var querySMs = func(hca string, port int) ([]discovery.SMRecord, error) {
	return discovery.QuerySMs(context.Background(), hca, port)
}

func splitSMKey(k string) (string, int) {
	i := strings.LastIndexByte(k, '/')
	port, _ := strconv.Atoi(k[i+1:])
	return k[:i], port
}

// querySM is discovery.QuerySM, replaced in tests.
var querySM = func(hca string, port int) (*discovery.SMInfo, error) {
	return discovery.QuerySM(context.Background(), hca, port)
}

// attachSM puts on each active InfiniBand port the subnet manager it sees,
// asking again once a minute and otherwise reusing the last answer.
func (a *Agent) attachSM(hcas []discovery.HCA) {
	a.mu.Lock()
	stale := time.Since(a.smAsked) > time.Minute || a.sm == nil
	a.mu.Unlock()
	if stale {
		sm, errs := map[string]*discovery.SMInfo{}, map[string]string{}
		_ = discovery.AttachSM(hcas, func(hca string, port int) (*discovery.SMInfo, error) {
			got, err := querySM(hca, port)
			if err != nil {
				errs[smKey(hca, port)] = err.Error()
				return nil, err
			}
			sm[smKey(hca, port)] = got
			return got, nil
		})
		// Every subnet manager on each fabric, through one port per master:
		// the same list through every port of one fabric is one query too
		// many for each.
		sms, smsErr := map[string][]discovery.SMRecord{}, map[string]string{}
		asked := map[string]bool{}
		for k, v := range sm {
			if asked[v.GUID] {
				continue
			}
			asked[v.GUID] = true
			hca, port := splitSMKey(k)
			if recs, err := querySMs(hca, port); err != nil {
				smsErr[k] = err.Error()
			} else {
				sms[k] = recs
			}
		}
		a.mu.Lock()
		for k, v := range sm {
			if old := a.sm[k]; old != nil && old.GUID != v.GUID {
				a.log.Warn("the master subnet manager changed", "port", k, "from", old.GUID, "to", v.GUID)
			}
		}
		a.sm, a.smErr, a.smAsked = sm, errs, time.Now()
		a.sms, a.smsErr = sms, smsErr
		a.mu.Unlock()
		return
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	for i := range hcas {
		for j := range hcas[i].Ports {
			if sm := a.sm[smKey(hcas[i].Name, hcas[i].Ports[j].Num)]; sm != nil && hcas[i].Ports[j].Active() {
				hcas[i].Ports[j].SM = sm
			}
		}
	}
}

// ensureDiags installs infiniband-diags on a node with adapters and without
// it. sminfo is what tells which physical fabric a port is on; without it,
// two separate fabrics given one name cannot be caught. Installed the way
// engine dependencies are, from the machine's own package sources.
func (a *Agent) ensureDiags(ctx context.Context) {
	a.mu.Lock()
	hasIB := len(a.inv.HCAs) > 0
	a.mu.Unlock()
	if !hasIB || a.prov == nil {
		return
	}
	if _, err := exec.LookPath("sminfo"); err == nil {
		return
	}
	a.log.Info("installing infiniband-diags, to tell which fabric each port is on")
	pkgs := map[string][]string{"deb": {"infiniband-diags"}, "rpm": {"infiniband-diags"}}
	if err := a.prov.InstallPackages(ctx, pkgs); err != nil {
		a.log.Warn("could not install infiniband-diags; fabrics cannot be told apart by their subnet manager",
			"err", err)
	}
}

// FeatureFabricTest is the agent answering FabricTestRequest.
const FeatureFabricTest = "fabric-test"

// runWriteBW is discovery.RunWriteBW, and haveWriteBW whether it is
// installed; both replaced in tests.
var (
	runWriteBW  = discovery.RunWriteBW
	haveWriteBW = func() bool { _, err := exec.LookPath("ib_write_bw"); return err == nil }
)

// fabricTest runs this node's side of a bandwidth test with another.
func (a *Agent) fabricTest(ctx context.Context, req *llmgwv1.FabricTestRequest) *llmgwv1.FabricTestReply {
	start := time.Now()
	reply := &llmgwv1.FabricTestReply{RequestId: req.GetRequestId()}
	defer func() { reply.TookMs = time.Since(start).Milliseconds() }()

	w := discovery.WriteBW{HCA: req.GetHca(), Port: int(req.GetPort()), GIDIndex: -1,
		TCPPort: int(req.GetTcpPort()), Seconds: int(req.GetSeconds())}
	if req.GetRole() == "client" {
		w.Peer = req.GetPeer()
	}
	a.mu.Lock()
	found := false
	for _, h := range a.inv.HCAs {
		for _, p := range h.Ports {
			if h.Name == w.HCA && p.Num == w.Port {
				found = true
				// RoCE needs its v2 GID named, as NCCL does; InfiniBand uses the LID.
				if p.LinkLayer == discovery.LinkEthernet && p.RoCEv2GID >= 0 {
					w.GIDIndex = p.RoCEv2GID
				}
			}
		}
	}
	a.mu.Unlock()
	if !found {
		reply.Error = fmt.Sprintf("this node has no port %s/%d", w.HCA, w.Port)
		return reply
	}
	if !haveWriteBW() && a.prov != nil {
		a.log.Info("installing perftest, for a bandwidth test with another node")
		pkgs := map[string][]string{"deb": {"perftest"}, "rpm": {"perftest"}}
		if err := a.prov.InstallPackages(ctx, pkgs); err != nil {
			reply.Error = "installing perftest: " + err.Error()
			return reply
		}
	}
	gbps, out, err := runWriteBW(ctx, w)
	reply.Gbps = gbps
	reply.Output = tailLines(out, 12)
	if err != nil {
		reply.Error = err.Error()
	}
	return reply
}

func tailLines(s string, n int) string {
	lines := strings.Split(strings.TrimRight(s, "\n"), "\n")
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return strings.Join(lines, "\n")
}

// discoverGPUFabric and discoverIMEX are replaced in tests.
var (
	discoverGPUFabric = discovery.DiscoverGPUFabric
	discoverIMEX      = discovery.DiscoverIMEX
)

// sampleNVLink reads, once a minute, which NVLink domain each GPU is joined
// to and whether IMEX is up. nvidia-smi -q takes a second or two, too long
// for every heartbeat, and a fabric that was joined an hour ago rarely
// changes in five seconds — but when a fabric manager restarts it does, for
// every card at once.
func (a *Agent) sampleNVLink(hb *llmgwv1.Heartbeat) {
	a.mu.Lock()
	gpus := slices.Clone(a.inv.GPUs)
	due := time.Since(a.nvlAsked) > time.Minute
	if due {
		a.nvlAsked = time.Now()
	}
	a.mu.Unlock()
	if len(gpus) == 0 || !due {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	fab, err := discoverGPUFabric(ctx, a.cfg.Discovery.NvidiaSMI)
	if err != nil || len(fab) == 0 {
		return
	}
	discovery.ApplyFabric(gpus, fab)
	for _, g := range gpus {
		if g.Fabric == nil {
			continue
		}
		hb.GpuFabric = append(hb.GpuFabric, &llmgwv1.GPUFabricState{
			Index: int32(g.Index), State: g.Fabric.State, Status: g.Fabric.Status,
			ClusterUuid: g.Fabric.ClusterUUID, CliqueId: g.Fabric.CliqueID, Bandwidth: g.Fabric.Bandwidth,
		})
	}
	m := discoverIMEX(ctx)
	hb.Imex = &llmgwv1.IMEX{State: m.State, Channels: int32(m.Channels)}
}