llmgw/internal/cli/mirror.go

package cli

import (
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/netip"
	"os"
	"slices"
	"strings"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// newMirrorCmd manages a directory of model weights served to the fleet.
//
// Operator tooling, run on the mirror host, the way `llmgw keys` is run on the
// machine holding the store. The gateway never writes here: a mirror that can
// be written to by the thing reading from it is a way to replace a model with
// another one that every node would then fetch without complaint.
func newMirrorCmd() *cobra.Command {
	var dir string

	c := &cobra.Command{
		Use:   "mirror",
		Short: "Hold model weights for the fleet to fetch from",
		Long: "One machine holding the model files, so nodes fetch over your own\n" +
			"network instead of the internet. Eight nodes each pulling forty gigabytes\n" +
			"from outside is a day's work, and on many corporate networks it is not\n" +
			"possible at all.\n\n" +
			"Files live at <model>/<version>/<file>, and a version is never\n" +
			"overwritten. That is what makes rolling back possible: changing the\n" +
			"catalog line back works because the old weights are still there.",
	}
	c.PersistentFlags().StringVar(&dir, "dir", "/srv/models", "the mirror's root directory")

	c.AddCommand(
		newMirrorAddCmd(&dir),
		newMirrorListCmd(&dir),
		newMirrorUnusedCmd(&dir),
		newMirrorRemoveCmd(&dir),
		newMirrorServeCmd(&dir),
		newMirrorInitCmd(&dir),
	)
	return c
}

func newMirrorAddCmd(dir *string) *cobra.Command {
	var model, version, as string

	c := &cobra.Command{
		Use:   "add <url>...",
		Short: "Download a model into the mirror",
		Long: "Downloads to a temporary name and renames when complete, so an\n" +
			"interrupted download never leaves something a catalog could be pointed\n" +
			"at. Reports the checksum: a mirror is the one place in this system where\n" +
			"a corrupted file is copied to every node rather than caught once.\n\n" +
			"Several URLs go into one version, as the files of one model: the shards\n" +
			"of a safetensors set with its config.json and tokenizer, for vLLM. The\n" +
			"version is offered to nodes only once it is whole — nothing still\n" +
			"downloading, every shard its index names present — and its source is\n" +
			"the version, ending in a slash.\n\n" +
			"Refuses to replace a version that is already here. Add a new one beside\n" +
			"it instead — that is what makes the old one still available when the new\n" +
			"one turns out to be worse.",
		Args: cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if model == "" || version == "" {
				return fmt.Errorf("--model and --version are required: they are the layout")
			}
			if as != "" && len(args) > 1 {
				return fmt.Errorf("--as names one file; with %d URLs each keeps its own name", len(args))
			}
			out := cmd.OutOrStdout()
			var last mirror.Added
			for _, source := range args {
				file := as
				if file == "" {
					file = mirror.LastPathElement(source)
				}
				e := mirror.Entry{Model: model, Version: version, File: file}
				fmt.Fprintf(out, "fetching %s\n", source)
				added, err := mirror.Add(cmd.Context(), *dir, source, e, func(done, total int64) {
					if total > 0 {
						fmt.Fprintf(out, "  %s of %s (%d%%)\n",
							fileSize(done), fileSize(total), done*100/total)
						return
					}
					fmt.Fprintf(out, "  %s so far\n", fileSize(done))
				})
				if err != nil {
					return err
				}
				fmt.Fprintf(out, "%s/%s/%s  %s\n  sha256 %s\n",
					added.Entry.Model, added.Entry.Version, added.Entry.File,
					fileSize(added.Entry.Bytes), added.SHA256)
				last = added
			}
			src := last.Entry.URL("http://<this-host>:8081")
			if len(args) > 1 {
				src = strings.TrimSuffix(mirror.Entry{Model: model, Version: version}.URL("http://<this-host>:8081"), "/") + "/"
				if m, err := mirror.ManifestFor(*dir, model, version); err != nil {
					fmt.Fprintf(out, "\nNot yet offered to nodes: %v\n", err)
				} else {
					fmt.Fprintf(out, "\n%d files, %s, whole.\n", len(m.Files), fileSize(m.Bytes))
				}
			}
			fmt.Fprintf(out, "\nPoint a model at it with:\n  source: %s\n", src)
			return nil
		},
	}
	c.Flags().StringVar(&model, "model", "", "the model's name, one directory level")
	c.Flags().StringVar(&version, "version", "", "the version, one directory level; a date works well")
	c.Flags().StringVar(&as, "as", "", "the filename to save as, for one URL; taken from the URL when omitted")
	return c
}

