llmgw/internal/control/adminapi/purge.go

package adminapi

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/url"
	"path/filepath"
	"slices"
	"strings"
	"time"

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

// A model removed with its files: the catalog entry, every version of it the
// mirror holds (and so, carried by its removal, a peer mirror's), and the
// copies the nodes fetched, once its replicas have stopped. What another
// model still uses stays.

// PurgeMirror is what a purge asks of the mirror.
type PurgeMirror interface {
	List(ctx context.Context) (mirror.ModelsResponse, error)
	Remove(ctx context.Context, model, version string) ([]string, error)
}

// SetPurgeMirror gives removing a model with its files the mirror.
func (a *API) SetPurgeMirror(m PurgeMirror) { a.purgeMirror = m }

// purgeWait is how long a purge waits for the model's replicas to stop
// before deleting the nodes' copies; after it, they are left for llmgw ctl
// weights --collect.
var purgeWait = 60 * time.Second

// PurgeView is what a model removed with its files took with it.
type PurgeView struct {
	Removed string `json:"removed"`
	// Mirror are the versions removed from the mirror, model/version.
	Mirror []string `json:"mirror,omitempty"`
	// Kept are versions another model still uses, left on the mirror.
	Kept  []string        `json:"kept,omitempty"`
	Nodes []CollectResult `json:"nodes,omitempty"`
	Note  string          `json:"note,omitempty"`
}

// purgeModel removes a model from the catalog and its files from the mirror
// and the nodes.
func (a *API) purgeModel(w http.ResponseWriter, r *http.Request, name string) {
	out, err := a.Purge(r.Context(), name, actorOf(r))
	switch {
	case errors.Is(err, errNoSuchModel):
		writeError(w, http.StatusNotFound, err.Error())
	case err != nil:
		writeError(w, http.StatusBadRequest, err.Error())
	default:
		writeJSON(w, http.StatusOK, out)
	}
}

// errNoSuchModel is a model the catalog does not have.
var errNoSuchModel = errors.New("no such model in the catalog")

// Purge removes a model from the catalog and its files from the mirror and
// the nodes, waiting a while for its replicas to stop before the nodes'.
func (a *API) Purge(ctx context.Context, name, by string) (PurgeView, error) {
	if a.models == nil {
		return PurgeView{}, errors.New("this control plane has no store")
	}
	m, ok := a.ctl.Catalog().Resolve(name)
	if !ok {
		return PurgeView{}, fmt.Errorf("%w: %s", errNoSuchModel, name)
	}
	name = m.Name
	others := []catalog.Model{}
	for _, o := range a.ctl.Catalog().Models() {
		if o.Name != name {
			others = append(others, o)
		}
	}
	if err := a.models.Remove(ctx, name); err != nil {
		return PurgeView{}, err
	}
	a.presetDeclined(ctx, name, by)
	out := PurgeView{Removed: name}

	sources := purgeSources(m)
	if a.purgeMirror == nil {
		out.Note = "this control plane has no mirror; its files there, if any, are left"
	} else {
		out.Mirror, out.Kept, out.Note = purgeFromMirror(ctx, a.purgeMirror, sources, others)
	}

	// The nodes' copies, once nothing runs on them.
	deadline := time.Now().Add(purgeWait)
	for a.runs(name) && time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return out, ctx.Err()
		case <-time.After(time.Second):
		}
	}
	if a.runs(name) {
		out.Note = strings.TrimSpace(out.Note + " Its replicas had not stopped yet: its copies on the nodes are " +
			"left, and llmgw ctl weights --collect deletes them.")
		return out, nil
	}
	out.Nodes = a.purgeFromNodes(ctx, sources, others)
	return out, nil
}

// runs says whether any replica of a model is still on a node.
func (a *API) runs(model string) bool {
	for _, rep := range a.ctl.Reconciler().Replicas() {
		if rep.Model == model {
			return true
		}
	}
	return false
}

