llmgw/internal/control/modelserve/updates.go

package modelserve

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"regexp"
	"slices"
	"strings"
	"sync/atomic"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/bench/eval"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
	"github.com/prometheus/client_golang/prometheus"
)

// Updates: a repository publishes new weights under the same name — a fixed
// chat template, a better quantisation — and nothing told anyone. A model
// served from the registry now remembers where it came from; this asks the
// registry, every so often, whether the files it runs have changed there,
// and moves a model to them on request: fetched into the mirror beside the
// old ones, checked, swapped in, its replicas replaced one at a time, and
// the old files kept so it can go back.

// Models is the catalog as an update changes it.
type Models interface {
	Declared(name string) bool
	SetSource(ctx context.Context, name, source string, origin *catalog.Origin, by string) (catalog.Model, error)
	SetOrigin(ctx context.Context, name string, origin *catalog.Origin) error
	SetContext(ctx context.Context, name string, contextLength, maxOutput int, profiles []catalog.Profile, by string) (catalog.Model, error)
	SetProfiles(ctx context.Context, name string, profiles []catalog.Profile, by string) (catalog.Model, error)
	// For a canary: the candidate added beside the model and taken out
	// again, and the share of requests it is given.
	Add(ctx context.Context, m catalog.Model, by string) error
	Remove(ctx context.Context, name string) error
	SetCanary(ctx context.Context, name string, c *catalog.Canary, by string) (catalog.Model, error)
}

// UpdateInfo is what the registry has for a model against what it runs.
type UpdateInfo struct {
	Model string `json:"model"`
	Repo  string `json:"repo,omitempty"`
	// Current is the commit the model runs, Latest the repository's now.
	Current string `json:"current,omitempty"`
	Latest  string `json:"latest,omitempty"`
	// Changed are the model's files whose weights differ at Latest, and
	// Missing those Latest no longer has at all.
	Changed []string `json:"changed,omitempty"`
	Missing []string `json:"missing,omitempty"`
	// Bytes is what an update downloads.
	Bytes int64 `json:"bytes,omitempty"`
	// Available is true when the model can be moved to Latest.
	Available bool `json:"available"`
	// Previous is the commit an update replaced, which it can go back to.
	Previous  string    `json:"previous,omitempty"`
	CheckedAt time.Time `json:"checked_at"`
	Err       string    `json:"error,omitempty"`
}

// place is where a model's weights sit on the mirror.
type place struct {
	model, version string
	files          []mirror.Status
}

// placeOf finds a source on the mirror: a file, or a directory of shards.
func placeOf(source string, list mirror.ModelsResponse) (place, bool) {
	base := strings.TrimSuffix(list.Base, "/") + "/"
	if list.Base == "" || !strings.HasPrefix(source, base) {
		return place{}, false
	}
	parts := strings.SplitN(strings.TrimPrefix(source, base), "/", 3)
	if len(parts) < 2 {
		return place{}, false
	}
	p := place{model: parts[0], version: parts[1]}
	file := ""
	if len(parts) == 3 {
		file = parts[2]
	}
	for _, st := range list.Models {
		if st.Model == p.model && st.Version == p.version && (file == "" || st.File == file) {
			p.files = append(p.files, st)
		}
	}
	return p, len(p.files) > 0
}

// parseResolve reads a registry download URL: base/owner/name/resolve/rev/path.
func parseResolve(u string) (repo, rev, path string, ok bool) {
	pu, err := url.Parse(u)
	if err != nil {
		return "", "", "", false
	}
	i := strings.Index(pu.Path, "/resolve/")
	if i < 0 {
		return "", "", "", false
	}
	repo = strings.Trim(pu.Path[:i], "/")
	if strings.Count(repo, "/") != 1 {
		return "", "", "", false
	}
	rev, path, ok = strings.Cut(pu.Path[i+len("/resolve/"):], "/")
	return repo, rev, path, ok && rev != "" && path != ""
}