func newMirrorListCmd(dir *string) *cobra.Command {
	return &cobra.Command{
		Use:   "list",
		Short: "Show what the mirror holds",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			entries, err := mirror.List(*dir)
			if err != nil {
				return err
			}
			if len(entries) == 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "The mirror at %s is empty.\n", *dir)
				return nil
			}
			w := tab()
			fmt.Fprintln(w, "MODEL\tVERSION\tFILE\tSIZE")
			var total int64
			for _, e := range entries {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", e.Model, e.Version, e.File, fileSize(e.Bytes))
				total += e.Bytes
			}
			_ = w.Flush()
			fmt.Fprintf(cmd.OutOrStdout(), "\n%s in %d files.\n", fileSize(total), len(entries))
			return nil
		},
	}
}

func newMirrorUnusedCmd(dir *string) *cobra.Command {
	var server, base string

	c := &cobra.Command{
		Use:   "unused",
		Short: "Show files no model in the catalog refers to",
		Long: "Compares what is here against every source in the catalog, so the\n" +
			"answer is what could be deleted rather than what looks old.\n\n" +
			"It refuses to answer when it cannot read the catalog. An empty catalog\n" +
			"and an unreachable one produce the same list — everything — and the\n" +
			"obvious next step after that list is deletion.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			entries, err := mirror.List(*dir)
			if err != nil {
				return err
			}
			sources, err := catalogSources(cmd, server)
			if err != nil {
				// Named as the reason rather than folded into a generic
				// failure: an operator who sees "no catalog" will fix the
				// address, and one who sees a list will start deleting.
				return fmt.Errorf("cannot read the catalog, so nothing can be called unused: %w", err)
			}

			unused := mirror.Unused(entries, sources, base)
			out := cmd.OutOrStdout()
			if len(unused) == 0 {
				fmt.Fprintf(out, "Every file here is referred to by a model. (%d files, %d catalog sources.)\n",
					len(entries), len(sources))
				return nil
			}
			w := tab()
			fmt.Fprintln(w, "MODEL\tVERSION\tFILE\tSIZE")
			var total int64
			for _, e := range unused {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", e.Model, e.Version, e.File, fileSize(e.Bytes))
				total += e.Bytes
			}
			_ = w.Flush()
			fmt.Fprintf(out, "\n%s in %d files that no model refers to.\n", fileSize(total), len(unused))
			fmt.Fprintln(out, "\nBefore deleting: a node that has to refetch one of these will not be\n"+
				"able to. Retire the catalog entry first, let the fleet settle, then remove.")
			return nil
		},
	}
	c.Flags().StringVar(&server, "server", "127.0.0.1:9090", "admin API address to read the catalog from")
	c.Flags().StringVar(&base, "base", "http://localhost:8081",
		"the address nodes use to reach this mirror, as it appears in catalog sources")
	return c
}

