llmgw/internal/control/routing.go

package control

import (
	"strconv"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
)

// RoutedReplica is a Ready replica as a router needs it: where it is, what it
// serves and how much it may take. Not its key — every control plane derives
// that from the replica's id and the cluster's secret, so none travels.
type RoutedReplica struct {
	ID       string  `json:"id"`
	Model    string  `json:"model"`
	Profile  string  `json:"profile"`
	Pool     string  `json:"pool"`
	Endpoint string  `json:"endpoint"`
	Weight   float64 `json:"weight"`
}

// Routing is what a standby that answers requests itself reads from the
// leader: the replicas it may send to, and which catalog they belong to.
type Routing struct {
	Replicas []RoutedReplica `json:"replicas"`
	// Catalog changes whenever the leader puts a catalog in force. Only
	// compared with the last one a standby saw, never with its own: a model
	// added, scaled or given a canary is in the store, and the standby reads
	// it again from there.
	Catalog string `json:"catalog"`
	// Rates are what every other control plane's router has counted against
	// each limit, the leader's among them: added to the asker's own, they
	// make a key's limits the cluster's.
	Rates *accounting.Rates `json:"rates,omitempty"`
}

// Activity is what a standby's router reports back: when it last sent each
// replica a request, and the on-demand models it was asked for. The leader
// evicts and loads by these, and without them a replica busy behind a
// standby looks idle to the one machine that decides.
type Activity struct {
	// From names the standby, so its counts replace its own last ones and
	// are not given back to it.
	From     string               `json:"from,omitempty"`
	LastUsed map[string]time.Time `json:"last_used,omitempty"`
	Want     []string             `json:"want,omitempty"`
	// Rates are its router's own counts against each limit.
	Rates *accounting.Rates `json:"rates,omitempty"`
}

// RateSharer is a router whose limits can be counted with other routers'.
type RateSharer interface {
	Rates() accounting.Rates
	SetOtherRates(accounting.Rates)
}

// ratesFresh is how long a standby's counts are believed after it last
// reported: a standby gone quiet stops counting against anyone's limits
// shortly after it stops taking requests.
const ratesFresh = 10 * time.Second

// sharedRates are the standbys' last counts, by name.
type sharedRates struct {
	mu   sync.Mutex
	by   map[string]rateReport
	seen bool // whether any standby has ever reported
}

type rateReport struct {
	rates accounting.Rates
	at    time.Time
}

// put keeps a standby's latest counts.
func (s *sharedRates) put(from string, r accounting.Rates, now time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.by == nil {
		s.by = map[string]rateReport{}
	}
	s.by[from] = rateReport{rates: r, at: now}
	s.seen = true
}

// sum adds every fresh standby's counts but except's; stale ones are
// forgotten. ok is false when no standby has ever reported, and there is
// nothing to put in force.
func (s *sharedRates) sum(except string, now time.Time) (accounting.Rates, bool) {
	out := accounting.NewRates(now)
	s.mu.Lock()
	defer s.mu.Unlock()
	for name, r := range s.by {
		if now.Sub(r.at) > ratesFresh {
			delete(s.by, name)
			continue
		}
		if name != except {
			out.Add(r.rates)
		}
	}
	return out, s.seen
}

// ReplicaFeed is a standby's copy of the leader's Ready replicas. ok is false
// when it has none worth using — standbys do not answer requests, or the
// leader has not been heard from lately — and the standby routes nothing.
type ReplicaFeed interface {
	Ready() (reps []RoutedReplica, ok bool)
}

// SetReplicaFeed lets a standby route to the leader's replicas.
func (c *Controller) SetReplicaFeed(f ReplicaFeed) { c.feed = f }

// Routing is this control plane's Ready replicas, for a standby to route to.
func (c *Controller) Routing() Routing { return c.RoutingFor("") }

// RoutingFor is Routing for the standby named from, with what every other
// router has counted against each limit.
func (c *Controller) RoutingFor(from string) Routing {
	out := Routing{Replicas: c.ready(), Catalog: c.CatalogRev()}
	if rs, ok := c.sink.(RateSharer); ok {
		others, _ := c.rates.sum(from, time.Now())
		others.Add(rs.Rates())
		out.Rates = &others
	}
	return out
}

// ready is the reconciler's Ready replicas, as a router needs them.
func (c *Controller) ready() []RoutedReplica {
	reps := c.rec.ReadyReplicas()
	out := make([]RoutedReplica, 0, len(reps))
	for _, r := range reps {
		out = append(out, RoutedReplica{
			ID: r.ID, Model: r.Model, Profile: r.Profile, Pool: r.Pool, Endpoint: r.Endpoint, Weight: r.Weight,
		})
	}
	return out
}

// shareRates puts the standbys' counts in force on this control plane's own
// router, while it leads. Once a tick, beside publishing.
func (c *Controller) shareRates() {
	rs, ok := c.sink.(RateSharer)
	if !ok || (c.lead != nil && !c.lead.IsLeader()) {
		return
	}
	if others, seen := c.rates.sum("", time.Now()); seen {
		rs.SetOtherRates(others)
	}
}

// CatalogRev names the catalog in force here. The process's owner id with a
// counter: a new leader, or a restarted one, is a different catalog as far as
// any standby knows.
func (c *Controller) CatalogRev() string {
	owner := ""
	if c.lead != nil {
		owner = c.lead.Owner()
	}
	return owner + "/" + strconv.FormatUint(c.catRev.Load(), 10)
}

// RecordActivity takes in what a standby's router did since its last report.
func (c *Controller) RecordActivity(a Activity) {
	if a.From != "" && a.Rates != nil {
		c.rates.put(a.From, *a.Rates, time.Now())
	}
	c.rec.RecordUsage(a.LastUsed)
	for _, m := range a.Want {
		c.rec.Want(m)
	}
}

// routable is what this process's router may send to: the replicas the
// reconciler holds Ready, or, on a standby following the leader, the leader's.
func (c *Controller) routable() []RoutedReplica {
	if c.feed != nil && c.lead != nil && !c.lead.IsLeader() {
		if reps, ok := c.feed.Ready(); ok {
			return reps
		}
		return nil
	}
	return c.ready()
}