// inferOrigin works a model's origin out from what the mirror recorded of
// where it fetched each file, for a model served before origins were kept.
func inferOrigin(p place) (*catalog.Origin, bool) {
	var o catalog.Origin
	for _, st := range p.files {
		repo, rev, path, ok := parseResolve(st.Source)
		if !ok || (o.Repo != "" && (repo != o.Repo || rev != o.Revision)) {
			return nil, false
		}
		o.Repo, o.Revision = repo, rev
		o.Files = append(o.Files, catalog.OriginFile{Path: path, SHA256: st.SHA256, Bytes: st.BytesDone})
	}
	return &o, o.Repo != ""
}

func originFiles(set hf.FileSet) []catalog.OriginFile {
	out := make([]catalog.OriginFile, 0, len(set.Files))
	for _, f := range set.Files {
		out = append(out, catalog.OriginFile{Path: f.Path, SHA256: f.SHA256, Bytes: f.Bytes})
	}
	return out
}

// Finder is a registry that can be searched, and read at a past commit: what
// finding the origin of weights fetched before origins were kept needs.
type Finder interface {
	Explore(ctx context.Context, q hf.Query) ([]hf.Model, error)
	RepoAt(ctx context.Context, id, rev string) (hf.Repo, error)
}

// originOf is a model's origin: recorded; or worked out from where the
// mirror fetched its files; or found on the registry by the files' names,
// the commit the mirror's directory is named after, and their sizes. What is
// worked out is recorded, so it is worked out once.
func (s *Server) originOf(ctx context.Context, m catalog.Model, list mirror.ModelsResponse) *catalog.Origin {
	if m.Origin != nil && m.Origin.Repo != "" {
		return m.Origin
	}
	s.mu.Lock()
	o, known := s.origins[m.Name]
	s.mu.Unlock()
	if known {
		return o
	}
	p, ok := placeOf(m.Source, list)
	if !ok {
		return nil
	}
	o, ok = inferOrigin(p)
	if !ok {
		o = s.findOrigin(ctx, p)
	}
	s.remember(ctx, m.Name, o)
	return o
}

// remember keeps an origin: in the catalog for a model added at runtime, and
// here for one from configuration, whose definition is not this one's to
// change. A nil origin is remembered too, so a model nobody can trace is not
// searched for at every check.
func (s *Server) remember(ctx context.Context, name string, o *catalog.Origin) {
	s.mu.Lock()
	s.origins[name] = o
	s.mu.Unlock()
	if o == nil || s.deps.Models == nil || s.deps.Models.Declared(name) {
		return
	}
	if err := s.deps.Models.SetOrigin(ctx, name, o); err != nil {
		s.deps.Log.Debug("could not record a model's origin", "model", name, "err", err)
	}
}

var hexRev = regexp.MustCompile(`^[0-9a-f]{7,40}$`)

// findOrigin looks for the repository a mirror's files came from: one that
// holds a file of the same name, and has the commit the mirror's directory
// is named after with every file at the same size. Searched by the mirror's
// name for the model and the file's own; among the repositories that match,
// the most downloaded, which is the original rather than a copy of it.
func (s *Server) findOrigin(ctx context.Context, p place) *catalog.Origin {
	f, ok := s.deps.Registry.(Finder)
	if !ok || !hexRev.MatchString(p.version) {
		return nil
	}
	file := p.files[0].File
	stem := strings.TrimSuffix(file, ".gguf")
	if i := strings.LastIndexAny(stem, "-."); i > 0 {
		stem = stem[:i] // the quantisation
	}
	tried := map[string]bool{}
	for _, q := range []string{p.model, stem} {
		found, err := f.Explore(ctx, hf.Query{Search: q, GGUF: true, Sort: hf.SortDownloads, Limit: 100, WithFiles: true})
		if err != nil {
			s.deps.Log.Debug("searching the registry for a model's origin", "model", p.model, "err", err)
			continue
		}
		for _, c := range found {
			if tried[c.ID] || !slices.ContainsFunc(c.Siblings, func(sb struct {
				RFilename string `json:"rfilename"`
			}) bool {
				return mirror.LastPathElement(sb.RFilename) == file
			}) {
				continue
			}
			tried[c.ID] = true
			if o := s.verifyOrigin(ctx, f, c.ID, p); o != nil {
				s.deps.Log.Info("found where a model's weights came from", "model", p.model, "repo", c.ID,
					"revision", catalog.ShortRevision(o.Revision))
				return o
			}
		}
	}
	return nil
}