func newMirrorRemoveCmd(dir *string) *cobra.Command {
	var model, version, file, server, base string
	var force bool

	c := &cobra.Command{
		Use:   "remove",
		Short: "Delete one version from the mirror",
		Long: "Checks the catalog first, and refuses if a model still refers to it —\n" +
			"a node that has to refetch a deleted file cannot, and the failure arrives\n" +
			"later, on the node, looking like something else.\n\n" +
			"Retire the catalog entry first, let the fleet settle, then remove.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if model == "" || version == "" || file == "" {
				return fmt.Errorf("--model, --version and --file are all required")
			}
			e := mirror.Entry{Model: model, Version: version, File: file}

			if !force {
				sources, err := catalogSources(cmd, server)
				if err != nil {
					return fmt.Errorf("cannot read the catalog to check whether this is still in use: %w\n"+
						"  Pass --force to remove it anyway", err)
				}
				if len(mirror.Unused([]mirror.Entry{e}, sources, base)) == 0 {
					return fmt.Errorf("a model still refers to %s/%s/%s; retire the catalog entry first",
						model, version, file)
				}
			}
			if err := mirror.Remove(*dir, e); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "removed %s/%s/%s\n", model, version, file)
			return nil
		},
	}
	c.Flags().StringVar(&model, "model", "", "the model directory")
	c.Flags().StringVar(&version, "version", "", "the version directory")
	c.Flags().StringVar(&file, "file", "", "the filename")
	c.Flags().StringVar(&server, "server", "127.0.0.1:9090", "admin API address to read the catalog from")
	c.Flags().StringVar(&base, "base", "http://localhost:8081", "the address nodes use to reach this mirror")
	c.Flags().BoolVar(&force, "force", false, "remove without checking the catalog")
	return c
}

func newMirrorServeCmd(dir *string) *cobra.Command {
	var listen, apiListen, base string
	var tlsc tlsconf.Config
	var allow, clients []string
	var allowPrivate bool

	c := &cobra.Command{
		Use:   "serve",
		Short: "Serve the mirror, and an API to manage it",
		Long: "Two listeners, on purpose.\n\n" +
			"The first serves the weights themselves, read-only and unauthenticated,\n" +
			"because that is what eight nodes pull forty gigabytes from and putting a\n" +
			"client certificate in front of it would mean re-certifying every agent's\n" +
			"fetcher to change how models are administered. Range requests work, which\n" +
			"is what makes an interrupted download on a node resume rather than restart.\n\n" +
			"The second is the management API the control plane calls: ask the mirror to\n" +
			"fetch a model, watch it arrive, remove one that is retired. It can be given\n" +
			"the fleet's CA and then only holders of a certificate from it get in.\n\n" +
			"The mirror does its own downloading. A control plane that had to pull the\n" +
			"bytes and push them here would carry every model twice over the network,\n" +
			"through the one machine least able to spare the bandwidth.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			root, err := os.Stat(*dir)
			if err != nil || !root.IsDir() {
				return fmt.Errorf("%s is not a directory to serve", *dir)
			}
			if err := tlsc.Validate(); err != nil {
				return err
			}
			if len(clients) > 0 && !tlsc.Enabled() {
				return fmt.Errorf("--api-client names a certificate, and needs --tls-ca, --tls-cert and --tls-key")
			}

			out := cmd.OutOrStdout()
			log := slog.Default()
			mgr := mirror.NewManager(*dir, log)
			if allowPrivate {
				mgr.AllowPrivateSources()
			}
			api := mirror.NewAPI(mgr, strings.TrimSuffix(base, "/"))

			// http.Dir refuses paths that escape the root, which is the whole
			// traversal question answered by the standard library rather than
			// by me.
			files := &http.Server{
				Addr:              listen,
				Handler:           readOnly(withManifest(*dir, http.FileServer(http.Dir(*dir)))),
				ReadHeaderTimeout: 10 * time.Second,
			}
			allowed, err := parseAllow(allow)
			if err != nil {
				return err
			}
			admin := &http.Server{
				Addr:              apiListen,
				Handler:           onlyClients(clients, apiGate(tlsc.Enabled(), allowed, api.Handler())),
				ReadHeaderTimeout: 10 * time.Second,
			}

			errs := make(chan error, 2)
			go func() { errs <- fmt.Errorf("serving weights: %w", files.ListenAndServe()) }()
			fmt.Fprintf(out, "serving %s on %s\n", *dir, listen)

			go func() { errs <- fmt.Errorf("serving the API: %w", serveAPI(admin, tlsc)) }()
			switch {
			case tlsc.Enabled():
				fmt.Fprintf(out, "management API on %s, mutual TLS against %s\n", apiListen, tlsc.CA)
				if len(clients) == 0 {
					// The same root signs every node's certificate, so without a
					// name any node could delete models here.
					log.Warn("the management API takes any certificate from " + tlsc.CA +
						", every node's included; name the control plane's with --api-client")
				}
			case len(allowed) > 0:
				fmt.Fprintf(out, "management API on %s, for %s and this machine only\n", apiListen, strings.Join(allow, ", "))
			default:
				// This machine only, which was once "anyone": an API that can
				// delete models and start downloads, open to every node that
				// could reach the port, is a choice the operator should make
				// with --api-allow rather than inherit.
				fmt.Fprintf(out, "management API on %s, for this machine only (--api-allow to add the control plane)\n", apiListen)
			}
			fmt.Fprintf(out, "catalog sources should point at %s\n", base)

			// The first failure ends it. A mirror serving bytes but refusing
			// management, or the reverse, is half a service pretending to be
			// whole, and the half that works keeps the alert quiet.
			return <-errs
		},
	}
	c.Flags().StringVar(&listen, "listen", ":8081", "address to serve the weights on")
	c.Flags().StringVar(&apiListen, "api-listen", ":8082", "address to serve the management API on")
	c.Flags().StringVar(&base, "base", "http://localhost:8081",
		"the address nodes use to reach this mirror, as it should appear in catalog sources")
	c.Flags().StringArrayVar(&allow, "api-allow", nil,
		"an address or network the management API answers, besides this machine; repeatable (default: this machine only)")
	c.Flags().StringArrayVar(&clients, "api-client", nil,
		"with mutual TLS, the common name of a certificate the management API answers; repeatable (default: any from --tls-ca)")
	c.Flags().BoolVar(&allowPrivate, "allow-private-sources", false,
		"let downloads reach loopback, private and link-local addresses")
	c.Flags().StringVar(&tlsc.CA, "tls-ca", "", "CA bundle that management clients must present a certificate from")
	c.Flags().StringVar(&tlsc.Cert, "tls-cert", "", "this mirror's certificate")
	c.Flags().StringVar(&tlsc.Key, "tls-key", "", "this mirror's private key")
	return c
}

