llmgw/internal/control/reconciler/adopt.go

package reconciler

import (
	"slices"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
)

// Adopt rebuilds replica state from shards an agent reports it is already
// running (plan §6.5).
//
// This is what makes a control plane restart survivable. Without it the
// reconciler sees an empty cluster, places duplicates, and the engines already
// running keep holding GPUs the inventory believes are free — so the duplicates
// fail to allocate, with an error naming the new process rather than the old.
//
// Shards whose model or profile is no longer in the catalog are returned for
// the caller to stop: nothing will ever claim them.
func (r *Reconciler) Adopt(nodeID string, shards []AdoptedShard) (unclaimed []AdoptedShard) {
	defer r.forgetShardsTheAgentDoesNotHave(nodeID, shards)

	byReplica := map[string][]AdoptedShard{}
	for _, s := range shards {
		byReplica[s.ReplicaID] = append(byReplica[s.ReplicaID], s)
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	for replicaID, group := range byReplica {
		if existing, known := r.replicas[replicaID]; known {
			// The other half of a replica we are still assembling. Merge rather
			// than choose: these shards belong to the same replica as the ones
			// already adopted, they are simply on a different node whose agent
			// registered a moment later. Treating them as a competing claim is
			// what tore multi-node replicas apart on every control plane
			// restart — one half adopted and failed for being short, the other
			// stopped as unclaimed when it finally arrived.
			if !existing.AdoptedAt.IsZero() && len(existing.Shards) < existing.WantShards {
				r.mergeAdopted(existing, group)
				continue
			}
			// Usually this is an agent reconnecting rather than a control plane
			// restart, and its status reports keep the state current.
			//
			// Not always. Replica ids are derived from the model name and a
			// counter, so a control plane that restarts and places before the
			// agents reconnect produces a *different* replica under the same
			// id, on a freshly allocated port. Skipping the adoption then
			// leaves the cluster pointing at a port nothing listens on, while
			// the engine that is actually serving is treated as a stray.
			//
			// The fleet is authoritative about what is running (plan §6.5), so
			// what we merely placed loses to what an agent reports — unless our
			// copy has a shard that has actually come up, in which case two
			// things really are running and the one we did not ask for is the
			// stray.
			if existing.hasLiveShard() {
				// Only what we do not recognise. A reconnecting agent reports
				// the shards we already know, and disowning those would tell it
				// to stop the very replica that is serving.
				known := map[string]struct{}{}
				for _, sh := range existing.Shards {
					known[sh.ID] = struct{}{}
				}
				var strays []AdoptedShard
				for _, s := range group {
					if _, ok := known[s.ShardID]; !ok {
						strays = append(strays, s)
					}
				}
				unclaimed = append(unclaimed, strays...)
				r.reservePorts(strays)
				continue
			}
			r.log.Warn("a running replica reported by the fleet replaces one we placed",
				"replica", replicaID, "placed_endpoint", existing.Endpoint)
			for _, s := range existing.Shards {
				r.inv.Release(s.ID)
			}
			delete(r.replicas, replicaID)
		}

		first := group[0]
		m, ok := r.cat.Load().Resolve(first.Model)
		if !ok {
			unclaimed = append(unclaimed, group...)
			r.reservePorts(group)
			continue
		}
		prof, ok := m.ProfileByName(first.Profile)
		if !ok {
			unclaimed = append(unclaimed, group...)
			r.reservePorts(group)
			continue
		}

		rep := &Replica{
			ID: replicaID, Model: m.Name, Profile: prof.Name,
			Fabric: first.Fabric, Pool: first.Pool,
			// From the agent, which was told it when the shard started. Empty
			// from an agent too old to report it, which reads as unknown.
			ProfileFP: first.ProfileFP,
			Weight:    prof.Weight, CreatedAt: r.now(),
			Endpoint: first.Endpoint,
			// From the profile, not from what was reported: the point is to
			// notice when what was reported is short.
			WantShards: max(prof.Nodes, 1),
			AdoptedAt:  r.now(),
		}
		r.mergeAdopted(rep, group)
		r.replicas[replicaID] = rep
		// An on-demand model that is up was asked for, by somebody this
		// control plane has not heard of: it started before a restart or a
		// change of leader. Without this it reads as wanted by nobody and is
		// drained on the next pass — every awake on-demand model, at every
		// failover. It gets a whole idle timeout from now instead.
		if _, asked := r.wanted[m.Name]; !asked && m.Scaling.OnDemand() {
			r.wanted[m.Name] = r.now()
		}
	}
	return unclaimed
}

// mergeAdopted folds shards an agent reports into a replica being reassembled.
//
// Called both for the first group, which creates the replica, and for every
// later group belonging to it, so a replica spanning N nodes is built up over N
// registrations. Shard ids already present are skipped: an agent that
// reconnects repeats what it is running, and counting those again would carry a
// replica past its shard count without a second node ever appearing.
func (r *Reconciler) mergeAdopted(rep *Replica, group []AdoptedShard) {
	have := make(map[string]struct{}, len(rep.Shards))
	for _, sh := range rep.Shards {
		have[sh.ID] = struct{}{}
	}
	for _, s := range group {
		if _, dup := have[s.ShardID]; dup {
			// Known already, but what it computes on may not be.
			for i := range rep.Shards {
				if rep.Shards[i].ID == s.ShardID && rep.Shards[i].Compute == nil {
					rep.Shards[i].Compute = s.Compute
				}
			}
			continue
		}
		have[s.ShardID] = struct{}{}
		rep.Shards = append(rep.Shards, Shard{
			ID: s.ShardID, NodeID: s.NodeID, Rank: s.Rank,
			GPUIndices: s.GPUIndices, HCAs: s.HCAs,
			HTTPPort: s.HTTPPort, State: s.State, Compute: s.Compute,
		})
		// Reclaim the GPUs and the memory, or the next placement hands them
		// out twice. The memory comes from the catalog rather than the agent:
		// what a profile asked for is not something a running process reports.
		r.inv.Allocate(s.NodeID, inventory.Allocation{
			ShardID: s.ShardID, ReplicaID: rep.ID, GPUIndices: s.GPUIndices,
			MemoryGB: r.profileMemory(rep.Model, rep.Profile),
			CPUCores: r.adoptedCores(rep.Model, rep.Profile),
			PowerW:   r.profileWatts(rep.Model, rep.Profile),
		})
		// The endpoint comes from the leader, which may be in any of the
		// groups — the rank 0 node is not necessarily the first to register.
		if s.Rank == 0 && s.Endpoint != "" {
			rep.Endpoint = s.Endpoint
		}
	}
	slices.SortFunc(rep.Shards, func(a, b Shard) int { return a.Rank - b.Rank })
	rep.deriveState(r.now())

	r.log.Info("adopted replica from the fleet",
		"replica", rep.ID, "model", rep.Model,
		"shards", len(rep.Shards), "want", rep.WantShards, "state", rep.State)
}

// forgetShardsTheAgentDoesNotHave fails shards we believe are on a node that the
// agent, having just told us everything it is running, did not mention.
//
// The mirror image of adoption, and the half that was missing. Adoption handles
// "the agent has something we do not know about"; this handles "we think the
// agent has something it does not". Both follow from the same rule: the fleet is
// authoritative about what is running (plan §6.5).
//
// Without it, an engine that dies while its agent survives leaves the control
// plane believing a replica is Ready forever. Observed on a real cluster: the
// engine was killed, the agent restarted and registered with nothing running,
// and the replica stayed Ready pointing at a dead port for as long as anyone
// watched. Traffic survived — the breaker took the endpoint out — but the
// cluster ran a third short with no sign of it.
//
// Safe against the obvious race. A registration follows a fresh connection, so
// anything we dispatched over a previous one is not running by definition;
// there is no window where a shard is legitimately started but unreported.
func (r *Reconciler) forgetShardsTheAgentDoesNotHave(nodeID string, running []AdoptedShard) {
	if nodeID == "" {
		return
	}
	reported := make(map[string]struct{}, len(running))
	for _, s := range running {
		reported[s.ShardID] = struct{}{}
	}

	var lost []string
	r.mu.Lock()
	for _, rep := range r.replicas {
		if rep.State == Stopped {
			continue
		}
		changed := false
		for i := range rep.Shards {
			sh := &rep.Shards[i]
			if sh.NodeID != nodeID || sh.State.Terminal() {
				continue
			}
			if _, ok := reported[sh.ID]; ok {
				continue
			}
			sh.State = supervisor.Failed
			sh.Err = "the agent reports it is not running"
			r.inv.Release(sh.ID)
			changed = true
			lost = append(lost, sh.ID)
		}
		if changed {
			rep.deriveState(r.now())
		}
	}
	r.mu.Unlock()

	for _, id := range lost {
		r.log.Warn("a shard we believed was running is not; the agent does not have it",
			"node", nodeID, "shard", id)
	}
}

// reservePorts records ports held by shards nothing owns.
//
// They are being stopped, but stopping is not instant and the engine keeps its
// socket until it exits. Handing the port to a new shard in the meantime makes
// that shard fail to bind, which fails its replica, which places another one on
// the same port — a loop that does not end on its own.
//
// Called with the lock held.
func (r *Reconciler) reservePorts(shards []AdoptedShard) {
	for _, s := range shards {
		if s.HTTPPort > 0 {
			r.reservedPorts[s.HTTPPort] = struct{}{}
		}
	}
}

// AdoptedShard is one running shard an agent reported at registration.
type AdoptedShard struct {
	NodeID     string
	ReplicaID  string
	ShardID    string
	Rank       int
	GPUIndices []int
	HCAs       []string
	HTTPPort   int
	State      supervisor.State
	Model      string
	Profile    string
	Pool       string
	Fabric     string
	ProfileFP  string
	// Endpoint is the leader's base URL, set on the rank 0 shard.
	Endpoint string
	// Compute is where the engine said it put the model, when the agent
	// kept that.
	Compute *Compute
}

// adoptedCores is what an adopted shard is taken to hold of its node's CPUs.
func (r *Reconciler) adoptedCores(model, profile string) int {
	p, _ := r.profileOf(model, profile)
	return cpuCores(p)
}