llmgw/internal/control/reconciler/replica.go

// Package reconciler drives the cluster toward the catalog's desired state
// (plan §6.3).
package reconciler

import (
	"fmt"
	"slices"
	"time"

	"go.opentelemetry.io/otel/trace"

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

// ReplicaState is where a replica stands.
type ReplicaState int

const (
	// Placing means it has been scheduled but no shard has started.
	Placing ReplicaState = iota
	// Starting means shards are launching or loading.
	Starting
	// Ready means it serves traffic.
	Ready
	// Draining means it takes no new requests and will be stopped.
	Draining
	// Stopped means it was shut down on purpose.
	Stopped
	// Failed means it died or never became ready.
	Failed
)

func (s ReplicaState) String() string {
	switch s {
	case Placing:
		return "placing"
	case Starting:
		return "starting"
	case Ready:
		return "ready"
	case Draining:
		return "draining"
	case Stopped:
		return "stopped"
	case Failed:
		return "failed"
	default:
		return "unknown"
	}
}

// Terminal reports whether the replica has finished for good.
func (s ReplicaState) Terminal() bool { return s == Stopped || s == Failed }

// Counts reports whether the replica should count toward the desired total.
// Draining replicas do not: they are on the way out, and counting them would
// stop a replacement from being started.
func (s ReplicaState) Counts() bool { return s == Placing || s == Starting || s == Ready }

// Shard is one node's part of a replica.
type Shard struct {
	ID         string
	NodeID     string
	Rank       int
	GPUIndices []int
	HCAs       []string
	HTTPPort   int
	State      supervisor.State
	Err        string

	// Fetch is the download this shard is waiting on, if any. Cleared when the
	// shard reaches any other state, so a finished download does not sit in
	// the listing looking like one still running.
	Fetch *FetchProgress

	// Compute is where the engine said it put the model as it became ready:
	// nil when it did not say.
	Compute *Compute
}

// Compute is where an engine runs its model, by its own account.
type Compute struct {
	// Backend is CUDA, Vulkan, Metal, ROCm, GPU or CPU.
	Backend string
	// Devices it named: "CUDA0: Quadro RTX 6000 (7.5)".
	Devices []string
	// Offloaded of Total layers went to the devices; 0 and 0 when unsaid.
	Offloaded, Total int
}

// FetchProgress is how far a shard's weights have got.
//
// Carried so an operator can tell a replica pulling a hundred gigabytes from
// one that is wedged. From outside they are the same thing — a shard that has
// been starting for ten minutes — and the difference is whether to wait or to
// go and look.
type FetchProgress struct {
	Source    string
	BytesDone int64
	// BytesTotal is 0 when the server did not say. A percentage would then be
	// invented, and an invented one that stops moving is worse than a byte
	// count that keeps rising.
	BytesTotal int64
}

// Percent returns progress as a percentage, and false when the total is
// unknown.
func (f FetchProgress) Percent() (int, bool) {
	if f.BytesTotal <= 0 {
		return 0, false
	}
	return int(f.BytesDone * 100 / f.BytesTotal), true
}

// Leader reports whether this shard serves HTTP for the replica.
func (s Shard) Leader() bool { return s.Rank == 0 }

// Replica is one running copy of a model. It is the unit the router sees.
type Replica struct {
	ID      string
	Model   string
	Profile string
	// ProfileFP is the profile's contents when this replica was placed.
	//
	// Compared with the catalog's to tell a replica running the definition in
	// force from one running whatever it was edited from. A running replica is
	// deliberately not moved when its profile changes — churn costs more than
	// the imperfection — so without this there is nothing that says the two
	// have diverged, and "why is this one on the wrong node" is answered from
	// memory of what somebody edited an hour ago.
	//
	// Empty on a replica adopted from an agent that predates it, which reads
	// as unknown rather than as stale.
	ProfileFP string
	// Fabric is set only for multi-node replicas.
	Fabric string
	Pool   string
	Shards []Shard
	State  ReplicaState
	// Endpoint is rank 0's base URL.
	Endpoint string
	// Weight is the profile's capacity weight, passed to the router.
	Weight float64
	// WantShards is how many shards this replica should have, from its profile.
	//
	// Without it a replica cannot tell "every shard I have is ready" from
	// "every shard I have left is ready". A control plane that adopted half of
	// a two-node replica reported it Ready with one shard and no endpoint, and
	// the router was handed a replica nothing could route to.
	WantShards int
	// AdoptedAt is set while a replica is being reassembled from what agents
	// report, and cleared once every shard has arrived.
	//
	// A multi-node replica does not arrive at once: each node has its own
	// agent, each agent registers separately, and a replica spanning two nodes
	// is therefore reported in two halves a second or so apart. Between them it
	// is legitimately short of shards, and failing it there condemned it before
	// its other half had a chance to register — after which that half was
	// stopped as unclaimed and the model rebuilt from nothing on every restart.
	AdoptedAt time.Time

	CreatedAt time.Time
	ReadyAt   time.Time
	// LastUsedAt is when a request was last routed to this replica, zero if
	// none has been since the router started.
	//
	// The difference between a replica that is holding eight cards for work and
	// one that is holding them out of habit. Nothing recorded it, so on a fleet
	// with more models than GPUs there was no way to answer the first question
	// anyone asks: what can I stop?
	LastUsedAt time.Time
	// DrainingAt is when the replica stopped taking new work, so the wait for a
	// replacement can be bounded.
	DrainingAt time.Time
	// PreemptedFor is the model this replica was stopped to make room for.
	PreemptedFor string
	// RetiringAt is set while the replica is being restarted: it keeps serving
	// but no longer counts toward the model's replicas, so a replacement is
	// placed beside it, and it is drained only once that replacement is ready.
	RetiringAt time.Time
	// Attempt counts how many times this model has failed to come up in a row,
	// carried across rebuilds so the backoff does not reset on every attempt.
	Attempt int
	Err     string

	// startSpan is the "replica start" span, open from placement until the
	// replica serves, fails or is stopped; nil once closed, and for a replica
	// adopted rather than started. startTrace names it for good.
	startSpan  trace.Span
	startTrace trace.SpanContext
}

// StartTrace names the span of this replica's start, for something that
// waited on it — a job serving a model — to link to. Invalid for a replica
// adopted rather than started, and whenever tracing is off.
func (r Replica) StartTrace() trace.SpanContext { return r.startTrace }

// NodeIDs lists the nodes this replica occupies.
func (r Replica) NodeIDs() []string {
	out := make([]string, 0, len(r.Shards))
	for _, s := range r.Shards {
		out = append(out, s.NodeID)
	}
	slices.Sort(out)
	return out
}

// OnNode reports whether any shard runs on a node.
func (r Replica) OnNode(id string) bool {
	return slices.ContainsFunc(r.Shards, func(s Shard) bool { return s.NodeID == id })
}

// deriveState recomputes the replica's state from its shards.
//
// A replica is only as healthy as its worst shard: one failed shard makes the
// whole thing useless, because the others are waiting on collectives that will
// never complete. That is why there is no path here that keeps a replica
// partially alive (plan §5.4).
// hasLiveShard reports whether any shard has actually got somewhere.
//
// Starting and Pending are not evidence of anything: they are what a replica
// looks like both while an engine loads a model and when nothing was ever
// started at all. Loading means the port is open, which only a real process
// does.
func (r *Replica) hasLiveShard() bool {
	for _, s := range r.Shards {
		if s.State == supervisor.Ready || s.State == supervisor.Loading {
			return true
		}
	}
	return false
}

// deriveState recomputes the replica's state from its shards.
//
// Takes the time rather than reading the clock, because ReadyAt is now
// something decisions are made against — a replica is safe from eviction for a
// while after it starts serving — and a field stamped from the wall clock
// inside a package with an injectable one cannot be reasoned about in a test.
// It was invisible while nothing compared ReadyAt to anything.
func (r *Replica) deriveState(now time.Time) {
	if r.State == Draining || r.State == Stopped {
		return
	}
	if len(r.Shards) == 0 {
		r.State = Placing
		return
	}

	ready := 0
	for _, s := range r.Shards {
		switch s.State {
		case supervisor.Failed:
			r.State = Failed
			if r.Err == "" {
				r.Err = fmt.Sprintf("shard %s failed: %s", s.ID, s.Err)
			}
			return
		case supervisor.Stopped:
			// A shard stopping on its own, while the replica was not draining,
			// is a failure of the replica.
			r.State = Failed
			if r.Err == "" {
				r.Err = fmt.Sprintf("shard %s stopped unexpectedly", s.ID)
			}
			return
		case supervisor.Ready:
			ready++
		}
	}

	// Both conditions. A replica missing a shard is not ready however healthy
	// the shards it still has look — and rank 0's absence in particular means
	// there is no endpoint to route to at all.
	if r.WantShards > 0 && len(r.Shards) >= r.WantShards {
		// Fully assembled. From here it is an ordinary replica: losing a shard
		// now is a failure, not a half that has yet to report.
		r.AdoptedAt = time.Time{}
	}
	if ready == len(r.Shards) && (r.WantShards == 0 || len(r.Shards) == r.WantShards) {
		if r.State != Ready {
			r.ReadyAt = now
		}
		r.State = Ready
		return
	}
	if r.WantShards > 0 && len(r.Shards) < r.WantShards {
		// Still assembling from the fleet: the rest of the shards are on nodes
		// whose agents have not registered yet. Starting rather than Failed,
		// which puts it under the ordinary startup timeout — so a half that
		// never arrives is still failed, just by the deadline that already
		// exists rather than a second before its sibling reports in.
		if !r.AdoptedAt.IsZero() {
			r.State = Starting
			return
		}
		r.State = Failed
		if r.Err == "" {
			r.Err = fmt.Sprintf("replica has %d of %d shards", len(r.Shards), r.WantShards)
		}
		return
	}
	r.State = Starting
}