llmgw/internal/web/admin_mirror.go

package web

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// MirrorAdmin is the mirror's management API as this page needs it.
//
// An interface rather than the client itself so the page can be tested without
// a mirror, and so a deployment without one leaves it nil and the page
// unmounted — the same shape the cluster dashboard uses.
type MirrorAdmin interface {
	List(ctx context.Context) (mirror.ModelsResponse, error)
	Fetch(ctx context.Context, req mirror.FetchRequest) (mirror.FetchResponse, error)
	Remove(ctx context.Context, model, version string) ([]string, error)
	// Addr is the mirror this points at, shown on the page because an operator
	// looking at an empty table needs to know which machine is empty.
	Addr() string
}

// Registry is the model registry the explorer browses. Nil leaves the explorer
// off, which is right for a deployment with no way out to the internet: a
// search box that can only time out is worse than no search box.
type Registry interface {
	Search(ctx context.Context, query string, limit int) ([]hf.Model, error)
	Repo(ctx context.Context, id string) (hf.Repo, error)
}

// repoSet is one thing worth fetching: a file, or every shard of one.
type repoSet struct {
	hf.FileSet
	// Sources are every file in the set, in the order the registry listed them.
	// A shard set goes as one action — a fifteen-shard model is not fifteen
	// decisions, it is one.
	Sources []repoSource
	// Model and Version are what the fetch form starts with. Editable, because
	// the name a catalog entry uses is the operator's to choose and the
	// registry's idea of it is only a good guess.
	Model   string
	Version string
}

type repoSource struct{ URL, Path string }

// mirrorRow is one model version as the page shows it.
type mirrorRow struct {
	mirror.Status
	// URL is what a catalog entry should say. The one string on this page
	// somebody will copy, so it is rendered rather than left to be assembled.
	URL     string
	Percent int
	// HasPercent is false when the upstream never sent a length. The bar is
	// then not drawn at all: a progress bar that has to invent its own total is
	// worse than a byte count that only goes up.
	HasPercent bool
}

type adminMirrorData struct {
	Addr string
	Base string
	Rows []mirrorRow

	Fetching int
	Present  int
	Failed   int

	// Unreachable separates "the mirror holds nothing" from "the mirror did not
	// answer". They render almost identically and mean opposite things: one is
	// a mirror waiting to be filled, the other is a machine to go and look at.
	Unreachable bool

	// The explorer.
	HasRegistry bool
	Query       string
	Results     []hf.Model
	Repo        *hf.Repo
	RepoSets    []repoSet
	RegistryErr string

	Message string
	Error   string
	CSRF    string
	Version string
	Now     string
	// RefreshSec is 0 when nothing is downloading. The page reloads only while
	// there is something to watch, because a reload while an operator is typing
	// a source URL into the form takes the URL with it.
	RefreshSec int
}

// RegisterAdminMirror mounts the model mirror under /admin/mirror.
//
// Fetching and removing, and nothing else. What the mirror holds is operator
// state rather than account state, but the reasoning is the same as
// /admin/users: these two are the work somebody does repeatedly and should not
// need a shell for, while anything that would let a stolen session reach
// further — pointing the control plane at a different mirror, say — stays in
// configuration where changing it means access to the host.
func (p *AccountPages) RegisterAdminMirror(mux *http.ServeMux, m MirrorAdmin, reg Registry) {
	p.hasMirror = true
	p.registry = reg

	mux.HandleFunc("GET /admin/mirror", p.adminOnly(
		func(w http.ResponseWriter, r *http.Request, _ store.User) {
			p.renderMirror(w, r, m, "", "")
		}))

	mux.HandleFunc("POST /admin/mirror/fetch", p.mirrorAction(m,
		func(ctx context.Context, admin store.User, r *http.Request) (string, error) {
			// One form posts one source or fifteen. A sharded model is one
			// decision, not one per shard, and the mirror fetches each file
			// independently anyway.
			sources, files := r.Form["source"], r.Form["file"]
			if len(sources) == 0 || strings.TrimSpace(sources[0]) == "" {
				return "", errors.New("a source URL is required")
			}
			model := strings.TrimSpace(r.FormValue("model"))
			version := strings.TrimSpace(r.FormValue("version"))

			var started int
			for i, src := range sources {
				src = strings.TrimSpace(src)
				if src == "" {
					continue
				}
				req := mirror.FetchRequest{Source: src, Model: model, Version: version}
				if i < len(files) {
					req.File = strings.TrimSpace(files[i])
				}
				if _, err := m.Fetch(ctx, req); err != nil {
					// Reported with what already went, so a set that fails on
					// its fourth shard does not read as though nothing started.
					return "", fmt.Errorf("%d of %d started, then %s: %w",
						started, len(sources), mirror.LastPathElement(src), err)
				}
				started++
			}
			// Named as started, not as done. The mirror answered 202 and the
			// bytes are still moving; a message saying "added" would invite
			// somebody to point the catalog at a partial file.
			return fmt.Sprintf("started downloading %s/%s (%s)",
				model, version, plural(started, "file", "files")), nil
		}))

	mux.HandleFunc("POST /admin/mirror/{model}/{version}/remove", p.mirrorAction(m,
		func(ctx context.Context, _ store.User, r *http.Request) (string, error) {
			model, ver := r.PathValue("model"), r.PathValue("version")
			removed, err := m.Remove(ctx, model, ver)
			if err != nil {
				return "", err
			}
			return "removed " + model + "/" + ver +
				" (" + plural(len(removed), "file", "files") + ")", nil
		}))
}

