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"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"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/tlsconf"
)

// 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
	// sensors reads heat and power: CPU and board from hwmon and RAPL, GPUs
	// from nvidia-smi.
	sensors *discovery.SensorSampler
	// wear reads what wears a GPU out: remapped rows, PCIe retries, XIDs.
	wear *discovery.HealthSampler

	// 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(),
		sensors:      discovery.NewSensorSampler(),
		wear:         discovery.NewHealthSampler(),
	}
	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
	}
	// An engine started by an agent older than the GPU check may be one
	// serving a GPU placement on the CPU. Stopped, so its replacement starts
	// through the check and fails there in words, if the build still sees
	// no GPU — rather than going on unnoticed after the upgrade that could
	// have caught it.
	if err := a.onItsGPUs(ctx, eng, o.Spec, sh.Status().PID); err != nil {
		a.log.Error("stopping an adopted engine that computes on the CPU though placed on GPUs",
			"shard", o.ShardID, "replica", o.ReplicaID, "err", err)
		_ = a.sup.Stop(ctx, o.ShardID)
		return
	}
	a.log.Info("adopted engine is serving",
		"shard", o.ShardID, "replica", o.ReplicaID)
}

// 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 = slices.DeleteFunc(provision.Installable(), a.disabled)
	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())
	hb.CpuCount = int32(runtime.NumCPU())

	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,
		})
	}
	if a.sensors != nil {
		se := a.sensors.Sample(context.Background())
		out.CpuTempC, out.CpuPowerW, out.InletTempC = se.CPUTempC, se.CPUPowerW, se.InletTempC
		var wear map[int]discovery.GPUWear
		if a.wear != nil && len(se.GPUs) > 0 {
			wear = map[int]discovery.GPUWear{}
			for _, w := range a.wear.Sample(context.Background()) {
				wear[w.Index] = w
			}
		}
		for _, g := range se.GPUs {
			gs := &llmgwv1.GPUSensor{
				Index: int32(g.Index), TempC: g.TempC, MemoryTempC: g.MemoryTempC,
				PowerW: g.PowerW, PowerLimitW: g.PowerLimitW, Throttle: g.Throttle,
				UtilizationPercent: g.UtilPercent, MemoryUtilizationPercent: g.MemUtilPercent,
				MemoryUsedMb: g.MemoryUsedMB, MemoryTotalMb: g.MemoryTotalMB,
				SmClockMhz: g.SMClockMHz, EccUncorrected: g.ECCUncorrected,
			}
			if w, ok := wear[g.Index]; ok {
				gs.HasWear = true
				gs.PcieReplays, gs.RetiredPagesSbe, gs.RetiredPagesDbe, gs.RetirementPending =
					w.PCIeReplays, w.RetiredSBE, w.RetiredDBE, w.RetirePending
				gs.RemappedRowsCorrectable, gs.RemappedRowsUncorrectable, gs.RemapPending, gs.RemapFailure =
					w.RemappedCorrectable, w.RemappedUncorrectable, w.RemapPending, w.RemapFailure
				gs.NvlinkCrcErrors, gs.NvlinkReplayErrors, gs.NvlinkRecoveryErrors = w.NVLinkCRC, w.NVLinkReplay, w.NVLinkRecovery
				gs.XidErrors, gs.LastXid, gs.LastXidText = w.XIDs, int32(w.LastXID), w.LastXIDText
				if !w.LastXIDAt.IsZero() {
					gs.LastXidUnix = w.LastXIDAt.Unix()
				}
			}
			out.Gpus = append(out.Gpus, gs)
		}
	}
	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_Upgrade:
		go a.upgrade(ctx, m.Upgrade, out)
	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
	}
	// Adapters are fetched like the model, and before the engine starts: it
	// loads them at start, and one missing is a replica that answers for the
	// base model under the adapter's name.
	for i, ad := range spec.Profile.Adapters {
		path, err := a.weights().Ensure(ctx, ad.Path)
		if err != nil {
			report(supervisor.Failed, fmt.Sprintf("adapter %s: %v", ad.Name, err), nil)
			return
		}
		spec.Profile.Adapters[i].Path = 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 {
		err = a.onItsGPUs(ctx, eng, spec, sh.Status().PID)
		if err != nil {
			// Serving, but not where it was placed: stopped, so the GPU it
			// holds in the inventory is not taken by a CPU process.
			_ = a.sup.Stop(ctx, shardID)
		}
	}
	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
}

// 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
}