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"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"maps"
	"math"
	"math/rand/v2"
	"net"
	"net/http"
	"net/netip"
	"os"
	"path/filepath"
	"runtime"
	"slices"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
	"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/train"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/firewall"
	"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
	"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// Agent runs on a GPU node.
type Agent struct {
	// upgrading is true while an upgrade downloads.
	upgrading atomic.Bool
	cfg       Config
	// answered says the current session heard from its control plane, which
	// starts the reconnect backoff again when it ends.
	answered atomic.Bool
	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
	// serverList is every control plane to dial: the configured ones, and
	// any the leader has named since. Read through servers().
	serverList []string
	// fwMu keeps one firewall change at a time; fwRunner runs its scripts
	// (nil: this machine's shell).
	fwMu     sync.Mutex
	fwRunner firewall.Runner
	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)

	// kv mirrors the KV caches of the engines that publish events, for the
	// router's precise index; kvSubs are their listeners by shard, guarded
	// by mu.
	kv     *kvevents.Mirror
	kvSubs map[string]*kvevents.Subscriber

	// settingsCheck says whether the cluster's settings for this agent
	// would load beside its agent.yaml; see SetSettingsCheck.
	settingsCheck func([]byte) error

	// trainer runs LoRA training jobs, one at a time (BACKLOG F4).
	trainer *train.Runner
}

