llmgw/internal/router/endpoint/endpoint.go

// Package endpoint tracks the upstream replicas the router can send to, along
// with their live state.
//
// In phase 1 the endpoint list comes from configuration. In phase 3 the control
// plane's EndpointSource (plan §6.4) will feed the same Registry; the scheduler
// and proxy must not notice the difference, which is why the API is shaped
// around wholesale list replacement from the start.
package endpoint

import (
	"maps"
	"slices"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
)

// Info is a replica's static description.
type Info struct {
	// ID uniquely identifies the replica and labels its metrics.
	ID string
	// URL is the engine's root address (e.g. http://gpu-a-01:8000), no trailing slash.
	URL string
	// APIKey is what the engine demands of every request, sent as a bearer
	// token. Empty for an engine that asks for none.
	APIKey string
	// Models lists the model names this replica serves.
	Models []string
	// Engine names the inference engine serving here, which decides what its
	// metrics are called. Empty means vLLM's names.
	Engine string
	// Pool names the (fabric, gpu_class) pool. Used for observability and pool
	// preference only; phase 1 may well have a single pool.
	Pool string
	// Weight is the capacity weight measured by ctl bench (§6.2.1). 0 means 1.
	Weight float64
	// Priority expresses pool preference; higher pools are tried first.
	Priority int
	// UpstreamModel is what the endpoint calls the model, when that differs
	// from the catalog's name: an external provider's. Requests are rewritten
	// to it. Empty sends the name as the client did.
	UpstreamModel string
	// Peer names the federated llmgw cluster behind an external endpoint,
	// empty for every other. PeerCluster is that cluster's own id, which a
	// request that has already passed through it carries in its hop header,
	// and PeerTags whether requests sent there name who sent them.
	Peer        string
	PeerCluster string
	PeerTags    bool
}

// EngineExternal marks an endpoint that is a provider outside the cluster.
// It publishes no metrics, so it is never scraped and never reads as stale.
const EngineExternal = "external"

// Serves reports whether the replica serves the given model.
func (i Info) Serves(model string) bool {
	return slices.Contains(i.Models, model)
}

// Stats is the router-side view of the engine metric contract (plan §7.6).
type Stats struct {
	// The three mandatory fields.
	Running      int
	Waiting      int
	KVCacheUsage float64

	// Optional.
	BlockSize    int
	NumGPUBlocks int

	// Raw counters, kept so derived rates can be computed.
	PromptTokens     float64
	GenerationTokens float64
	SuccessTotal     float64
}

// Snapshot is a complete view of one endpoint at scheduling time.
// It is a value type so the scheduler can work without holding a lock.
type Snapshot struct {
	Info
	Stats Stats

	// DrainRate is the observed completion rate in requests/sec, exponentially
	// smoothed. This is AIBrix's load-balance signal: Running alone cannot tell
	// that a replica is slow.
	DrainRate float64

	// Inflight counts requests the router has dispatched here that have not
	// finished. It compensates for scrape lag (§7.4); without it, herding is
	// unavoidable within a 2s scrape interval.
	Inflight int

	// RunningCapacity is the highest concurrent request count ever observed on
	// this endpoint — in practice the engine's batch capacity, learned rather
	// than configured. Without it, requests the router has dispatched but the
	// engine has not yet reported are indistinguishable from queue depth, and
	// during a scrape interval that makes a perfectly healthy engine look
	// overloaded.
	RunningCapacity int

	// Stale means no successful scrape arrived within StaleAfter. Stale
	// endpoints are not dropped, only scored down.
	Stale bool
	// Age is the time since the last successful scrape.
	Age time.Duration
	// Breaker is the circuit breaker's current state, read without side effects.
	Breaker breaker.State
}

// Config holds Registry settings.
type Config struct {
	// StaleAfter is how long an endpoint may go without a successful scrape
	// before being considered stale.
	StaleAfter time.Duration
	// Breaker holds per-endpoint circuit breaker settings.
	Breaker breaker.Config
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	return Config{
		StaleAfter: 10 * time.Second,
		Breaker:    breaker.DefaultConfig(),
	}
}

// live is one endpoint's mutable state.
type live struct {
	info Info

	stats      Stats
	scrapedAt  time.Time
	lastSucc   float64
	drainRate  float64
	maxRunning int

	inflight int
	brk      *breaker.Breaker
	// lastUsed is when a request was last dispatched here.
	//
	// The fact nothing recorded. Whether a replica is earning the cards it
	// holds is the first question on a fleet with more models than GPUs, and
	// the control plane had no way to ask it: it knew what was placed and
	// whether it was healthy, never whether anyone wanted it.
	//
	// Recorded here because Acquire is already the one place a request meets an
	// endpoint, and already holds the lock. One timestamp per dispatch, and
	// collection is one pass over the replicas rather than anything per
	// request — which is what makes it affordable on a fleet of hundreds.
	lastUsed time.Time
}

// Registry holds the live state of all endpoints. It is safe for concurrent use.
type Registry struct {
	cfg Config
	now func() time.Time

	mu sync.RWMutex
	m  map[string]*live
}

// NewRegistry creates an empty registry.
func NewRegistry(cfg Config) *Registry {
	if cfg.StaleAfter <= 0 {
		cfg.StaleAfter = 10 * time.Second
	}
	return &Registry{cfg: cfg, now: time.Now, m: make(map[string]*live)}
}