// verifyOrigin reads a repository at the mirror's commit and says whether
// every file the mirror holds is there at the same size.
func (s *Server) verifyOrigin(ctx context.Context, f Finder, id string, p place) *catalog.Origin {
	repo, err := f.RepoAt(ctx, id, p.version)
	if err != nil || !strings.HasPrefix(repo.Revision, p.version) {
		return nil
	}
	o := &catalog.Origin{Repo: repo.ID, Revision: repo.Revision}
	for _, st := range p.files {
		i := slices.IndexFunc(repo.Files, func(g hf.File) bool {
			return mirror.LastPathElement(g.Path) == st.File && g.Bytes == st.BytesDone
		})
		if i < 0 {
			return nil
		}
		g := repo.Files[i]
		o.Files = append(o.Files, catalog.OriginFile{Path: g.Path, SHA256: g.SHA256, Bytes: g.Bytes})
	}
	return o
}

// SetOriginOf records, by hand, the repository a model's weights came from;
// the commit is the one the mirror's directory is named after, and the files
// must be there at the sizes the mirror holds.
func (s *Server) SetOriginOf(ctx context.Context, name, repoID string) (*catalog.Origin, error) {
	m, err := s.current(name)
	if err != nil {
		return nil, err
	}
	list, err := s.deps.Mirror.List(ctx)
	if err != nil {
		return nil, fmt.Errorf("the mirror: %w", err)
	}
	p, ok := placeOf(m.Source, list)
	if !ok {
		return nil, fmt.Errorf("%s does not run from the mirror", name)
	}
	f, ok := s.deps.Registry.(Finder)
	if !ok {
		return nil, errors.New("this registry cannot be read at a past commit")
	}
	o := s.verifyOrigin(ctx, f, repoID, p)
	if o == nil {
		return nil, fmt.Errorf("%s has no commit %s holding %s at the size the mirror has", repoID, p.version,
			p.files[0].File)
	}
	s.remember(ctx, name, o)
	s.checkOne(ctx, m, o)
	return o, nil
}

// compare says what the registry's repository has against an origin.
func compare(m catalog.Model, o *catalog.Origin, repo hf.Repo) UpdateInfo {
	u := UpdateInfo{Model: m.Name, Repo: o.Repo, Current: o.Revision, Latest: repo.Revision, CheckedAt: time.Now()}
	if o.Previous != nil {
		u.Previous = o.Previous.Revision
	}
	if sameRevision(repo.Revision, o.Revision) {
		return u
	}
	for _, f := range o.Files {
		i := slices.IndexFunc(repo.Files, func(g hf.File) bool { return g.Path == f.Path })
		if i < 0 {
			u.Missing = append(u.Missing, f.Path)
			continue
		}
		g := repo.Files[i]
		if (f.SHA256 != "" && g.SHA256 != "" && !strings.EqualFold(f.SHA256, g.SHA256)) ||
			(f.SHA256 == "" && f.Bytes != 0 && f.Bytes != g.Bytes) {
			u.Changed = append(u.Changed, f.Path)
			u.Bytes += g.Bytes
		}
	}
	u.Available = len(u.Changed) > 0 && len(u.Missing) == 0
	return u
}