// 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(),
		serverList:   loadControlPlanes(cfg.Servers, cfg.TLS.Enabled()),
		kv:           kvevents.NewMirror(),
		kvSubs:       map[string]*kvevents.Subscriber{},
	}
	a.prov = &provision.Installer{Root: cfg.EnginesDir, Log: log, OnStatus: a.provisionStatus}
	runAs := cfg.Superv.RunAs
	if runAs == "" {
		runAs = supervisor.DefaultRunAs
	}
	// Where the engines are, as the installer has it when the file says
	// nothing.
	enginesDir := cfg.EnginesDir
	if enginesDir == "" {
		enginesDir = provision.DefaultRoot
	}
	a.trainer = train.New(train.Env{EnginesDir: enginesDir, WorkDir: filepath.Join(cfg.Superv.WorkDir, "trainer"),
		RunAs: runAs, Log: log})
	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 {
		// Read each time round: the leader may have named another since.
		servers := a.servers()
		server := servers[next%len(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++
		// A session that was answered was a connection that worked, however
		// it ended: the backoff starts again rather than carrying on from
		// the last outage. Without this every failover the agent lived
		// through made the next one slower — the third waited half a minute
		// on a control plane that had been ready for twenty seconds.
		if a.answered.Swap(false) {
			attempt = 0
		}
		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(servers) > 1 && next%len(servers) != 0 {
			delay = 0
		}
		a.log.Warn("disconnected from the control plane, retrying",
			"err", err, "server", server,
			"next", servers[next%len(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()
			// The engine is still publishing to the port the last agent
			// listened on; listen there again.
			if o.Spec.KVEventsPort > 0 {
				spec := o.Spec
				a.listenKV(&spec, o.Spec.KVEventsPort)
			}
			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 := a.engineBase(o.Spec)
		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
	}
	// What it computes on, for one started before the agent kept that — or
	// kept as the CPU for a GPU placement by an agent that did not follow
	// an engine into a process group of its own, which the check above has
	// just found on its GPUs.
	if c := a.sup.ComputeOf(o.ShardID); c == nil || (len(o.Spec.GPUIndices) > 0 && c.Backend == "CPU") {
		if c, ok := a.gpuCompute(ctx, o.Spec.GPUIndices, sh.Status().PID); ok {
			a.sup.SetCompute(o.ShardID, c)
		}
	}
	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.
// unready are the shards a registration reported as not yet ready.
func unready(running []*llmgwv1.RunningShard) map[string]bool {
	out := map[string]bool{}
	for _, r := range running {
		if r.GetState() != agentsvc.StateToWire(supervisor.Ready) {
			out[r.GetShardId()] = true
		}
	}
	return out
}

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,
			Compute:    computeToWire(a.sup.ComputeOf(st.ID)),
			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
}

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

// tracerName names the spans the agent opens.
const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/agent"

// 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)
	host.MemBandwidthGBs = a.memBandwidth()
	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,
		// What is measured by now, not what the startup scan found before
		// the measurement ended: a reconnect would otherwise say 0.
		MemBandwidthGbs: a.memBandwidth(),
	}

	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()
	go a.watchExits(ctx, out, unready(reg.RunningShards))
	go a.forwardKV(ctx, out)
	go a.reportCompute(ctx, out)

	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.answered.Store(true)
		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())
	hb.MemBandwidthGbs = a.memBandwidth()

	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()

	// Measured now, not carried from registration: a card is shared with
	// whatever else runs on the machine, and that comes and goes while the
	// agent runs. The registration figure stays only for a node where
	// nvidia-smi cannot be asked.
	if len(inv.GPUs) > 0 {
		var shards []int
		for _, st := range a.sup.Statuses() {
			if st.PID > 0 {
				shards = append(shards, st.PID)
			}
		}
		hb.GpuForeignMemoryMb, hb.GpuMemoryUsedMb = a.foreignGPUMemory(context.Background(), shards)
	}
	if len(hb.GpuMemoryUsedMb) != len(inv.GPUs) {
		hb.GpuForeignMemoryMb, hb.GpuMemoryUsedMb = nil, nil
		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
		}
		_ = os.Remove(weights.ChecksumPath(abs))
		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) {
	// The control plane's span that sent this, when it is tracing: what is
	// done here for a start or a stop goes into the same trace.
	tctx := obs.Extract(ctx, msg.GetTrace())
	switch m := msg.Msg.(type) {
	case *llmgwv1.ControlMessage_Start:
		go a.startShard(tctx, m.Start, out)
	case *llmgwv1.ControlMessage_Stop:
		go a.stopShard(tctx, m.Stop.GetShardId(), out)
	case *llmgwv1.ControlMessage_Drain:
		_, span := obs.Tracer(tracerName).Start(tctx, "drain shard",
			trace.WithSpanKind(trace.SpanKindConsumer),
			trace.WithAttributes(attribute.String("llmgw.shard", m.Drain.GetShardId())))
		err := a.sup.Drain(m.Drain.GetShardId())
		if err != nil {
			a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err)
		}
		obs.End(span, 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_ControlPlanes:
		a.setControlPlanes(m.ControlPlanes)
	case *llmgwv1.ControlMessage_Firewall:
		go a.applyFirewall(m.Firewall)
	case *llmgwv1.ControlMessage_AgentSettings:
		go a.applySettings(m.AgentSettings)
	case *llmgwv1.ControlMessage_TrainStart:
		go a.trainJob(ctx, m.TrainStart, out)
	case *llmgwv1.ControlMessage_TrainCancel:
		a.trainer.Cancel(m.TrainCancel.GetJobId())
	case *llmgwv1.ControlMessage_RestartAgent:
		go func() {
			a.log.Info("restarting, as the control plane asks; the engines keep serving",
				"reason", m.RestartAgent.GetReason())
			// A moment for the log line to leave before the process does.
			time.Sleep(300 * time.Millisecond)
			exitForUpgrade()
		}()
	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)
		})
	}
}

// bindAddr is where engines listen: the management address the control
// plane reaches them on, not every address the node has — a public one
// included. Empty, on every address, when that is a name rather than an
// address this machine holds.
func (a *Agent) bindAddr() string {
	ip, err := netip.ParseAddr(a.cfg.MgmtAddr)
	if err != nil || ip.IsUnspecified() || !localAddr(ip) {
		return ""
	}
	return ip.String()
}

// engineBase is where the agent asks an engine whether it is ready: where it
// listens, loopback when that is every address.
func (a *Agent) engineBase(s engines.ShardSpec) string {
	host := s.BindAddr
	if host == "" {
		host = "127.0.0.1"
	}
	return "http://" + net.JoinHostPort(host, strconv.Itoa(s.HTTPPort))
}