// purgeSources are the files a model fetched: its weights, its adapters,
// and the weights its last update replaced.
func purgeSources(m catalog.Model) []string {
	out := m.WeightSources()
	for _, ad := range m.Adapters {
		out = append(out, ad.Source)
	}
	if m.Origin != nil && m.Origin.Previous != nil && m.Origin.Previous.Source != "" {
		out = append(out, m.Origin.Previous.Source)
	}
	return out
}

// inUse is every file the other models fetch.
func inUse(others []catalog.Model) []string {
	var out []string
	for _, o := range others {
		out = append(out, purgeSources(o)...)
	}
	return out
}

// purgeFromMirror removes from the mirror every version of the model
// directories the sources are in, but those another model uses.
func purgeFromMirror(ctx context.Context, mc PurgeMirror, sources []string, others []catalog.Model) (removed, kept []string, note string) {
	held, err := mc.List(ctx)
	if err != nil {
		return nil, nil, "the mirror could not be read (" + err.Error() + "): its files there are left"
	}
	dirs := map[string]bool{}
	for _, s := range sources {
		if model, _, ok := mirrorPlace(s); ok {
			dirs[model] = true
		}
	}
	used := map[string]bool{}
	for _, s := range inUse(others) {
		if model, version, ok := mirrorPlace(s); ok {
			used[model+"/"+version] = true
		}
	}
	seen := map[string]bool{}
	for _, st := range held.Models {
		key := st.Model + "/" + st.Version
		if !dirs[st.Model] || seen[key] {
			continue
		}
		seen[key] = true
		if used[key] {
			kept = append(kept, key)
			continue
		}
		if _, err := mc.Remove(ctx, st.Model, st.Version); err != nil {
			note = strings.TrimSpace(note + " " + key + " could not be removed from the mirror: " + err.Error() + ".")
			continue
		}
		removed = append(removed, key)
	}
	slices.Sort(removed)
	slices.Sort(kept)
	return removed, kept, note
}

// mirrorPlace is the model directory and version a mirror URL names:
// BASE/MODEL/VERSION/FILE.
func mirrorPlace(source string) (model, version string, ok bool) {
	u, err := url.Parse(source)
	if err != nil || (u.Scheme != "http" && u.Scheme != "https") {
		return "", "", false
	}
	// A mirror serves from its root: a registry's address is longer.
	parts := strings.Split(strings.Trim(u.Path, "/"), "/")
	if len(parts) != 3 {
		return "", "", false
	}
	model, version = parts[0], parts[1]
	if mirror.ValidName(model) != nil || mirror.ValidName(version) != nil {
		return "", "", false
	}
	return model, version, true
}

// purgeFromNodes deletes the nodes' copies of the sources: the files llmgw
// fetched for them under any mirror's address, that no other model wants.
func (a *API) purgeFromNodes(ctx context.Context, sources []string, others []catalog.Model) []CollectResult {
	rm := a.ctl.WeightRemover()
	if rm == nil {
		return nil
	}
	names := map[string]bool{}
	for _, s := range sources {
		names[filepath.Base(s)] = true
	}
	wanted := inUse(others)
	var out []CollectResult
	for _, n := range a.ctl.Inventory().Nodes() {
		size := map[string]int64{}
		have := make([]string, 0, len(n.Weights))
		for _, w := range n.Weights {
			have = append(have, w.Path)
			size[w.Path] = w.SizeBytes
		}
		var paths []string
		for _, p := range weights.Orphans(have, wanted) {
			// The cache name is a hash of the address, then the file's own
			// name: the address may be either mirror's, the name is the same.
			if _, file, ok := strings.Cut(filepath.Base(p), "-"); ok && names[file] {
				paths = append(paths, p)
			}
		}
		if len(paths) == 0 {
			continue
		}
		res := CollectResult{Node: n.ID, Files: len(paths)}
		for _, p := range paths {
			res.Bytes += size[p]
		}
		if err := rm.RemoveWeights(ctx, n.ID, paths); err != nil {
			res.Error, res.Files, res.Bytes = err.Error(), 0, 0
			slog.Warn("a purged model's copies could not be deleted on a node", "node", n.ID, "err", err)
		}
		out = append(out, res)
	}
	return out
}