// SetClock replaces the time source, for tests.
func (r *Registry) SetClock(now func() time.Time) {
	r.mu.Lock()
	r.now = now
	r.mu.Unlock()
}

// Replace reconciles the endpoint list with the given set. Endpoints that
// remain keep their live state, new ones get a breaker, removed ones are dropped.
func (r *Registry) Replace(infos []Info) {
	r.mu.Lock()
	defer r.mu.Unlock()

	seen := make(map[string]struct{}, len(infos))
	for _, in := range infos {
		seen[in.ID] = struct{}{}
		if l, ok := r.m[in.ID]; ok {
			l.info = in
			continue
		}
		r.m[in.ID] = &live{info: in, brk: breaker.New(r.cfg.Breaker)}
	}
	for id := range r.m {
		if _, ok := seen[id]; !ok {
			delete(r.m, id)
		}
	}
}

// IDs returns the registered endpoint IDs in sorted order.
func (r *Registry) IDs() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return slices.Sorted(maps.Keys(r.m))
}

// Infos returns the registered endpoint descriptions, sorted by ID.
func (r *Registry) Infos() []Info {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]Info, 0, len(r.m))
	for _, l := range r.m {
		out = append(out, l.info)
	}
	slices.SortFunc(out, func(a, b Info) int { return strcmp(a.ID, b.ID) })
	return out
}

// Observe records a successful scrape and updates the drain rate.
func (r *Registry) Observe(id string, s Stats) {
	r.mu.Lock()
	defer r.mu.Unlock()

	l, ok := r.m[id]
	if !ok {
		return
	}
	now := r.now()
	if !l.scrapedAt.IsZero() {
		dt := now.Sub(l.scrapedAt).Seconds()
		// A counter going backwards means the engine restarted; skip the derivative.
		if dt > 0 && s.SuccessTotal >= l.lastSucc {
			rate := (s.SuccessTotal - l.lastSucc) / dt
			const alpha = 0.3 // exponential smoothing factor
			l.drainRate = alpha*rate + (1-alpha)*l.drainRate
		} else if s.SuccessTotal < l.lastSucc {
			l.drainRate = 0
		}
	}
	l.stats = s
	l.lastSucc = s.SuccessTotal
	l.scrapedAt = now
	l.maxRunning = max(l.maxRunning, s.Running)
}

// Snapshots returns a view of every endpoint serving the given model, or of all
// endpoints when model is empty. Breaker state is read, never consumed.
func (r *Registry) Snapshots(model string) []Snapshot {
	r.mu.RLock()
	defer r.mu.RUnlock()

	now := r.now()
	out := make([]Snapshot, 0, len(r.m))
	for _, l := range r.m {
		if model != "" && !l.info.Serves(model) {
			continue
		}
		var age time.Duration
		stale := l.info.Engine != EngineExternal
		if !l.scrapedAt.IsZero() {
			age = now.Sub(l.scrapedAt)
			stale = age > r.cfg.StaleAfter
		}
		out = append(out, Snapshot{
			Info:            l.info,
			Stats:           l.stats,
			DrainRate:       l.drainRate,
			Inflight:        l.inflight,
			RunningCapacity: l.maxRunning,
			Stale:           stale,
			Age:             age,
			Breaker:         l.brk.State(),
		})
	}
	slices.SortFunc(out, func(a, b Snapshot) int { return strcmp(a.ID, b.ID) })
	return out
}

// Acquire asks permission to send a request to an endpoint.
//
// It returns (nil, false) if the circuit breaker refuses. On success the
// inflight counter is incremented and the returned function must be called to
// report the outcome: release(err == nil).
func (r *Registry) Acquire(id string) (release func(ok bool), granted bool) {
	r.mu.Lock()
	l, exists := r.m[id]
	r.mu.Unlock()
	if !exists {
		return nil, false
	}
	if !l.brk.Allow() {
		return nil, false
	}

	r.mu.Lock()
	l.inflight++
	l.lastUsed = r.now()
	r.mu.Unlock()

	var once sync.Once
	return func(ok bool) {
		once.Do(func() {
			r.mu.Lock()
			l.inflight--
			r.mu.Unlock()
			if ok {
				l.brk.Success()
			} else {
				l.brk.Failure()
			}
		})
	}, true
}

// LastUsed reports when each endpoint last had a request dispatched to it.
//
// Endpoints that have never been used are absent rather than zero-valued: "not
// since this router started" and "not for an hour" are different facts, and a
// zero time would quietly become the oldest idle replica in any comparison —
// so a router restart would nominate everything for eviction at once.
func (r *Registry) LastUsed() map[string]time.Time {
	r.mu.RLock()
	defer r.mu.RUnlock()

	out := make(map[string]time.Time, len(r.m))
	for id, l := range r.m {
		if l.lastUsed.IsZero() {
			continue
		}
		out[id] = l.lastUsed
	}
	return out
}

// ScrapeFailed reports a failed metrics scrape.
//
// A scrape failure is not a request failure, so it does not touch the breaker:
// the endpoint may be serving fine while only its metrics endpoint is down.
// It will go stale on its own and score lower (§7.6).
func (r *Registry) ScrapeFailed(string) {}

// Inflight returns an endpoint's current inflight count, for tests and metrics.
func (r *Registry) Inflight(id string) int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	if l, ok := r.m[id]; ok {
		return l.inflight
	}
	return 0
}

func strcmp(a, b string) int {
	switch {
	case a < b:
		return -1
	case a > b:
		return 1
	default:
		return 0
	}
}