// sameRevision compares commits either of which may be shortened.
func sameRevision(a, b string) bool {
	if a == "" || b == "" {
		return a == b
	}
	return strings.HasPrefix(a, b) || strings.HasPrefix(b, a)
}

// CheckUpdates asks the registry about every model with an origin.
func (s *Server) CheckUpdates(ctx context.Context) {
	if s.deps.Current == nil {
		return
	}
	list, err := s.deps.Mirror.List(ctx)
	if err != nil {
		s.deps.Log.Warn("checking for model updates: the mirror", "err", err)
		return
	}
	seen := map[string]bool{}
	for _, m := range s.deps.Current() {
		o := s.originOf(ctx, m, list)
		if o == nil {
			continue
		}
		seen[m.Name] = true
		s.checkOne(ctx, m, o)
	}
	s.mu.Lock()
	for name := range s.updates {
		if !seen[name] {
			delete(s.updates, name)
		}
	}
	s.mu.Unlock()
}

func (s *Server) checkOne(ctx context.Context, m catalog.Model, o *catalog.Origin) UpdateInfo {
	repo, err := s.deps.Registry.Repo(ctx, o.Repo)
	u := UpdateInfo{Model: m.Name, Repo: o.Repo, Current: o.Revision, CheckedAt: time.Now()}
	if err != nil {
		u.Err = err.Error()
	} else {
		u = compare(m, o, repo)
		if u.Available {
			s.deps.Log.Info("a model's repository has new weights", "model", m.Name, "repo", o.Repo,
				"current", catalog.ShortRevision(o.Revision), "latest", catalog.ShortRevision(repo.Revision))
		}
	}
	s.mu.Lock()
	prev, had := s.updates[m.Name]
	s.updates[m.Name] = u
	s.mu.Unlock()
	// Told once per new revision, not at every check that finds it still there.
	if u.Available && (!had || !prev.Available || prev.Latest != u.Latest) {
		s.deps.Notify.Notify(notify.UpdateAvailable, m.Name+"@"+u.Latest,
			fmt.Sprintf("%s has new weights at %s (%s); `llmgw ctl model update %s` moves to them",
				m.Name, o.Repo, catalog.ShortRevision(u.Latest), m.Name),
			map[string]any{"model": m.Name, "repo": o.Repo, "current": o.Revision, "latest": u.Latest})
	}
	return u
}

// CheckModel asks the registry about one model now.
func (s *Server) CheckModel(ctx context.Context, name string) (UpdateInfo, error) {
	m, err := s.current(name)
	if err != nil {
		return UpdateInfo{}, err
	}
	list, err := s.deps.Mirror.List(ctx)
	if err != nil {
		return UpdateInfo{}, fmt.Errorf("the mirror: %w", err)
	}
	// Asked by hand: one that could not be traced is looked for again.
	s.mu.Lock()
	if o, ok := s.origins[name]; ok && o == nil {
		delete(s.origins, name)
	}
	s.mu.Unlock()
	o := s.originOf(ctx, m, list)
	if o == nil {
		return UpdateInfo{}, fmt.Errorf("where %s's weights came from is not known: give its repository with "+
			"`llmgw ctl model origin %s OWNER/REPO`", name, name)
	}
	return s.checkOne(ctx, m, o), nil
}

// Updates is what the last checks found, by model.
func (s *Server) Updates() map[string]UpdateInfo {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make(map[string]UpdateInfo, len(s.updates))
	for k, v := range s.updates {
		out[k] = v
	}
	return out
}

// WatchUpdates checks every so often until ctx ends; zero never checks.
func (s *Server) WatchUpdates(ctx context.Context, every time.Duration) {
	if every <= 0 {
		return
	}
	// A little after starting, when the catalog is in force and the mirror
	// answering, rather than in the middle of starting.
	wait := 2 * time.Minute
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
		}
		s.CheckUpdates(ctx)
		wait = every
	}
}

