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"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"log/slog"
	"os"
	"strings"
	"sync"
	"time"

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

// 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 {
	// engineChoices are the engine builds chosen per node; nil without a
	// store to keep them in.
	engineChoices *EngineChoices
	// facility is the holds, budgets and operator labels placement reads.
	facility *Facility
	cfg      Config
	log      *slog.Logger
	inv      *inventory.Store
	cat      *catalog.Live
	rec      *reconciler.Reconciler
	lead     *Leadership
	disp     reconciler.Dispatcher
	sink     EndpointSink
	usage    UsageSource
	peers    Peers

	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{
		facility: NewFacility(inv, log),
		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())
	// And what the replicas are doing, when the source can say: what adds
	// and removes replicas between a model's min and max.
	if l, ok := c.usage.(LoadSource); ok {
		c.rec.RecordLoad(l.ModelLoad())
	}
}

// Peers says which federated clusters' models may be sent requests now.
//
// Federation keeps it: whether each peer answered its last check, and
// whether it still offers the model. A peer model is in the catalog either
// way — a local model may name it as a fallback — and is published only
// while this says yes, which is what takes a peer that is down out of
// routing within a publish.
type Peers interface {
	// PeerRoute is the peer's own cluster id and whether requests sent to
	// it carry tags, and false when the model is not to be sent anything.
	PeerRoute(peer, upstreamModel string) (cluster string, sendTags, ok bool)
}

// SetPeers supplies the federation's view of its peers. Without one, no
// peer model is published.
func (c *Controller) SetPeers(p Peers) {
	c.mu.Lock()
	c.peers = p
	c.mu.Unlock()
}

// LoadSource reports each model's replicas' load, keyed by catalog name.
type LoadSource interface {
	ModelLoad() map[string]reconciler.Load
}

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

// ModelHealth is how each model in the catalog stands.
func (c *Controller) ModelHealth() []reconciler.ModelHealth { return c.rec.ModelHealth() }

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

// Leadership is what SetLeadership set; nil when there is nobody to compete
// with.
func (c *Controller) Leadership() *Leadership { return c.lead }

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

// Upgrader can tell a node's agent to upgrade itself.
type Upgrader interface {
	Upgrade(ctx context.Context, nodeID, version, releases string) error
}

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

// LogReader can ask a node for the recent lines of one of its logs.
type LogReader interface {
	Logs(ctx context.Context, nodeID, source string, lines int) (out, sources []string, err error)
}

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

// FabricTester returns the dispatcher's bandwidth test capability, or nil.
func (c *Controller) FabricTester() fabrictest.Asker {
	r, _ := c.disp.(fabrictest.Asker)
	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)
		info := endpoint.Info{
			ID:       rep.ID,
			URL:      strings.TrimRight(rep.Endpoint, "/"),
			Models:   names,
			APIKey:   c.engineKey(rep.ID),
			Engine:   prof.Engine,
			Pool:     rep.Pool,
			Weight:   rep.Weight,
			Priority: prof.Priority,
		}
		if m.Shared != nil {
			// One replica answering for every member of its environment:
			// offered under their names, each sent its own workflow and
			// checkpoint. Its own hidden name nobody asks for.
			info.Models = nil
			info.Comfy = map[string]endpoint.ComfyRoute{}
			ok := true
			for _, sm := range m.Shared.Members {
				w, err := c.workflow(sm.Workflow)
				if err != nil {
					c.log.Warn("a shared ComfyUI member's workflow could not be read; its replica is not offered",
						"model", sm.Name, "replica", rep.ID, "err", err)
					ok = false
					break
				}
				route := endpoint.ComfyRoute{Workflow: w, Checkpoint: sm.Checkpoint}
				for _, n := range append([]string{sm.Name}, sm.Aliases...) {
					info.Models = append(info.Models, n)
					info.Comfy[n] = route
				}
				if info.Workflow == "" {
					info.Workflow, info.UpstreamModel = w, sm.Checkpoint
				}
			}
			if ok {
				out = append(out, info)
			}
			continue
		}
		if prof.Engine == catalog.EngineComfyUI {
			workflow, ckpt := m.ComfyWorkflowFor(prof)
			w, err := c.workflow(workflow)
			if err != nil {
				c.log.Warn("a ComfyUI model's workflow could not be read; its replica is not offered",
					"model", m.Name, "replica", rep.ID, "err", err)
				continue
			}
			info.Workflow, info.UpstreamModel = w, ckpt
		}
		out = append(out, info)
	}
	return append(out, c.externalEndpoints()...)
}