// mirrorAction wraps a change: admin session, CSRF, audit line, re-render.
//
// The audit line is not optional. Removing a version deletes weights every node
// may be about to fetch, and the only record that it happened on purpose is
// this log entry naming who clicked it.
func (p *AccountPages) mirrorAction(m MirrorAdmin,
	do func(ctx context.Context, admin store.User, r *http.Request) (string, error),
) http.HandlerFunc {
	return p.adminOnly(func(w http.ResponseWriter, r *http.Request, admin store.User) {
		sess, _, _ := p.session(r)
		if !p.accounts.CheckCSRF(sess, r.FormValue("csrf")) {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		msg, err := do(r.Context(), admin, r)
		if err != nil {
			p.log.Warn("a mirror action was refused",
				"admin", admin.Email, "mirror", m.Addr(), "err", err)
			p.renderMirror(w, r, m, "", err.Error())
			return
		}
		p.log.Info("the mirror was changed",
			"admin", admin.Email, "mirror", m.Addr(), "what", msg)
		p.renderMirror(w, r, m, msg, "")
	})
}

func (p *AccountPages) renderMirror(w http.ResponseWriter, r *http.Request,
	m MirrorAdmin, msg, errMsg string,
) {
	d := adminMirrorData{
		Addr:    m.Addr(),
		Message: msg,
		Error:   errMsg,
		CSRF:    csrfFor(r),
		Version: p.version,
		Now:     time.Now().UTC().Format(time.RFC3339),
	}

	p.explore(r, &d)

	got, err := m.List(r.Context())
	if err != nil {
		p.log.Error("could not read the mirror", "mirror", m.Addr(), "err", err)
		d.Unreachable = true
		if d.Error == "" {
			d.Error = err.Error()
		}
		render(w, "admin_mirror", d)
		return
	}

	d.Base = got.Base
	for _, st := range got.Models {
		row := mirrorRow{
			Status: st,
			URL: mirror.Entry{Model: st.Model, Version: st.Version, File: st.File}.
				URL(got.Base),
		}
		if pc, ok := st.Percent(); ok {
			row.Percent, row.HasPercent = pc, true
		}
		switch st.State {
		case mirror.Fetching:
			d.Fetching++
		case mirror.Failed:
			d.Failed++
		default:
			d.Present++
		}
		d.Rows = append(d.Rows, row)
	}

	// Downloads first, then failures, then what is simply here. The two states
	// somebody opened this page for are at the top; the rest is a list they
	// scroll. Within a state the order is fixed so rows do not shuffle between
	// reloads while a byte count ticks.
	sort.SliceStable(d.Rows, func(i, j int) bool {
		a, b := d.Rows[i], d.Rows[j]
		if ra, rb := mirrorRank(a.State), mirrorRank(b.State); ra != rb {
			return ra < rb
		}
		if a.Model != b.Model {
			return a.Model < b.Model
		}
		if a.Version != b.Version {
			return a.Version < b.Version
		}
		return a.File < b.File
	})

	if d.Fetching > 0 {
		d.RefreshSec = 5
	}
	render(w, "admin_mirror", d)
}

func mirrorRank(s mirror.State) int {
	switch s {
	case mirror.Fetching:
		return 0
	case mirror.Failed:
		return 1
	default:
		return 2
	}
}

func plural(n int, one, many string) string {
	if n == 1 {
		return "1 " + one
	}
	return strconv.Itoa(n) + " " + many
}

// explore fills in the registry half of the page: a search, or one repository.
//
// A failure here is reported in its own line rather than as the page's error.
// The registry is somebody else's server on the other side of the internet, and
// it being down says nothing about the mirror — which is the thing this page is
// actually for, and which must keep working when Hugging Face does not.
func (p *AccountPages) explore(r *http.Request, d *adminMirrorData) {
	if p.registry == nil {
		return
	}
	d.HasRegistry = true
	d.Query = strings.TrimSpace(r.URL.Query().Get("q"))

	if repo := strings.TrimSpace(r.URL.Query().Get("repo")); repo != "" {
		got, err := p.registry.Repo(r.Context(), repo)
		if err != nil {
			d.RegistryErr = err.Error()
			return
		}
		d.Repo = &got
		var weights []hf.File
		for _, f := range got.Files {
			if !f.IsWeights() {
				// Thirty files of which four are worth gigabytes. The rest are
				// a README and a tokenizer, and listing them buries the ones
				// somebody came for.
				continue
			}
			weights = append(weights, f)
		}
		for _, set := range hf.Group(weights) {
			rs := repoSet{
				FileSet: set,
				Model:   defaultModelName(got.ID),
				// The revision, not the date: two fetches of the same commit
				// land in the same place, and a repository re-uploaded under
				// the same branch lands somewhere else.
				Version: got.ShortRevision(),
			}
			for _, f := range set.Files {
				rs.Sources = append(rs.Sources,
					repoSource{URL: got.DownloadURL("", f.Path), Path: f.Path})
			}
			d.RepoSets = append(d.RepoSets, rs)
		}
		return
	}

	if d.Query != "" {
		got, err := p.registry.Search(r.Context(), d.Query, 20)
		if err != nil {
			d.RegistryErr = err.Error()
			return
		}
		d.Results = got
	}
}

// defaultModelName is the catalog name a repository suggests.
//
// The last element, lowercased, with the quantisation marker dropped: the owner
// is not part of what anybody calls the model, and "-GGUF" describes the file
// format rather than the model, which is already clear from the file being
// fetched.
func defaultModelName(repoID string) string {
	name := mirror.LastPathElement(repoID)
	name = strings.TrimSuffix(name, "-GGUF")
	name = strings.TrimSuffix(name, "-gguf")
	return strings.ToLower(name)
}