llmgw/internal/control/reconciler/restart.go

package reconciler

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"time"

	"go.opentelemetry.io/otel/attribute"

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

// retire finishes restarting replicas.
//
// A retiring replica keeps serving until its replacement is ready, and is then
// drained like any other. When no replacement could be placed beside it, what
// happens depends on what is in the way:
//
//   - the replica itself — one node, or every card that fits held: it is
//     stopped now, and its replacement takes its place. Waiting would change
//     nothing; it used to wait DrainGrace anyway, five minutes reading only
//     "restart replacing", which looked like a restart that never would.
//   - more than it: the replacement fits nowhere even with it stopped — its
//     profile was edited to need more than any node has free. Stopping it
//     would leave the model with nothing, so the restart is given up, and
//     says why where the replica is shown.
//   - something being moved out of the way (a preemption, an idle model
//     stopping): waited for.
//   - nothing tried this pass (a backoff): the old rule, stopped after
//     DrainGrace with no replacement in sight.
//
// A restart asked for in place (ctl restart --in-place) stops first, always.
// Each stop is said in the log, since it is the one case with an outage.
func (r *Reconciler) retire(ctx context.Context, m catalog.Model, ids []string, readyLive, live, desired int,
	outcome scaleOutcome, parked bool,
) {
	for _, id := range ids {
		r.mu.Lock()
		rep, ok := r.replicas[id]
		var since time.Time
		if ok {
			since = rep.RetiringAt
		}
		inPlace := r.restartInPlace[id]
		r.mu.Unlock()
		if !ok {
			continue
		}
		switch {
		case readyLive >= desired:
			r.drainForRestart(ctx, id, "restarted: its replacement serves")
		case inPlace:
			r.log.Warn("restarting a replica in place, as asked: stopped before its replacement starts", "replica", id)
			r.stopForRestart(ctx, id, "restarted in place, as asked")
		case live >= desired:
			// A replacement is placed and loading; it is waited for.
		case parked:
			r.refuseRestart(ctx, id, "its replacements fail to start; the model is parked (llmgw ctl retry "+m.Name+")")
		case outcome == noRoom && r.fitsWithout(m, nil, []victim{{ID: id, Model: m.Name}}):
			r.log.Warn("restarting a replica in place: it is the only thing in its replacement's way",
				"replica", id, "model", m.Name)
			r.stopForRestart(ctx, id, "restarted in place: it was the only thing in its replacement's way")
		case outcome == noRoom || outcome == overLimit:
			r.refuseRestart(ctx, id, "its replacement fits on no node even with this replica stopped "+
				"(llmgw ctl explain-placement "+m.Name+")")
		case outcome == notTried && r.now().Sub(since) > r.cfg.DrainGrace:
			r.log.Warn("restarting a replica in place: no replacement beside it",
				"replica", id, "after", r.cfg.DrainGrace)
			r.stopForRestart(ctx, id, "restarted in place: no replacement beside it within "+r.cfg.DrainGrace.String())
		}
	}
}

// stopForRestart drains a replica and stops it at once, and places its
// replacement in what it freed.
//
// At once, not on the next pass: the pass goes on to the other models, and
// one waiting for room took the freed cards first. Seen on the reference
// cluster with a judge waiting for the only GPU: the model restarted there
// lost its card to the judge.
func (r *Reconciler) stopForRestart(ctx context.Context, id, why string) {
	r.mu.Lock()
	rep, ok := r.replicas[id]
	model := ""
	if ok {
		model = rep.Model
	}
	r.mu.Unlock()
	r.drainForRestart(ctx, id, why)
	r.stopReplica(ctx, id, why)
	m, found := r.cat.Load().Resolve(model)
	if !ok || !found || m.Name != model {
		return
	}
	r.mu.Lock()
	var avoid []string
	live := 0
	for _, o := range r.replicas {
		// As reconcileModel counts: one being retired is not one it has.
		if o.Model == model && o.State.Counts() && o.RetiringAt.IsZero() {
			live++
			avoid = append(avoid, o.NodeIDs()...)
		}
	}
	r.mu.Unlock()
	if live >= r.desiredFor(m) {
		return
	}
	if outcome := r.scaleUp(ctx, m, avoid); outcome != placed {
		r.log.Warn("the replacement of a replica restarted in place was not placed in its room",
			"replica", id, "model", model, "outcome", outcome.String())
	}
}

