llmgw/internal/control/reconciler/preempt.go

package reconciler

import (
	"context"
	"fmt"
	"slices"
	"sort"
	"time"

	"go.opentelemetry.io/otel/attribute"

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

// hold keeps a model that gave up room from taking it back.
type hold struct {
	Until time.Time
	For   string
}

// maxVictims bounds how many replicas one placement may stop. Past this the
// model is better reported as not fitting than paid for with half a cluster.
const maxVictims = 3

// victim is a replica that could be stopped for another model.
type victim struct {
	ID       string
	Model    string
	Class    string
	AboveMin bool
	Tier     int
	Idle     time.Duration
	// Size is what stopping it frees, for choosing the smaller of two
	// otherwise equal victims.
	Size quota.Amount
}

// preemptFor stops replicas of lower classes when a model cannot otherwise be
// placed, and says whether it did.
//
// What may be taken from whom:
//
//   - A lower class gives way to a higher one: opportunistic to standard,
//     standard to critical.
//   - Within a class, only a replica above its model's minimum, and only for
//     a model still below its own. Room a model grew into under load is not a
//     promise; its minimum is.
//   - A critical model's minimum is never taken.
//
// Victims only ever come from a lower rank, or from above a minimum for a
// model below one, so two models cannot take turns stopping each other.
//
// The set is tried on a copy of the inventory first, and only stopped if it
// would let the model be placed: stopping replicas that do not make room is
// the one outcome worse than not preempting.
func (r *Reconciler) preemptFor(ctx context.Context, m catalog.Model, avoid []string) bool {
	cands := r.victims(m)
	if len(cands) == 0 {
		return false
	}
	// Candidates in order until the model fits; then every one whose
	// stopping turns out not to be needed is kept after all. Without that
	// second step an idle replica on the wrong node, tried first, would be
	// stopped along with the one that actually made the room.
	var chosen []victim
	for _, v := range cands {
		if len(chosen) == maxVictims {
			break
		}
		chosen = append(chosen, v)
		if !r.fitsWithout(m, avoid, chosen) {
			continue
		}
		for i := len(chosen) - 1; i >= 0 && len(chosen) > 1; i-- {
			fewer := append(append([]victim(nil), chosen[:i]...), chosen[i+1:]...)
			if r.fitsWithout(m, avoid, fewer) {
				chosen = fewer
			}
		}
		r.preempt(ctx, m, chosen)
		return true
	}
	return false
}

// fitsWithout reports whether m could be placed with these replicas gone,
// tried on a copy of the inventory.
func (r *Reconciler) fitsWithout(m catalog.Model, avoid []string, gone []victim) bool {
	trial := r.inv.Clone()
	for _, v := range gone {
		trial.ReleaseReplica(v.ID)
	}
	sched := placement.New(trial)
	res := sched.Place(r.orderedProfiles(m), avoid)
	if res.Placement == nil && len(avoid) > 0 {
		res = sched.Place(r.orderedProfiles(m), nil)
	}
	return res.Placement != nil
}

// preempt drains the chosen replicas and holds their models back.
func (r *Reconciler) preempt(ctx context.Context, m catalog.Model, chosen []victim) {
	now := r.now()
	var ids []string
	r.mu.Lock()
	for _, v := range chosen {
		if rep, ok := r.replicas[v.ID]; ok {
			rep.PreemptedFor = m.Name
		}
		ids = append(ids, v.ID)
		r.held[v.Model] = hold{Until: now.Add(r.cfg.PreemptHold), For: m.Name}
		r.preempted[v.Model]++
	}
	r.preempting[m.Name] = ids
	r.mu.Unlock()

	victims := make([]string, 0, len(chosen))
	for _, v := range chosen {
		r.log.Warn("preempting a replica to make room",
			"replica", v.ID, "model", v.Model, "class", v.Class, "above_min", v.AboveMin,
			"for", m.Name, "for_class", m.EffectiveClass())
		desc := fmt.Sprintf("%s (%s, %s)", v.ID, v.Model, v.Class)
		if v.AboveMin {
			desc = fmt.Sprintf("%s (%s, %s, above its minimum)", v.ID, v.Model, v.Class)
		}
		victims = append(victims, desc)
	}
	ctx, span := startSpan(ctx, "preempt",
		attribute.String("llmgw.for", m.Name),
		attribute.String("llmgw.for_class", m.EffectiveClass()),
		attribute.StringSlice("llmgw.victims", victims),
		attribute.String("llmgw.hold", r.cfg.PreemptHold.String()))
	defer span.End()
	r.scaleDown(ctx, ids, "preempted for "+m.Name)
}

// awaitingPreemption reports whether replicas stopped for a model are still
// on their way out, so it waits for them rather than stopping more.
func (r *Reconciler) awaitingPreemption(model string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	ids := r.preempting[model]
	ids = slices.DeleteFunc(slices.Clone(ids), func(id string) bool {
		_, ok := r.replicas[id]
		return !ok
	})
	if len(ids) == 0 {
		delete(r.preempting, model)
		return false
	}
	r.preempting[model] = ids
	return true
}

// isHeld reports whether a model gave up room and may not take it back yet.
// The hold ends early once the model it was given to has what it asked for.
//
// Today a preempted replica is only removed by its own model's drain loop,
// which runs after that model has decided not to scale up, so the room it
// frees reaches the model it was freed for first anyway. The hold is what
// keeps that true for any other way a replica can go — a stop reported
// between passes, a reap at the start of one — rather than it resting on the
// order of lines in reconcileModel.
func (r *Reconciler) isHeld(model string) bool {
	r.mu.Lock()
	h, ok := r.held[model]
	r.mu.Unlock()
	if !ok {
		return false
	}
	if r.now().After(h.Until) {
		r.mu.Lock()
		delete(r.held, model)
		r.mu.Unlock()
		return false
	}
	if to, ok := r.cat.Load().Resolve(h.For); ok && r.liveCount(h.For) >= r.desiredFor(to) {
		r.mu.Lock()
		delete(r.held, model)
		r.mu.Unlock()
		return false
	}
	return true
}

// liveCount is how many replicas of a model hold capacity.
func (r *Reconciler) liveCount(model string) int {
	r.mu.Lock()
	defer r.mu.Unlock()
	n := 0
	for _, rep := range r.replicas {
		if rep.Model == model && rep.State.Counts() && rep.RetiringAt.IsZero() {
			n++
		}
	}
	return n
}

// victims lists the replicas that may be stopped for m, in the order they
// should be: idle on-demand ones first, then what other tenants borrowed past
// their guarantees when m's owner is within its own, then what lower classes
// grew into under load, then the same within m's class, then lower classes'
// minimums; within each, the longest unused first.
//
// A replica within its owner's guarantee is taken only for a model of the
// same owner: a guarantee that another tenant's higher class could take is
// not one.
func (r *Reconciler) victims(m catalog.Model) []victim {
	cat := r.cat.Load()
	now := r.now()
	myRank := catalog.ClassRank(m.EffectiveClass())
	belowMin := r.liveCount(m.Name) < r.minFor(m)

	pol := r.policy()
	var borrowed map[string]bool
	reclaim := false
	if pol != nil {
		hs := r.Holdings()
		borrowed = pol.Borrowed(hs)
		reclaim = pol.Guaranteed(m.Owner, pol.Usage(hs), profileNeed(m))
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	// Per model, its live replicas newest first: the newest are the ones
	// above its minimum.
	byModel := map[string][]*Replica{}
	for _, rep := range r.replicas {
		if rep.Model == m.Name || !rep.State.Counts() || !rep.RetiringAt.IsZero() || rep.PreemptedFor != "" {
			continue
		}
		byModel[rep.Model] = append(byModel[rep.Model], rep)
	}

	var out []victim
	for model, reps := range byModel {
		other, ok := cat.Resolve(model)
		if !ok {
			continue
		}
		class := other.EffectiveClass()
		rank := catalog.ClassRank(class)
		sort.Slice(reps, func(i, j int) bool { return reps[i].CreatedAt.After(reps[j].CreatedAt) })
		above := max(len(reps)-r.minFor(other), 0)
		for i, rep := range reps {
			aboveMin := i < above
			tier := -1
			switch {
			case rank < myRank && other.Scaling.OnDemand() && idleFor(rep, now) >= r.cfg.EvictionGrace:
				tier = 0
			case rank < myRank && aboveMin:
				tier = 2
			case rank == myRank && aboveMin && belowMin:
				tier = 3
			case rank < myRank && class != catalog.ClassCritical:
				tier = 4
			}
			if pol != nil && other.Owner != m.Owner {
				switch {
				case borrowed[rep.ID] && reclaim && (class != catalog.ClassCritical || aboveMin):
					// Past its owner's guarantee, wanted by an owner within
					// its own: taken back whatever the classes.
					if tier != 0 {
						tier = 1
					}
				case !borrowed[rep.ID] && guaranteed(pol, other.Owner):
					continue // held within its owner's guarantee
				}
			}
			if tier < 0 {
				continue
			}
			// Just started: stopping it throws away a model load that has
			// only now been paid for.
			if !rep.ReadyAt.IsZero() && now.Sub(rep.ReadyAt) < r.cfg.MinLifetime {
				continue
			}
			if len(nodesOf(rep)) > 1 {
				continue // see makeRoomFor: a multi-node replica is its own decision
			}
			prof, _ := other.ProfileByName(rep.Profile)
			var size quota.Amount
			for _, sh := range rep.Shards {
				size = size.Plus(quota.Amount{GPUs: len(sh.GPUIndices), CPUCores: cpuCores(prof), MemoryGB: prof.MinMemoryGB})
			}
			out = append(out, victim{ID: rep.ID, Model: model, Class: class,
				AboveMin: aboveMin, Tier: tier, Idle: idleFor(rep, now), Size: size})
		}
	}
	sort.SliceStable(out, func(i, j int) bool {
		a, b := out[i], out[j]
		if a.Tier != b.Tier {
			return a.Tier < b.Tier
		}
		if a.Idle != b.Idle {
			return a.Idle > b.Idle
		}
		// The smaller first: the least served is lost if it is enough, and
		// the pruning after keeps a larger one only when it is needed.
		if a.Size.GPUs != b.Size.GPUs {
			return a.Size.GPUs < b.Size.GPUs
		}
		if a.Size.MemoryGB != b.Size.MemoryGB {
			return a.Size.MemoryGB < b.Size.MemoryGB
		}
		return a.ID < b.ID
	})
	return out
}

// idleFor is how long a replica has gone unused; zero when nobody has said,
// which is not the same as idle.
func idleFor(rep *Replica, now time.Time) time.Duration {
	if rep.LastUsedAt.IsZero() {
		return 0
	}
	return now.Sub(rep.LastUsedAt)
}

func nodesOf(rep *Replica) map[string]struct{} {
	out := map[string]struct{}{}
	for _, sh := range rep.Shards {
		out[sh.NodeID] = struct{}{}
	}
	return out
}

// guaranteed reports whether an owner has a guarantee anywhere up its tree.
func guaranteed(pol *quota.Policy, owner string) bool {
	for _, q := range pol.Governing(owner) {
		if q.Guarantee.Any() {
			return true
		}
	}
	return false
}