func (s *Server) current(name string) (catalog.Model, error) {
	if s.deps.Current == nil {
		return catalog.Model{}, errors.New("this control plane cannot read its catalog here")
	}
	for _, m := range s.deps.Current() {
		if m.Name == name {
			return m, nil
		}
	}
	return catalog.Model{}, fmt.Errorf("no model called %s", name)
}

// StartUpdate moves a model to its repository's latest weights, as a job.
// Held to a quality suite, the new weights run first as a canary given none
// of its requests, are asked the suite, and replace the old only at its
// score.
func (s *Server) StartUpdate(name, by string, ch Change) (Job, error) {
	if s.deps.Models == nil || s.deps.Restart == nil {
		return Job{}, errors.New("updating needs the store, which this control plane lacks")
	}
	m, err := s.current(name)
	if err != nil {
		return Job{}, err
	}
	if s.deps.Models.Declared(name) {
		return Job{}, fmt.Errorf("%s is declared in server.yaml: serve the new weights there, by its source, "+
			"or they would revert on the next restart", name)
	}
	g, skipped := gateFor(m, ch)
	if g != nil && m.Canary != nil {
		return Job{}, fmt.Errorf("%s has a canary, %s: an update held to a quality suite tries the new weights as "+
			"one. Promote it (`llmgw ctl model promote %s --eval ...`) or take it out first", name, m.Canary.Model, name)
	}
	if err := s.checkGate(canaryName(name), g); err != nil {
		return Job{}, err
	}
	s.mu.Lock()
	for _, j := range s.jobs {
		if j.Kind == "update" && j.Model == name && j.State == Running {
			s.mu.Unlock()
			return Job{}, fmt.Errorf("%s is already being updated: job %s", name, j.ID)
		}
		if g != nil && (j.Kind == "canary" || j.Kind == "promote") && j.Model == name && j.State == Running {
			s.mu.Unlock()
			return Job{}, fmt.Errorf("%s has a job running: %s", name, j.ID)
		}
	}
	s.seq++
	j := &Job{ID: fmt.Sprintf("update-%d-%d", time.Now().Unix(), s.seq), Kind: "update", Model: name,
		Request: Request{Name: name, By: by}, State: Running, Started: time.Now()}
	s.keepJob(j)
	s.mu.Unlock()
	if g != nil {
		go s.runGatedUpdate(j, g)
	} else {
		go s.runUpdate(j, skipped)
	}
	return s.snapshot(j), nil
}

// runGatedUpdate is an update held to a quality suite: the new weights as a
// canary with no share, asked the suite, then promoted or taken out.
func (s *Server) runGatedUpdate(j *Job, g *eval.Request) {
	ctx := context.Background()
	name := j.Model
	p, ok := s.prepare(ctx, j, name)
	if !ok {
		return
	}
	cand, ok := s.addCandidate(ctx, j, p, catalog.Canary{Percent: 0})
	if !ok {
		return
	}
	if !s.gate(ctx, j, cand, g, false, func() string {
		if _, err := s.AbortCanary(ctx, name, "llmgw"); err != nil {
			return fmt.Sprintf("the new weights could not be taken out (%v): `llmgw ctl model canary %s --abort`", err, name)
		}
		return fmt.Sprintf("the new weights were taken out again; %s still runs %s", name, catalog.ShortRevision(p.o.Revision))
	}) {
		return
	}
	if s.promoteTo(ctx, j) {
		s.finish(j)
	}
}

// prepared is a model's new weights, fetched and checked: what an update and
// a canary both start from.
type prepared struct {
	m      catalog.Model
	o      *catalog.Origin
	repo   hf.Repo
	set    hf.FileSet
	source string
}