// refuseRestart gives up a replica's restart, leaving it serving, and keeps
// why for whoever looks at the replica.
func (r *Reconciler) refuseRestart(ctx context.Context, id, why string) {
	r.mu.Lock()
	if rep, ok := r.replicas[id]; ok {
		rep.RetiringAt = time.Time{}
	}
	delete(r.restartInPlace, id)
	r.restartRefused[id] = why
	r.mu.Unlock()
	r.log.Warn("not restarting a replica: it keeps serving", "replica", id, "because", why)
	_, span := startSpan(ctx, "restart refused",
		attribute.String("llmgw.replica", id), attribute.String("llmgw.reason", why))
	span.End()
}

func (r *Reconciler) drainForRestart(ctx context.Context, id, why string) {
	r.mu.Lock()
	rep, ok := r.replicas[id]
	var shards []Shard
	if ok {
		rep.State, rep.DrainingAt, rep.RetiringAt = Draining, r.now(), time.Time{}
		shards = slices.Clone(rep.Shards)
	}
	delete(r.restartInPlace, id)
	r.mu.Unlock()
	if ok {
		r.drainShards(ctx, id, shards, why)
	}
}

// nextRestartLocked pops the next queued replica of a model that is still
// ready to be restarted. Called with mu held.
func (r *Reconciler) nextRestartLocked(model string) (string, bool) {
	q := r.restartQueue[model]
	for len(q) > 0 {
		id := q[0]
		q = q[1:]
		if rep, ok := r.replicas[id]; ok && rep.State == Ready && rep.RetiringAt.IsZero() {
			r.restartQueue[model] = q
			return id, true
		}
	}
	delete(r.restartQueue, model)
	return "", false
}

// RestartPace is how many of a model's replicas restart at once.
type RestartPace struct {
	// MaxSurge are replaced beside themselves at once: their replacements
	// placed first, then they stop. Zero is one.
	MaxSurge int
	// MaxUnavailable are stopped first at once, their replacements starting
	// in their place: an outage for each while it loads, needing no room.
	MaxUnavailable int
}

// width is how many may be restarting at once.
func (p RestartPace) width() int {
	surge := p.MaxSurge
	if surge <= 0 && p.MaxUnavailable <= 0 {
		surge = 1
	}
	return surge + max(p.MaxUnavailable, 0)
}

// RestartPaced queues replicas like Restart (or RestartInPlace), as many of
// each model's at once as pace says.
func (r *Reconciler) RestartPaced(pace RestartPace, inPlace bool, ids ...string) error {
	if err := r.restart(inPlace, ids); err != nil {
		return err
	}
	r.mu.Lock()
	for _, id := range ids {
		if rep, ok := r.replicas[id]; ok {
			r.restartPace[rep.Model] = pace
		}
	}
	r.mu.Unlock()
	r.Notify()
	return nil
}

// ErrNoSuchReplica means no ready replica has that id.
var ErrNoSuchReplica = errors.New("no ready replica with that id")

// Restart queues replicas to be replaced, each by a new one started beside it
// before it is stopped — to pick up an edited profile, an engine key, a new
// engine version. Replicas of one model are replaced one at a time.
func (r *Reconciler) Restart(ids ...string) error { return r.restart(false, ids) }

// RestartInPlace queues replicas to be replaced by stopping each first and
// then starting its replacement: for when there is knowingly no room beside
// it, and the outage while it loads is accepted.
func (r *Reconciler) RestartInPlace(ids ...string) error { return r.restart(true, ids) }

func (r *Reconciler) restart(inPlace bool, ids []string) error {
	r.mu.Lock()
	for _, id := range ids {
		rep, ok := r.replicas[id]
		if !ok || rep.State != Ready {
			r.mu.Unlock()
			return fmt.Errorf("%w: %s", ErrNoSuchReplica, id)
		}
		delete(r.restartRefused, id)
		if inPlace {
			r.restartInPlace[id] = true
		}
		if !slices.Contains(r.restartQueue[rep.Model], id) && rep.RetiringAt.IsZero() {
			r.restartQueue[rep.Model] = append(r.restartQueue[rep.Model], id)
		}
	}
	r.mu.Unlock()
	r.Notify()
	return nil
}

// RestartState says where a replica's restart stands: "queued",
// "replacing", "in place", "refused: why", or "" when none was asked for.
func (r *Reconciler) RestartState(id string) string {
	r.mu.Lock()
	defer r.mu.Unlock()
	rep, ok := r.replicas[id]
	switch {
	case !ok:
		return ""
	case !rep.RetiringAt.IsZero() && r.restartInPlace[id]:
		return "in place"
	case !rep.RetiringAt.IsZero():
		return "replacing"
	case slices.Contains(r.restartQueue[rep.Model], id):
		return "queued"
	case r.restartRefused[id] != "":
		return "refused: " + r.restartRefused[id]
	}
	return ""
}