llmgw/internal/control/reconciler/demand.go

package reconciler

import (
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)

// RecordUsage stamps replicas with when each was last sent a request.
//
// Replicas absent from the map are left alone rather than zeroed: a router that
// has just restarted knows nothing about what happened before it, and treating
// that as "never used" would make every replica look abandoned at once —
// exactly when a fleet is least able to afford acting on that.
func (r *Reconciler) RecordUsage(lastUsed map[string]time.Time) {
	if len(lastUsed) == 0 {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	for id, at := range lastUsed {
		rep, ok := r.replicas[id]
		if !ok {
			continue
		}
		// Only forwards. Two routers reporting, or a stale snapshot arriving
		// late, must not move a replica's last use backwards.
		if at.After(rep.LastUsedAt) {
			rep.LastUsedAt = at
		}
	}
	r.noteUsage(lastUsed)
}

// IdleFor reports how long a replica has gone without a request, and whether
// that is known at all.
//
// The second return is not a detail. "Idle for an hour" and "nobody has told me
// yet" are different states, and anything that acts on idleness — eviction,
// most obviously — has to refuse to act on the second.
func (r *Reconciler) IdleFor(id string) (time.Duration, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	rep, ok := r.replicas[id]
	if !ok || rep.LastUsedAt.IsZero() {
		return 0, false
	}
	return r.now().Sub(rep.LastUsedAt), true
}

// Want marks an on-demand model as asked for, so the next pass places it.
//
// Called by whatever fields the request — the router, in practice. It returns
// whether the model is one this applies to, so a caller can tell "being loaded"
// from "there is genuinely nothing serving this", which are different answers
// to give a client.
func (r *Reconciler) Want(model string) bool {
	m, ok := r.cat.Load().Resolve(model)
	if !ok || !m.Scaling.OnDemand() {
		return false
	}
	r.mu.Lock()
	r.wanted[m.Name] = r.now()
	// A placement that failed before this request says nothing about now:
	// forgotten, so the pass this triggers decides afresh and Unservable
	// does not answer from a stale result.
	if res, ok := r.lastExplain[m.Name]; ok && res.Placement == nil {
		delete(r.lastExplain, m.Name)
	}
	r.mu.Unlock()
	// Immediately, not on the next tick. Someone is waiting on this request,
	// and five seconds of ticker is five seconds they spend staring at nothing.
	r.Notify()
	return true
}

// desiredFor is how many replicas a model should have right now.
//
// For an ordinary model that is its configured minimum. For an on-demand one it
// is one while somebody is asking and none once they stop, which is the whole
// feature: a cluster with more models than GPUs can hold every model in its
// catalog and only the wanted ones in its cards.
func (r *Reconciler) desiredFor(m catalog.Model) int {
	// Stopped by the operators: nothing runs, whatever the catalog asks.
	if r.stopped.Load() {
		return 0
	}
	base := r.baseFor(m)
	if base == 0 {
		return 0
	}
	r.mu.Lock()
	extra := r.extra[m.Name]
	r.mu.Unlock()
	return min(base+extra, max(r.maxFor(m), base))
}

// minFor is a model's minimum now: its schedule's window's, or its own
// outside every window.
func (r *Reconciler) minFor(m catalog.Model) int {
	if w, ok := m.Scaling.Window(r.localNow()); ok {
		return w.Min
	}
	return m.Scaling.Min
}

// maxFor is the ceiling load scaling reaches now.
func (r *Reconciler) maxFor(m catalog.Model) int {
	if w, ok := m.Scaling.Window(r.localNow()); ok && w.Max > 0 {
		return w.Max
	}
	return m.Scaling.Max
}

// localNow is now in the time zone schedules are written in.
func (r *Reconciler) localNow() time.Time {
	if r.cfg.Location != nil {
		return r.now().In(r.cfg.Location)
	}
	return r.now()
}

// baseFor is what a model runs without load: its minimum, or for an on-demand
// model one while somebody is asking and none once they stop. A schedule's
// window keeps its minimum running however quiet it is.
func (r *Reconciler) baseFor(m catalog.Model) int {
	if m.IsExternal() {
		return 0 // served elsewhere: nothing to place, whatever its scaling says
	}
	if lo := r.minFor(m); !m.Scaling.OnDemand() || lo > 0 {
		return lo
	}
	r.mu.Lock()
	last, asked := r.wanted[m.Name]
	r.mu.Unlock()
	if !asked {
		return 0
	}
	if r.now().Sub(last) > r.idleTimeout(m) {
		return 0
	}
	return 1
}

// idleTimeout is how long an on-demand model stays up unasked-for.
func (r *Reconciler) idleTimeout(m catalog.Model) time.Duration {
	if m.Scaling.IdleTimeout > 0 {
		return m.Scaling.IdleTimeout
	}
	return r.cfg.IdleTimeout
}

// noteUsage keeps an on-demand model alive while it is being used.
//
// Without this a model would be evicted exactly IdleTimeout after the request
// that loaded it, however busy it had been since — the timer would measure
// time since the first request rather than the last.
func (r *Reconciler) noteUsage(lastUsed map[string]time.Time) {
	byModel := map[string]time.Time{}
	for id, at := range lastUsed {
		rep, ok := r.replicas[id]
		if !ok {
			continue
		}
		if at.After(byModel[rep.Model]) {
			byModel[rep.Model] = at
		}
	}
	for model, at := range byModel {
		if cur, ok := r.wanted[model]; ok && at.After(cur) {
			r.wanted[model] = at
		}
	}
}

// Load is what a model's replicas are doing, as the router sees it.
type Load struct {
	Replicas int
	Running  int
	Waiting  int
}

// RecordLoad adds or takes away replicas between a model's min and max as its
// load says.
//
// Max was a ceiling nothing reached: every model ran exactly its minimum
// however long its queue. Now requests waiting at the engines — every slot of
// every replica busy — held for ScaleUpAfter add one replica, and a quiet spell
// of ScaleDownAfter takes one of those back. One step at a time, and not while
// the last one is still starting: on a small CPU fleet each replica is minutes
// of weights and gigabytes of memory, and a spike that is over by the time it
// loads should not leave three behind.
//
// With a benchmark stored for the profile a model runs on, its measured
// capacity is used too: more requests per replica than were measured to meet
// the SLO is busy, even with the engines' queues empty — an engine will take
// requests into its slots long after latency has gone — and a replica is only
// taken back when the rest would still be under scaleDownHeadroom of what they
// were measured to hold, so scaling down does not land on the edge.
func (r *Reconciler) RecordLoad(loads map[string]Load) {
	now := r.now()
	// The profile each model's replicas run on, for its measurement; read
	// before the loop, which holds the lock.
	running := map[string]string{}
	r.mu.Lock()
	for _, rep := range r.replicas {
		if rep.State.Counts() {
			running[rep.Model] = rep.Profile
		}
	}
	r.mu.Unlock()
	measured := map[string]Measured{}
	for model, profile := range running {
		if m, ok := r.measured(model, profile); ok {
			measured[model] = m
		}
	}
	for _, m := range r.cat.Load().Models() {
		base := r.baseFor(m)
		r.mu.Lock()
		ceiling := r.maxFor(m)
		if base == 0 || ceiling <= base {
			delete(r.extra, m.Name)
			delete(r.busySince, m.Name)
			delete(r.quietSince, m.Name)
			r.mu.Unlock()
			continue
		}
		starting := false
		for _, rep := range r.replicas {
			if rep.Model == m.Name && (rep.State == Placing || rep.State == Starting) {
				starting = true
			}
		}
		l := loads[m.Name]
		perReplica := float64(l.Waiting) / float64(max(l.Replicas, 1))
		extra := r.extra[m.Name]
		meas, hasMeas := measured[m.Name]
		overMeasured := hasMeas &&
			float64(l.Running+l.Waiting) > float64(meas.MaxConcurrency*max(l.Replicas, 1))
		fitsInFewer := !hasMeas ||
			float64(l.Running) <= scaleDownHeadroom*float64(meas.MaxConcurrency*max(l.Replicas-1, 0))
		switch {
		case perReplica >= r.cfg.ScaleUpWaiting || overMeasured:
			delete(r.quietSince, m.Name)
			since, ok := r.busySince[m.Name]
			if !ok {
				r.busySince[m.Name] = now
				break
			}
			if now.Sub(since) >= r.cfg.ScaleUpAfter && base+extra < ceiling && !starting {
				r.extra[m.Name] = extra + 1
				r.busySince[m.Name] = now
				r.log.Info("scaling up for load", "model", m.Name, "replicas", base+extra+1,
					"waiting_per_replica", perReplica, "in_flight", l.Running+l.Waiting,
					"measured_per_replica", meas.MaxConcurrency, "for", now.Sub(since).Round(time.Second))
			}
		case l.Waiting == 0 && extra > 0 && fitsInFewer:
			delete(r.busySince, m.Name)
			since, ok := r.quietSince[m.Name]
			if !ok {
				r.quietSince[m.Name] = now
				break
			}
			if now.Sub(since) >= r.cfg.ScaleDownAfter {
				r.extra[m.Name] = extra - 1
				r.quietSince[m.Name] = now
				r.log.Info("scaling down after a quiet spell", "model", m.Name,
					"replicas", base+extra-1, "quiet_for", now.Sub(since).Round(time.Second))
			}
		default:
			delete(r.busySince, m.Name)
			delete(r.quietSince, m.Name)
		}
		r.mu.Unlock()
	}
}

// scaleDownHeadroom is how full, of their measured capacity, the replicas left
// after a scale-down may be. Below one so that the load that was there a
// minute ago does not put them straight back over.
const scaleDownHeadroom = 0.7