// prepare reads where a model came from, fetches its repository's newer
// files into the mirror beside the ones in use, and checks the nodes can
// load them. It fails the job and says false when any of that cannot be.
func (s *Server) prepare(ctx context.Context, j *Job, name string) (prepared, bool) {

	// 1. The registry: which files, at which commit.
	s.step(j, "registry", Running, "reading where "+name+" came from")
	m, err := s.current(name)
	if err != nil {
		s.fail(j, "registry", err)
		return prepared{}, false
	}
	list, err := s.deps.Mirror.List(ctx)
	if err != nil {
		s.fail(j, "registry", fmt.Errorf("the mirror: %w", err))
		return prepared{}, false
	}
	o := s.originOf(ctx, m, list)
	p, onMirror := placeOf(m.Source, list)
	if o == nil || !onMirror {
		s.fail(j, "registry", fmt.Errorf("%s did not come from the registry through the mirror", name))
		return prepared{}, false
	}
	repo, err := s.deps.Registry.Repo(ctx, o.Repo)
	if err != nil {
		s.fail(j, "registry", err)
		return prepared{}, false
	}
	s.mu.Lock()
	j.Request.Repo = repo.ID
	s.mu.Unlock()
	u := compare(m, o, repo)
	switch {
	case sameRevision(repo.Revision, o.Revision):
		s.fail(j, "registry", fmt.Errorf("%s already runs %s's latest commit, %s", name, repo.ID, catalog.ShortRevision(repo.Revision)))
		return prepared{}, false
	case len(u.Missing) > 0:
		s.fail(j, "registry", fmt.Errorf("%s at %s no longer has %s; choose one of its files and serve it under a new name",
			repo.ID, repo.ShortRevision(), strings.Join(u.Missing, ", ")))
		return prepared{}, false
	case len(u.Changed) == 0:
		s.fail(j, "registry", fmt.Errorf("%s has a newer commit, %s, but the files %s runs are the same in it",
			repo.ID, repo.ShortRevision(), name))
		return prepared{}, false
	case repo.Gated.Is():
		s.fail(j, "registry", fmt.Errorf("%s is gated: it needs an account's approval the mirror does not have", repo.ID))
		return prepared{}, false
	}
	set := hf.FileSet{}
	for _, f := range o.Files {
		g := repo.Files[slices.IndexFunc(repo.Files, func(g hf.File) bool { return g.Path == f.Path })]
		set.Files = append(set.Files, g)
		set.Bytes += g.Bytes
	}
	set.Name, set.Shards = set.Files[0].Path, len(set.Files)
	s.step(j, "registry", Done, fmt.Sprintf("%s: %s → %s, %d of %d file(s) changed, %.1f GB",
		repo.ID, catalog.ShortRevision(o.Revision), repo.ShortRevision(), len(u.Changed), len(o.Files), gb(set.Bytes)))

	// 2. The mirror, beside the files in use.
	source, err := s.mirror(ctx, j, repo, set, p.model)
	if err != nil {
		s.fail(j, "mirror", err)
		return prepared{}, false
	}

	// 3. Whether the nodes load it, and whether it still fits what is
	// reserved for it.
	s.step(j, "check", Running, "reading its header from the mirror")
	detail := "no header to read"
	if u := headerURL(source, set); u != "" {
		v := gguf.Check(ctx, u, s.deps.LlamaVersion)
		if !v.Servable {
			s.fail(j, "check", fmt.Errorf("llama.cpp %s on the nodes cannot run the new weights: %s", s.deps.LlamaVersion, v.Reason))
			return prepared{}, false
		}
		detail = orUnknown(v.Header.Architecture)
		if len(m.Profiles) > 0 {
			ctxPool := 16384
			for i, a := range m.Profiles[0].Args {
				if a == "--ctx-size" && i+1 < len(m.Profiles[0].Args) {
					fmt.Sscan(m.Profiles[0].Args[i+1], &ctxPool)
				}
			}
			need := int((v.Header.MemoryBytes(set.Bytes, ctxPool) + 1<<30 - 1) >> 30)
			if have := m.Profiles[0].MinMemoryGB; need > have && have > 0 {
				detail += fmt.Sprintf("; needs about %d GB and %d is reserved: raise it if replicas fail", need, have)
			}
		}
	}
	s.step(j, "check", Done, detail)
	return prepared{m: m, o: o, repo: repo, set: set, source: source}, true
}

