llmgw/internal/cli/mirror.go

package cli

import (
	"fmt"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

// 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),
	)
	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" +
			"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.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			source := args[0]
			file := as
			if file == "" {
				file = lastPathElement(source)
			}
			if model == "" || version == "" {
				return fmt.Errorf("--model and --version are required: they are the layout")
			}

			e := mirror.Entry{Model: model, Version: version, File: file}
			out := cmd.OutOrStdout()
			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, "\n%s/%s/%s  %s\n  sha256 %s\n",
				added.Entry.Model, added.Entry.Version, added.Entry.File,
				fileSize(added.Entry.Bytes), added.SHA256)
			fmt.Fprintf(out, "\nPoint a model at it with:\n  source: %s\n",
				added.Entry.URL("http://<this-host>:8081"))
			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; 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 string

	c := &cobra.Command{
		Use:   "serve",
		Short: "Serve the mirror over HTTP",
		Long: "Read-only by construction: this handler has no path that writes, so a\n" +
			"mirror cannot be used to replace a model with another one that every node\n" +
			"would then fetch without complaint.\n\n" +
			"Range requests are supported, which is what makes an interrupted download\n" +
			"on a node resume rather than start again.\n\n" +
			"nginx does this too, and better if you already run it. This exists so a\n" +
			"mirror is one command rather than another component to install.",
		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)
			}
			// http.Dir refuses paths that escape the root, which is the whole
			// traversal question answered by the standard library rather than
			// by me.
			srv := &http.Server{
				Addr:              listen,
				Handler:           readOnly(http.FileServer(http.Dir(*dir))),
				ReadHeaderTimeout: 10 * time.Second,
			}
			fmt.Fprintf(cmd.OutOrStdout(), "serving %s on %s\n", *dir, listen)
			return srv.ListenAndServe()
		},
	}
	c.Flags().StringVar(&listen, "listen", ":8081", "address to serve on")
	return c
}

// 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
}

// lastPathElement is the filename a URL ends in.
func lastPathElement(u string) string {
	u = strings.SplitN(u, "?", 2)[0]
	if i := strings.LastIndex(u, "/"); i >= 0 && i < len(u)-1 {
		return u[i+1:]
	}
	return u
}