// serveAPI starts the management listener, with mutual TLS when it is
// configured.
//
// NextProtos is cleared from what tlsconf builds. That configuration is for the
// agent stream, which needs HTTP/2 for bidirectional streams and so pins "h2";
// this is an ordinary REST API, and pinning h2 would fail the handshake for
// every HTTP/1.1 client — curl included — with an error that says nothing about
// ALPN. Empty lets net/http offer both.
func serveAPI(srv *http.Server, c tlsconf.Config) error {
	if !c.Enabled() {
		return srv.ListenAndServe()
	}
	cfg, err := c.ServerConfig()
	if err != nil {
		return err
	}
	cfg.NextProtos = nil
	srv.TLSConfig = cfg
	return srv.ListenAndServeTLS("", "")
}

// readOnly refuses anything but a read.
//
// http.FileServer already ignores other methods, but saying so here means the
// property is stated rather than inherited — and a handler added later cannot
// quietly make the mirror writable.
func readOnly(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet && r.Method != http.MethodHead {
			http.Error(w, "the mirror is read-only", http.StatusMethodNotAllowed)
			return
		}
		h.ServeHTTP(w, r)
	})
}

// catalogSources reads every model's source from a running control plane.
func catalogSources(cmd *cobra.Command, server string) ([]string, error) {
	var models []adminapi.ModelView
	if err := newCtlClient(server).get(cmd.Context(), "/api/v1/models", &models); err != nil {
		return nil, err
	}
	out := make([]string, 0, len(models))
	for _, m := range models {
		out = append(out, m.Source)
	}
	return out, nil
}