func (s *Server) runUpdate(j *Job, skipped bool) {
	ctx := context.Background()
	name := j.Model
	p, ok := s.prepare(ctx, j, name)
	if !ok {
		return
	}
	s.gate(ctx, j, name, nil, skipped, nil)
	m, o, repo, set, source := p.m, p.o, p.repo, p.set, p.source

	// 4. The catalog: the new source, and the old one kept to go back to.
	s.step(j, "catalog", Running, "moving "+name+" to "+repo.ShortRevision())
	next := &catalog.Origin{Repo: repo.ID, Revision: repo.Revision, Files: originFiles(set),
		Previous: &catalog.PreviousOrigin{Source: m.Source, Revision: o.Revision, Files: o.Files}}
	if _, err := s.deps.Models.SetSource(ctx, name, source, next, j.Request.By); err != nil {
		s.fail(j, "catalog", err)
		return
	}
	s.step(j, "catalog", Done, "the catalog names "+source)

	// 5. The replicas, replaced one at a time.
	if err := s.replace(ctx, j, name); err != nil {
		s.fail(j, "replicas", fmt.Errorf("%w. `llmgw ctl model rollback %s` goes back to %s", err, name,
			catalog.ShortRevision(o.Revision)))
		return
	}
	s.mu.Lock()
	s.updates[name] = UpdateInfo{Model: name, Repo: repo.ID, Current: repo.Revision, Latest: repo.Revision,
		Previous: o.Revision, CheckedAt: time.Now()}
	s.mu.Unlock()
	s.finish(j)
}

// replace restarts a model's replicas onto the catalog in force and waits
// for the new ones, then asks one to answer.
func (s *Server) replace(ctx context.Context, j *Job, name string) error {
	var old []string
	running := 0
	for _, r := range s.deps.Fleet.Replicas() {
		if r.Model != name {
			continue
		}
		running++
		if r.State == reconciler.Ready {
			old = append(old, r.ID)
		}
	}
	if running == 0 {
		s.step(j, "replicas", Skipped, "none running; the next to start loads the new weights")
		return nil
	}
	if len(old) == 0 {
		s.step(j, "replicas", Skipped, "none ready; those starting load the new weights when they next start")
		return nil
	}
	if err := s.deps.Restart(old...); err != nil {
		return err
	}
	s.step(j, "replicas", Running, fmt.Sprintf("replacing %d, one at a time", len(old)))
	deadline := time.Now().Add(s.deps.ReadyTimeout)
	for {
		left, ready := 0, []reconciler.Replica{}
		for _, r := range s.deps.Fleet.Replicas() {
			if r.Model != name {
				continue
			}
			if slices.Contains(old, r.ID) {
				left++
			} else if r.State == reconciler.Ready {
				ready = append(ready, r)
			}
		}
		for _, h := range s.deps.Fleet.ModelHealth() {
			if h.Model == name && h.Parked {
				return fmt.Errorf("the new weights failed to start %d times: %s", h.Failures, h.Reason)
			}
		}
		if left == 0 && len(ready) > 0 {
			for _, r := range ready {
				s.linkReplica(j, r)
			}
			s.step(j, "replicas", Done, fmt.Sprintf("%d replaced", len(old)))
			s.step(j, "smoke test", Running, "asking "+ready[0].ID+" to say hello")
			reply, tps, err := s.smoke(ctx, ready[0], name, s.embeds(name))
			if err != nil {
				s.step(j, "smoke test", Failed, err.Error())
				return err
			}
			s.mu.Lock()
			j.Reply, j.TokensPerSecond = reply, tps
			s.mu.Unlock()
			s.step(j, "smoke test", Done, fmt.Sprintf("answered at %.1f tokens/s: %q", tps, reply))
			return nil
		}
		s.step(j, "replicas", Running, fmt.Sprintf("%d of %d replaced, %d new ready", len(old)-left, len(old), len(ready)))
		if time.Now().After(deadline) {
			return fmt.Errorf("the replicas were not replaced after %s", s.deps.ReadyTimeout)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(s.deps.Poll):
		}
	}
}