// localAddr reports whether an address is one of this machine's.
func localAddr(ip netip.Addr) bool {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return false
	}
	for _, a := range addrs {
		if p, err := netip.ParsePrefix(a.String()); err == nil && p.Addr().Unmap() == ip.Unmap() {
			return true
		}
	}
	return false
}

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

	// This node's whole part of the start, from the command arriving to the
	// engine answering, beneath the control plane's span that sent it. The
	// steps — weights, launch, readiness — are its children.
	ctx, span := obs.Tracer(tracerName).Start(ctx, "start shard",
		trace.WithSpanKind(trace.SpanKindConsumer),
		trace.WithAttributes(
			attribute.String("llmgw.node", a.cfg.NodeID),
			attribute.String("llmgw.replica", replicaID),
			attribute.String("llmgw.shard", shardID),
			attribute.Int("llmgw.rank", int(m.GetRank())),
			attribute.String("llmgw.model", m.GetProfile().GetServedModelName()),
			attribute.String("llmgw.profile", m.GetProfile().GetName()),
			attribute.String("llmgw.engine", m.GetProfile().GetEngine()),
			attribute.String("llmgw.gpu_class", m.GetProfile().GetGpuClass()),
			attribute.IntSlice("llmgw.gpus", int32Ints(m.GetGpuIndices())),
		))
	var failed error
	defer func() { obs.End(span, failed) }()

	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) {
		if st == supervisor.Failed {
			failed = errors.New(errMsg)
		}
		send(&llmgwv1.ShardStatus{
			State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail,
		})
	}

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

	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.
	// Checked again here, whatever sent the profile: this is the process
	// about to run as root.
	if err := engines.UnsafeArgs(spec.Profile.Engine, spec.Profile.Args); err != nil {
		report(supervisor.Failed, err.Error(), nil)
		return
	}
	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, Phase: p.Phase,
			},
		})
	})
	if path, err := a.weights().EnsureChecked(ctx, spec.Profile.ModelPath, spec.Profile.AllowPickle, spec.Profile.ModelSHA256); err != nil {
		stop()
		report(supervisor.Failed, err.Error(), nil)
		return
	} else {
		stop()
		spec.Profile.ModelPath = path
	}
	// A shared ComfyUI replica's members' weights, each fetched like the
	// model's, and handed to the gate where they landed.
	if raw := spec.Profile.Env[engines.SharedModelsEnv]; raw != "" {
		var ms []map[string]string
		if err := json.Unmarshal([]byte(raw), &ms); err != nil {
			report(supervisor.Failed, "the shared replica's models are not readable: "+err.Error(), nil)
			return
		}
		for i, m := range ms {
			stop := a.watchFetch(m["source"], func(p weights.Progress) {
				send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Starting),
					Fetch: &llmgwv1.FetchProgress{Source: p.Source, BytesDone: p.BytesDone, BytesTotal: p.BytesTotal,
						Phase: p.Phase}})
			})
			path, err := a.weights().Ensure(ctx, m["source"], m["allow_pickle"] == "true")
			stop()
			if err != nil {
				report(supervisor.Failed, fmt.Sprintf("model %s: %v", m["name"], err), nil)
				return
			}
			ms[i]["path"] = path
		}
		b, _ := json.Marshal(ms)
		spec.Profile.Env = maps.Clone(spec.Profile.Env)
		spec.Profile.Env[engines.SharedModelsEnv] = string(b)
	}
	// 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, spec.Profile.AllowPickle)
		if err != nil {
			report(supervisor.Failed, fmt.Sprintf("adapter %s: %v", ad.Name, err), nil)
			return
		}
		spec.Profile.Adapters[i].Path = path
	}
	// Adapters that load on demand are loaded by the router, by a path that
	// has to mean the same file on every node: a link of the adapter's name
	// in the directory engines run in.
	if spec.Profile.RuntimeLoRA() {
		if err := engines.LinkRuntimeAdapters(a.sup.EngineDir(), spec.Profile.Adapters); err != nil {
			report(supervisor.Failed, err.Error(), nil)
			return
		}
	}

	a.listenKV(&spec, 0)
	cmds, err := eng.Commands(spec)
	if err != nil {
		a.stopKV(shardID, replicaID)
		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
		}
		a.stopKV(shardID, replicaID)
		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 := a.engineBase(spec)
		err = a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout)
	} else {
		err = a.sup.WaitAlive(ctx, sh)
	}
	// Where it put the model, by its own account: said with the readiness,
	// and a GPU placement it computes on the CPU refused even where
	// nvidia-smi cannot be asked.
	compute, said, cerr := engineCompute(eng, spec.GPUIndices, sh.Logs(512))
	if err == nil && cerr != nil {
		err = cerr
		_ = a.sup.Stop(ctx, shardID)
	}
	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
	}
	// The engine's own account when it gave one, with its layers; else what
	// nvidia-smi sees the shard holding.
	if !said {
		compute, said = a.gpuCompute(ctx, spec.GPUIndices, sh.Status().PID)
	}
	if said {
		// Kept with the shard, so an agent restarted, or a control plane
		// taking over, still learns it.
		a.sup.SetCompute(shardID, compute)
		send(&llmgwv1.ShardStatus{State: agentsvc.StateToWire(supervisor.Ready), Compute: computeToWire(&compute)})
	} else {
		report(supervisor.Ready, "", nil)
	}
	// Its dying later is reported by the session's watchExits.
}