// externalEndpoints are the models served outside the cluster, as endpoints:
// one each, with the provider's key read from its file every time, so a key
// rotated on disk is used from the next publish on.
func (c *Controller) externalEndpoints() []endpoint.Info {
	c.mu.Lock()
	peers := c.peers
	c.mu.Unlock()
	var out []endpoint.Info
	for _, m := range c.cat.Load().Models() {
		if !m.IsExternal() {
			continue
		}
		var cluster string
		var tags bool
		if m.IsPeer() {
			ok := false
			if peers != nil {
				cluster, tags, ok = peers.PeerRoute(m.External.Peer, m.External.Model)
			}
			if !ok {
				continue
			}
		}
		key := ""
		if f := m.External.APIKeyFile; f != "" {
			b, err := secrets.Read(context.Background(), f)
			if err != nil {
				c.log.Warn("an external model's key could not be read; not offered", "model", m.Name, "err", err)
				continue
			}
			key = strings.TrimSpace(string(b))
		}
		// A peer's endpoint says so in its id, which is what usage records
		// as the replica that answered.
		id := "external:" + m.Name
		if m.IsPeer() {
			id = "peer:" + m.Name
		}
		engine, workflow := endpoint.EngineExternal, ""
		if m.External.API == catalog.APIComfyUI {
			w, err := c.workflow(m.External.Workflow)
			if err != nil {
				c.log.Warn("a ComfyUI model's workflow could not be read; not offered", "model", m.Name, "err", err)
				continue
			}
			engine, workflow = endpoint.EngineComfyUI, w
		}
		out = append(out, endpoint.Info{
			ID: id, URL: strings.TrimRight(m.External.URL, "/"),
			Models: append([]string{m.Name}, m.Aliases...), APIKey: key,
			Engine: engine, Pool: "external", Weight: 1,
			UpstreamModel: m.External.Model,
			Peer:          m.External.Peer, PeerCluster: cluster, PeerTags: tags,
			Workflow: workflow,
		})
	}
	return out
}

// workflow is a ComfyUI workflow as the router is sent it: a built-in's
// name as it is, a file read here, on the control plane the catalog says it
// is on, and sent as the workflow itself.
func (c *Controller) workflow(w string) (string, error) {
	if !strings.ContainsAny(w, "/.") {
		return w, nil
	}
	b, err := os.ReadFile(w)
	return string(b), err
}

// 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|%s|%s|%s|%t|%x|%x;", i.ID, i.URL, i.Engine, i.Pool, i.Models, i.Weight,
			i.Priority, i.UpstreamModel, i.Peer, i.PeerCluster, i.PeerTags, sha256.Sum256([]byte(i.APIKey)),
			sha256.Sum256([]byte(i.Workflow)))
	}
	return b.String()
}

// EngineKey is the key a replica's engine demands, empty when engines run
// without one.
func (c *Controller) EngineKey(replicaID string) string { return c.engineKey(replicaID) }

func (c *Controller) engineKey(replicaID string) string {
	if c.cfg.Reconciler.EngineKey == nil {
		return ""
	}
	return c.cfg.Reconciler.EngineKey(replicaID)
}

// EngineKeys derives each replica's engine key from one secret.
//
// HMAC rather than a stored random key per replica: the router needs the key
// for every replica it sends to, including ones adopted after a restart that
// this process never started, and a second control plane has to agree. A
// function of a shared secret and the replica id gives all three without a
// table to keep in step. A key taken from one node opens that replica only.
func EngineKeys(secret []byte) func(replicaID string) string {
	return func(replicaID string) string {
		m := hmac.New(sha256.New, secret)
		m.Write([]byte("llmgw engine key\x00" + replicaID))
		return "llmgw-" + hex.EncodeToString(m.Sum(nil))[:40]
	}
}

// QuotaUse is every capacity quota beside what its scope holds.
func (c *Controller) QuotaUse() []reconciler.QuotaUse { return c.rec.QuotaUse() }

// SetEngineChoices gives the controller the engine builds chosen per node.
func (c *Controller) SetEngineChoices(e *EngineChoices) { c.engineChoices = e }

// Facility is where domains are held and nodes labelled; never nil.
func (c *Controller) Facility() *Facility { return c.facility }

// EngineChoices are the engine builds chosen per node, or nil.
func (c *Controller) EngineChoices() *EngineChoices { return c.engineChoices }

// ModelOwners maps each model to the part of the tenancy tree its capacity
// counts against; a model owned by nobody is not in it.
func (c *Controller) ModelOwners() map[string]string {
	out := map[string]string{}
	if cat := c.Catalog(); cat != nil {
		for _, m := range cat.Models() {
			if m.Owner != "" {
				out[m.Name] = m.Owner
			}
		}
	}
	return out
}