llmgw/internal/control/control.go

// Package control wires the control plane together: inventory, catalog,
// placement and reconciliation, and publishes the result to the router.
package control

import (
	"context"
	"fmt"
	"log/slog"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

// EndpointSink receives the endpoint list. The router implements it.
//
// It is an interface so the control plane does not depend on the router: when
// the two split into separate processes, this becomes a watch stream and
// nothing above it changes (plan §3.1).
type EndpointSink interface {
	SetEndpoints(ctx context.Context, infos []endpoint.Info)
}

// CatalogSink receives catalog changes. The router implements it.
//
// The catalog used to be fixed for the life of the process, so the router could
// derive its view once and keep it. Now that a model can be added while the
// cluster runs, the router has to be told — and it is told rather than asked,
// because a push happens once per change and a pull would happen once per
// request.
type CatalogSink interface {
	SetCatalog(c *catalog.Catalog)
}

// UsageSource reports when each endpoint was last sent a request.
//
// The return path of EndpointSink, and an interface for the same reason: the
// control plane decides what should run, the router is the only thing that
// knows whether anyone wants it, and neither should have to be in the same
// process as the other to say so.
//
// Deliberately a snapshot rather than an event stream. On a fleet with hundreds
// of replicas, per-request notification is a firehose carrying one bit of
// information — the newest timestamp wins — and a map read once a tick carries
// exactly the same answer for a fixed cost.
type UsageSource interface {
	LastUsed() map[string]time.Time
}

// Config tunes the control plane.
type Config struct {
	Reconciler reconciler.Config
	// PublishInterval is how often the endpoint list is republished even when
	// nothing appears to have changed. Events drive the common case; this is
	// the backstop.
	PublishInterval time.Duration
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	return Config{
		Reconciler:      reconciler.DefaultConfig(),
		PublishInterval: time.Second,
	}
}

// Controller owns the control plane's state.
type Controller struct {
	cfg   Config
	log   *slog.Logger
	inv   *inventory.Store
	cat   *catalog.Live
	rec   *reconciler.Reconciler
	lead  *Leadership
	disp  reconciler.Dispatcher
	sink  EndpointSink
	usage UsageSource

	mu   sync.Mutex
	last string // fingerprint of the last published list

	stop chan struct{}
	done chan struct{}
	once sync.Once
}

// New builds a controller.
func New(cfg Config, cat *catalog.Live, disp reconciler.Dispatcher, sink EndpointSink, log *slog.Logger) *Controller {
	if log == nil {
		log = slog.Default()
	}
	if cfg.PublishInterval <= 0 {
		cfg.PublishInterval = time.Second
	}
	inv := inventory.NewStore()
	return &Controller{
		cfg:  cfg,
		log:  log,
		inv:  inv,
		cat:  cat,
		rec:  reconciler.New(cfg.Reconciler, cat, inv, disp, log),
		disp: disp,
		sink: sink,
		stop: make(chan struct{}),
		done: make(chan struct{}),
	}
}

// Inventory exposes the node inventory, which the agent service feeds.
func (c *Controller) Inventory() *inventory.Store { return c.inv }

// SetUsageSource supplies where request activity is read from.
//
// A setter rather than a constructor argument because a control plane without
// one is fully functional — it simply cannot say whether a replica is idle,
// which is how it behaved until now.
func (c *Controller) SetUsageSource(u UsageSource) { c.usage = u }

// collectUsage stamps replicas with when they were last asked for anything.
//
// Once a tick, over the replicas rather than the requests. On a fleet of
// hundreds this is a map lookup each; there is no version of it that grows with
// traffic.
func (c *Controller) collectUsage() {
	if c.usage == nil {
		return
	}
	c.rec.RecordUsage(c.usage.LastUsed())
}

// ReplaceCatalog puts a new catalog in force, without a restart.
//
// The whole point of the change: adding a model was an operator procedure with
// downtime — edit a file on the control plane, restart, wait out the settling
// window, watch adoption churn. Nothing about that was necessary. The
// reconciler already compares desired against actual every pass, so a new
// catalog is just a different desired, and it is picked up on the next tick
// without anything being torn down.
//
// The catalog is validated by its own constructor, so an invalid one cannot
// reach here; what can reach here is a valid catalog that removes a model
// something is still serving, and that is handled the way it always has been —
// the reconciler stops replicas whose model is gone, and Publish drops
// endpoints for them first, so traffic stops being routed before anything is
// killed.
func (c *Controller) ReplaceCatalog(ctx context.Context, next *catalog.Catalog) {
	c.cat.Store(next)
	if cs, ok := c.sink.(CatalogSink); ok {
		cs.SetCatalog(next)
	}
	// Endpoints first, then wake the loop: a model that just left the catalog
	// should stop receiving traffic before its replicas are stopped, and
	// Publish is what takes it out of the router's list.
	c.Publish(ctx)
	c.rec.Notify()

	c.log.Info("catalog replaced", "models", len(next.Models()))
}

// Nodes, FreeGPUs and Replicas satisfy fleetmetrics.Source, so what the cluster
// is made of can be scraped from the one process that already knows.
func (c *Controller) Nodes() []inventory.Node        { return c.inv.Nodes() }
func (c *Controller) FreeGPUs(nodeID string) []int   { return c.inv.FreeGPUs(nodeID) }
func (c *Controller) Replicas() []reconciler.Replica { return c.rec.Replicas() }

// Reconciler exposes the reconciliation loop, for status reports and ctl.
func (c *Controller) Reconciler() *reconciler.Reconciler { return c.rec }

// Catalog exposes the model catalog.
func (c *Controller) Catalog() *catalog.Catalog { return c.cat.Load() }

// SetLeadership makes this control plane compete for the right to place work,
// instead of assuming it is the only one.
//
// Wired after construction because deciding leadership needs a store, and the
// store is opened by the router. Left unset, this process places
// unconditionally — which is what a single control plane wants, and what every
// deployment that existed before this had.
func (c *Controller) SetLeadership(l *Leadership) {
	c.lead = l
	c.rec.SetLeadership(l)
}

// IsLeader reports whether this process is placing work. True when leadership
// was never configured, because then there is nobody to compete with.
func (c *Controller) IsLeader() bool { return c.lead == nil || c.lead.IsLeader() }

// WeightRemover can ask a node to delete cached model files.
//
// Separate from Dispatcher because it is not part of scheduling and not every
// dispatcher has it — the ones tests use do not. A caller gets nil rather than
// a method that silently does nothing.
type WeightRemover interface {
	RemoveWeights(ctx context.Context, nodeID string, paths []string) error
}

// WeightRemover returns the dispatcher's deletion capability, or nil.
func (c *Controller) WeightRemover() WeightRemover {
	r, _ := c.disp.(WeightRemover)
	return r
}

// Start runs the reconciler and begins publishing endpoints.
func (c *Controller) Start(ctx context.Context) {
	if c.lead != nil {
		go c.lead.Run(ctx)
	}
	c.rec.Start(ctx)

	go func() {
		defer close(c.done)
		t := time.NewTicker(c.cfg.PublishInterval)
		defer t.Stop()

		c.Publish(ctx)
		for {
			select {
			case <-c.stop:
				return
			case <-ctx.Done():
				return
			case <-t.C:
				// Before publishing, so a replica's idleness is as fresh as
				// the endpoint list it travels beside.
				c.collectUsage()
				c.Publish(ctx)
			}
		}
	}()
}

// Close stops the control plane.
func (c *Controller) Close() {
	c.once.Do(func() {
		close(c.stop)
		<-c.done
		c.rec.Close()
	})
}

// Publish pushes the current Ready replicas to the sink.
//
// It is exported so tests can drive it deterministically instead of racing a
// ticker, and it skips the push when nothing changed: replacing the list is
// cheap but resynchronizes the metric pollers, and doing that every second for
// no reason adds churn to every engine's /metrics.
func (c *Controller) Publish(ctx context.Context) {
	infos := c.endpoints()

	fp := fingerprint(infos)
	c.mu.Lock()
	unchanged := fp == c.last
	c.last = fp
	c.mu.Unlock()
	if unchanged {
		return
	}

	c.log.Info("publishing endpoints", "count", len(infos))
	c.sink.SetEndpoints(ctx, infos)
}

// endpoints converts Ready replicas into the router's view.
func (c *Controller) endpoints() []endpoint.Info {
	ready := c.rec.ReadyReplicas()
	out := make([]endpoint.Info, 0, len(ready))

	for _, rep := range ready {
		m, ok := c.cat.Load().Resolve(rep.Model)
		if !ok {
			// The model left the catalog while a replica was still running.
			// Withdraw it here first, so traffic stops before the reconciler
			// stops the engine on its next pass.
			continue
		}
		// Clients may address a model by any of its names, and the router
		// matches endpoints by the resolved name plus aliases.
		names := append([]string{m.Name}, m.Aliases...)

		prof, _ := m.ProfileByName(rep.Profile)
		out = append(out, endpoint.Info{
			ID:       rep.ID,
			URL:      strings.TrimRight(rep.Endpoint, "/"),
			Models:   names,
			Engine:   prof.Engine,
			Pool:     rep.Pool,
			Weight:   rep.Weight,
			Priority: prof.Priority,
		})
	}
	return out
}

// fingerprint summarizes a list so an unchanged one can be skipped.
func fingerprint(infos []endpoint.Info) string {
	var b strings.Builder
	for _, i := range infos {
		fmt.Fprintf(&b, "%s|%s|%s|%s|%v|%g|%d;", i.ID, i.URL, i.Engine, i.Pool, i.Models, i.Weight, i.Priority)
	}
	return b.String()
}