// watchExits reports, for as long as this session lasts, every shard that
// was ready and has failed: the ones this agent started, and the ones it
// adopted from the agent before it.
//
// Per session rather than per shard. A watcher tied to the stream a shard
// was started on stopped with that stream, so after a reconnect nothing
// reported the engine dying; and an adopted engine had no watcher at all —
// on node06 one was killed and its replica stayed "ready" for the control
// plane, routed to and failing, indefinitely.
//
// It also reports the shards the registration said were not ready yet once
// they are. An adopted engine is health-checked after the agent registers,
// and a shard started on the stream before is reported on that stream; so
// when the registration came first nothing ever told the control plane, and
// the replica stayed "starting" while it served — on node05, after a
// control plane restart, the only replica of a model.
func (a *Agent) watchExits(ctx context.Context, out chan<- *llmgwv1.AgentMessage, pending map[string]bool) {
	t := time.NewTicker(time.Second)
	defer t.Stop()
	ready := map[string]bool{}
	for {
		for _, st := range a.sup.Statuses() {
			switch st.State {
			case supervisor.Ready:
				ready[st.ID] = true
				if !pending[st.ID] {
					continue
				}
				delete(pending, st.ID)
				sh, err := a.sup.Get(st.ID)
				if err != nil {
					continue
				}
				msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
					NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID,
					State:   agentsvc.StateToWire(st.State),
					Compute: computeToWire(a.sup.ComputeOf(st.ID)),
				}}}
				select {
				case out <- msg:
				case <-ctx.Done():
					return
				}
			case supervisor.Failed:
				if !ready[st.ID] {
					continue // never ready here: its start reported it
				}
				delete(ready, st.ID)
				sh, err := a.sup.Get(st.ID)
				if err != nil {
					continue
				}
				msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
					NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID,
					State: agentsvc.StateToWire(st.State), Error: st.Err,
					ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20),
				}}}
				select {
				case out <- msg:
				case <-ctx.Done():
					return
				}
			default:
				delete(ready, st.ID)
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

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()

	ctx, span := obs.Tracer(tracerName).Start(ctx, "stop shard",
		trace.WithSpanKind(trace.SpanKindConsumer),
		trace.WithAttributes(
			attribute.String("llmgw.node", a.cfg.NodeID),
			attribute.String("llmgw.replica", replicaID),
			attribute.String("llmgw.shard", shardID),
		))
	err := a.sup.Stop(ctx, shardID)
	if err != nil {
		a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
	}
	a.stopKV(shardID, replicaID)
	obs.End(span, 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,
				ReadIdleTimeout: streamIdle,
				PingTimeout:     streamPing,
				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,
		ReadIdleTimeout: streamIdle,
		PingTimeout:     streamPing,
		// Bounded, as the plaintext branch's is: dialling a control plane
		// whose machine is off waits out the kernel's connect retries,
		// about two minutes, before the next address is tried.
		DialTLSContext: func(ctx context.Context, network, addr string, c *tls.Config) (net.Conn, error) {
			d := tls.Dialer{NetDialer: &net.Dialer{Timeout: timeout}, Config: c}
			return d.DialContext(ctx, network, addr)
		},
	}}, nil
}