// Rollback moves a model back to what its last update replaced, if those
// files are still on the mirror, and restarts its replicas onto them. The
// update is kept as the one to go forward to again.
func (s *Server) Rollback(ctx context.Context, name, by string) (string, error) {
	if s.deps.Models == nil || s.deps.Restart == nil {
		return "", errors.New("going back needs the store, which this control plane lacks")
	}
	m, err := s.current(name)
	if err != nil {
		return "", err
	}
	if m.Origin == nil || m.Origin.Previous == nil {
		return "", fmt.Errorf("%s has not been updated, so there is nothing to go back to", name)
	}
	prev := m.Origin.Previous
	list, err := s.deps.Mirror.List(ctx)
	if err != nil {
		return "", fmt.Errorf("the mirror: %w", err)
	}
	if _, ok := placeOf(prev.Source, list); !ok {
		return "", fmt.Errorf("the files of %s are no longer on the mirror", catalog.ShortRevision(prev.Revision))
	}
	back := &catalog.Origin{Repo: m.Origin.Repo, Revision: prev.Revision, Files: prev.Files,
		Previous: &catalog.PreviousOrigin{Source: m.Source, Revision: m.Origin.Revision, Files: m.Origin.Files}}
	if _, err := s.deps.Models.SetSource(ctx, name, prev.Source, back, by); err != nil {
		return "", err
	}
	var ids []string
	for _, r := range s.deps.Fleet.Replicas() {
		if r.Model == name && r.State == reconciler.Ready {
			ids = append(ids, r.ID)
		}
	}
	msg := fmt.Sprintf("%s runs %s again", name, catalog.ShortRevision(prev.Revision))
	if len(ids) > 0 {
		if err := s.deps.Restart(ids...); err != nil {
			return "", err
		}
		msg += fmt.Sprintf("; %d replica(s) are being replaced", len(ids))
	}
	s.mu.Lock()
	if u, ok := s.updates[name]; ok {
		u.Current, u.Previous = prev.Revision, m.Origin.Revision
		u.Available = u.Latest != "" && u.Latest != prev.Revision
		s.updates[name] = u
	}
	s.mu.Unlock()
	return msg, nil
}

// UpdateCollector publishes llmgw_model_update_available, 1 for each model
// whose repository has new weights. Registered when the metrics are, before
// there is a server to read: Set gives it one.
type UpdateCollector struct{ s atomic.Pointer[Server] }

var updateDesc = prometheus.NewDesc("llmgw_model_update_available",
	"1 when the registry repository a model was served from has published different weights for its files.",
	[]string{"model", "repo"}, nil)

// Set gives the collector the server whose checks it reports.
func (c *UpdateCollector) Set(s *Server) { c.s.Store(s) }

// Describe implements prometheus.Collector.
func (c *UpdateCollector) Describe(ch chan<- *prometheus.Desc) { ch <- updateDesc }

// Collect implements prometheus.Collector.
func (c *UpdateCollector) Collect(ch chan<- prometheus.Metric) {
	s := c.s.Load()
	if s == nil {
		return
	}
	for _, u := range s.Updates() {
		v := 0.0
		if u.Available {
			v = 1
		}
		ch <- prometheus.MustNewConstMetric(updateDesc, prometheus.GaugeValue, v, u.Model, u.Repo)
	}
}