// withManifest answers a request for a version with the list of files in it.
//
// A directory otherwise gets FileServer's HTML index, and something that has to
// know what a version consists of would be left parsing it. A model is
// routinely several files — sharded safetensors, or a GGUF split across parts —
// and a catalog entry names one source, so without a way to ask "what is in
// this version" a sharded model can be mirrored and never served.
//
// On the byte listener rather than the management API on purpose. Fetching
// weights is the one thing every node does, and making it depend on the
// administrative port would undo the split that keeps that port closeable.
func withManifest(root string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		model, version, ok := versionPath(r.URL.Path)
		if !ok {
			next.ServeHTTP(w, r)
			return
		}
		m, err := mirror.ManifestFor(root, model, version)
		switch {
		case errors.Is(err, mirror.ErrIncomplete):
			// Said, not hidden behind a 404: a node asking now is told why
			// it cannot have the version yet, and tries again later.
			w.Header().Set("Retry-After", "60")
			http.Error(w, err.Error(), http.StatusConflict)
			return
		case err != nil:
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(m)
	})
}

// versionPath matches "/model/version/" and nothing else.
//
// The trailing slash is the whole signal: "/m/v/file.gguf" is a file and
// "/m/v/" is the set it belongs to. Requiring it means a source URL says which
// of the two it means without a flag beside it.
func versionPath(p string) (model, version string, ok bool) {
	if !strings.HasSuffix(p, "/") {
		return "", "", false
	}
	parts := strings.Split(strings.Trim(p, "/"), "/")
	if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
		return "", "", false
	}
	return parts[0], parts[1], true
}

// parseAllow reads addresses and networks: 10.0.0.140, 10.0.0.0/24.
func parseAllow(list []string) ([]netip.Prefix, error) {
	var out []netip.Prefix
	for _, s := range list {
		for _, one := range strings.Split(s, ",") {
			one = strings.TrimSpace(one)
			if one == "" {
				continue
			}
			if p, err := netip.ParsePrefix(one); err == nil {
				out = append(out, p.Masked())
				continue
			}
			a, err := netip.ParseAddr(one)
			if err != nil {
				return nil, fmt.Errorf("--api-allow %q is not an address or a network", one)
			}
			out = append(out, netip.PrefixFrom(a.Unmap(), a.Unmap().BitLen()))
		}
	}
	return out, nil
}

// apiGate decides who reaches the management API. With mutual TLS the
// certificate decides, narrowed by --api-allow when given. Without it, only
// this machine and the addresses allowed: nothing allowed used to mean anyone,
// and every node that could reach the port could delete models.
func apiGate(mtls bool, allowed []netip.Prefix, next http.Handler) http.Handler {
	if mtls && len(allowed) == 0 {
		return next
	}
	return onlyFrom(allowed, next)
}

// onlyClients answers only certificates with one of these common names, when
// any are given. The handshake has already checked the certificate against
// the root; this narrows which of the root's certificates may manage.
func onlyClients(names []string, next http.Handler) http.Handler {
	if len(names) == 0 {
		return next
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.TLS == nil || len(r.TLS.VerifiedChains) == 0 ||
			!slices.Contains(names, r.TLS.VerifiedChains[0][0].Subject.CommonName) {
			http.Error(w, "this mirror is managed by its control plane only", http.StatusForbidden)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// onlyFrom answers only this machine and the addresses allowed: the management
// API can delete models and start downloads, and a mirror on its own machine
// is reached over the network by one control plane and no one else.
func onlyFrom(allowed []netip.Prefix, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		host, _, err := net.SplitHostPort(r.RemoteAddr)
		a, perr := netip.ParseAddr(host)
		if err != nil || perr != nil {
			http.Error(w, "forbidden", http.StatusForbidden)
			return
		}
		a = a.Unmap()
		ok := a.IsLoopback()
		for _, p := range allowed {
			ok = ok || p.Contains(a)
		}
		if !ok {
			http.Error(w, "this mirror is managed by its control plane only", http.StatusForbidden)
			return
		}
		next.ServeHTTP(w, r)
	})
}