// A stream to a control plane whose machine stopped dead — powered off, cut
// off — is never closed: nothing sends the FIN, and the agent's heartbeats sit
// in the socket's buffer until the kernel gives up retransmitting, fifteen
// minutes later. Found by switching off a leader: the standby took over in 14
// seconds and the node stayed with the dead one. An HTTP/2 ping after five
// quiet seconds, given three to answer, finds it in eight: by the time the
// standby leads, as it does within the lease's ten seconds, the node is on
// its way to it. At fifteen, with the lease at fifteen too, the new leader
// went without its nodes, and so its judges, for half a minute.
const (
	streamIdle = 5 * time.Second
	streamPing = 3 * time.Second
)

// 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. On RoCE,
// which GID entry routes is a property of these adapters. And a MIG instance
// is opened by its UUID, which only this node knows.
func (a *Agent) localizeSpec(spec *engines.ShardSpec) {
	if a.cfg.MgmtInterface != "" {
		spec.MgmtInterface = a.cfg.MgmtInterface
	}
	a.mu.Lock()
	hcas, gpus := a.inv.HCAs, a.inv.GPUs
	a.mu.Unlock()
	spec.GPUUUIDs = migUUIDs(gpus, spec.GPUIndices)
	if len(spec.HCAs) == 0 {
		return
	}
	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
			}
		}
	}
}

// migUUIDs names a shard's devices by UUID when one of them is a MIG
// instance, and is nil otherwise: whole cards keep being given by index, as
// they always were. A device this node does not have is left to fail where
// it always has, in the engine, rather than guessed at here.
func migUUIDs(gpus []discovery.GPU, indices []int) []string {
	var out []string
	mig := false
	for _, i := range indices {
		k := slices.IndexFunc(gpus, func(g discovery.GPU) bool { return g.Index == i })
		if k < 0 || gpus[k].UUID == "" {
			return nil
		}
		mig = mig || gpus[k].IsMIG()
		out = append(out, gpus[k].UUID)
	}
	if !mig {
		return nil
	}
	return out
}

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

// reportCompute tells the control plane, once a session, what each running
// shard computes on as soon as that is known: an adopted engine's is found
// after registration, which carried what was known then.
func (a *Agent) reportCompute(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
	// What was sent for each shard, so a compute that changes is sent again:
	// an adopted shard's kept one is sent at once, and adopting it may then
	// find it was kept wrong.
	sent := map[string]string{}
	t := time.NewTicker(5 * time.Second)
	defer t.Stop()
	for range 60 {
		for _, st := range a.sup.Statuses() {
			if st.State.Terminal() {
				continue
			}
			c := a.sup.ComputeOf(st.ID)
			if c == nil {
				continue
			}
			said := fmt.Sprint(*c)
			if sent[st.ID] == said {
				continue
			}
			sh, err := a.sup.Get(st.ID)
			if err != nil {
				continue
			}
			// No state: a report of compute alone, not a transition.
			msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{
				NodeId: a.cfg.NodeID, ReplicaId: sh.Spec.ReplicaID, ShardId: st.ID, Compute: computeToWire(c),
			}}}
			select {
			case out <- msg:
				sent[st.ID] = said
			case <-ctx.Done():
				return
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// memBandwidth is the node's memory read rate in GB/s, 0 until measured.
// Measured in the background, the first time it is asked for: it reads half
// a gigabyte, which must not hold up an agent reaching its control plane,
// and the heartbeat carries it when it is known. Not at once but after a
// minute or so, chosen at random, and again every hour, keeping the best:
// VMs on one host that all restart together measure each other, and read a
// third less than any of them gets alone. Once a process rather than an
// agent, so tests running many agents in one read it once.
func (a *Agent) memBandwidth() float64 {
	bwOnce.Do(func() {
		go func() {
			time.Sleep(bandwidthDelay())
			for {
				if bw := measureBandwidth(); bw > math.Float64frombits(bwBits.Load()) {
					bwBits.Store(math.Float64bits(bw))
					a.log.Info("this node's memory", "read_gbs", math.Round(bw*10)/10)
				}
				time.Sleep(time.Hour)
			}
		}()
	})
	return math.Float64frombits(bwBits.Load())
}

// bandwidthDelay is how long an agent waits before measuring, so agents
// started together measure apart.
var bandwidthDelay = func() time.Duration { return 10*time.Second + rand.N(110*time.Second) }

var (
	bwOnce sync.Once
	bwBits atomic.Uint64
	// measureBandwidth is replaced in tests, which have no use for it.
	measureBandwidth = discovery.MeasureMemoryBandwidth
)