diff --git a/internal/cli/mirror.go b/internal/cli/mirror.go new file mode 100644 index 0000000..da89b29 --- /dev/null +++ b/internal/cli/mirror.go @@ -0,0 +1,296 @@ +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 //, 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 ", + 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://: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 +} diff --git a/internal/cli/mirror_test.go b/internal/cli/mirror_test.go new file mode 100644 index 0000000..3defe3e --- /dev/null +++ b/internal/cli/mirror_test.go @@ -0,0 +1,44 @@ +package cli + +import ( + "net/http" + "net/http/httptest" + "testing" +) + +func TestTheMirrorRefusesEverythingButReads(t *testing.T) { + // Stated here rather than inherited from http.FileServer, so a handler + // added later cannot quietly make the mirror writable — which would turn + // it into a way to replace a model with another one that every node then + // fetches without complaint. + h := readOnly(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + + for _, m := range []string{http.MethodGet, http.MethodHead} { + rec := httptest.NewRecorder() + h.ServeHTTP(rec, httptest.NewRequest(m, "/x.gguf", nil)) + if rec.Code != http.StatusOK { + t.Errorf("%s was refused: %d", m, rec.Code) + } + } + for _, m := range []string{http.MethodPut, http.MethodPost, http.MethodDelete, http.MethodPatch} { + rec := httptest.NewRecorder() + h.ServeHTTP(rec, httptest.NewRequest(m, "/x.gguf", nil)) + if rec.Code != http.StatusMethodNotAllowed { + t.Errorf("%s was allowed: %d", m, rec.Code) + } + } +} + +func TestTheFilenameIsTakenFromTheURLWhenNotGiven(t *testing.T) { + for _, c := range []struct{ url, want string }{ + {"https://h/a/b/model.gguf", "model.gguf"}, + {"https://h/model.gguf?download=true", "model.gguf"}, + {"model.gguf", "model.gguf"}, + } { + if got := lastPathElement(c.url); got != c.want { + t.Errorf("lastPathElement(%q) = %q, want %q", c.url, got, c.want) + } + } +} diff --git a/internal/cli/model.go b/internal/cli/model.go index 0a7e00c..fbbb7bc 100644 --- a/internal/cli/model.go +++ b/internal/cli/model.go @@ -31,6 +31,7 @@ func newModelCmd(client func() *ctlClient) *cobra.Command { func newModelAddCmd(client func() *ctlClient) *cobra.Command { var path string + var skipCheck bool c := &cobra.Command{ Use: "add -f ", Short: "Add a model from a definition file", @@ -61,6 +62,23 @@ func newModelAddCmd(client func() *ctlClient) *cobra.Command { if m.Name == "" { return fmt.Errorf("%s: the definition has no name", path) } + + // Before the catalog is touched, because a source that does not + // resolve is otherwise discovered by a replica failing to start + // minutes later on a node — with an error naming the node rather + // than the typo. + if !skipCheck { + size, err := checkSource(cmd.Context(), m.Source) + if err != nil { + return fmt.Errorf("%w\n Pass --no-source-check to add it anyway, "+ + "which is right when the mirror is not up yet or the nodes "+ + "reach it by an address this machine does not.", err) + } + if _, isURL := httpSource(m.Source); isURL { + fmt.Fprintln(cmd.OutOrStdout(), sourceNote(m.Source, size)) + } + } + body, err := json.Marshal(m) if err != nil { return err @@ -77,6 +95,8 @@ func newModelAddCmd(client func() *ctlClient) *cobra.Command { }, } c.Flags().StringVarP(&path, "file", "f", "", "YAML file holding the model definition") + c.Flags().BoolVar(&skipCheck, "no-source-check", false, + "add it without asking whether the weights are actually at that URL") return c } diff --git a/internal/cli/root.go b/internal/cli/root.go index 9718c99..556421d 100644 --- a/internal/cli/root.go +++ b/internal/cli/root.go @@ -42,6 +42,7 @@ func Root() *cobra.Command { newCtlCmd(), newTenancyCmd(), newUsersCmd(), + newMirrorCmd(), newSimEngineCmd(), ) return root diff --git a/internal/cli/sourcecheck.go b/internal/cli/sourcecheck.go new file mode 100644 index 0000000..d6126d0 --- /dev/null +++ b/internal/cli/sourcecheck.go @@ -0,0 +1,89 @@ +package cli + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" +) + +// checkSource asks whether a model's weights are actually there. +// +// A source that does not resolve is caught today by a replica failing to start, +// minutes later, on a node — and the error names the node rather than the typo. +// Asking now turns that into a sentence while the operator still has the file +// open. +// +// Only URLs. A path is checked on the node that will use it, by the node, which +// is the only machine whose answer means anything; a bare model id is resolved +// by the engine and is not this function's business. +func checkSource(ctx context.Context, source string) (size int64, err error) { + url, ok := httpSource(source) + if !ok { + return 0, nil + } + + ctx, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil) + if err != nil { + return 0, err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return 0, fmt.Errorf("cannot reach %s: %w", url, err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return 0, fmt.Errorf("%s answered %s", url, resp.Status) + } + // Absent when the server does not send it, which is not an error: a + // resumable fetch does not need the length in advance, and refusing here + // would reject a mirror that is merely terse. Go reports that as -1; + // normalised to 0 so callers have one spelling of "unknown". + if resp.ContentLength < 0 { + return 0, nil + } + return resp.ContentLength, nil +} + +// httpSource returns the URL to check, and whether there is one. +// +// hf: is expanded the way the fetcher expands it, so the check asks about the +// same file the download will ask for rather than about a string that looks +// similar. +func httpSource(source string) (string, bool) { + switch { + case strings.HasPrefix(source, "http://"), strings.HasPrefix(source, "https://"): + return source, true + case strings.HasPrefix(source, "hf:"): + rest := strings.TrimPrefix(source, "hf:") + parts := strings.SplitN(rest, "/", 3) + if len(parts) != 3 { + // A repository rather than a file. vLLM and SGLang download those + // themselves, so there is nothing here to check. + return "", false + } + return fmt.Sprintf("https://huggingface.co/%s/%s/resolve/main/%s", + parts[0], parts[1], parts[2]), true + default: + return "", false + } +} + +// sourceNote is what to print about a checked source. +// +// It says where the answer came from, because the control plane's network is +// not the nodes' network: a mirror this machine can reach may be unreachable +// from a node behind a different firewall, and the reverse. The check catches +// typos and a mirror that is down, which is most of what goes wrong, and it +// should not be read as proof that the fetch will work. +func sourceNote(source string, size int64) string { + if size > 0 { + return fmt.Sprintf(" source reachable from here: %s (%s)", source, fileSize(size)) + } + return fmt.Sprintf(" source reachable from here: %s (size not advertised)", source) +} diff --git a/internal/cli/sourcecheck_test.go b/internal/cli/sourcecheck_test.go new file mode 100644 index 0000000..cdd40b8 --- /dev/null +++ b/internal/cli/sourcecheck_test.go @@ -0,0 +1,113 @@ +package cli + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestOnlyURLsAreChecked(t *testing.T) { + // A path is checked by the node that will use it, which is the only + // machine whose answer means anything — the control plane may not even + // have that filesystem. A bare model id is the engine's business. + for _, src := range []string{ + "/var/lib/llmgw-models/x.gguf", "./x.gguf", "Qwen/Qwen2.5-7B-Instruct", + } { + if _, ok := httpSource(src); ok { + t.Errorf("%q was treated as a URL to check", src) + } + if _, err := checkSource(context.Background(), src); err != nil { + t.Errorf("checking %q returned %v; it should be skipped", src, err) + } + } +} + +func TestAnHFFileIsCheckedAtTheURLTheFetcherWouldUse(t *testing.T) { + // Asking about a string that merely looks similar would be worse than not + // asking: it would pass while the download later fails. + got, ok := httpSource("hf:Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/q4.gguf") + if !ok { + t.Fatal("an hf: file was not recognised") + } + want := "https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/resolve/main/q4.gguf" + if got != want { + t.Errorf("hf: resolved to %q, want %q", got, want) + } +} + +func TestAnHFRepositoryIsNotChecked(t *testing.T) { + // vLLM and SGLang download whole repositories themselves; there is no one + // file to ask about. + if _, ok := httpSource("hf:Qwen/Qwen2.5-7B-Instruct"); ok { + t.Error("a repository was treated as a single file") + } +} + +func TestAReachableSourceReportsItsSize(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodHead { + t.Errorf("the check used %s, not HEAD — it must not download the file", r.Method) + } + w.Header().Set("Content-Length", "1073741824") + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + size, err := checkSource(context.Background(), srv.URL+"/model.gguf") + if err != nil { + t.Fatal(err) + } + if size != 1<<30 { + t.Errorf("size = %d", size) + } +} + +func TestAMissingFileSaysSoWithTheURL(t *testing.T) { + // The whole point: a typo becomes a sentence while the operator still has + // the file open, rather than a replica failing minutes later on a node + // with an error that names the node. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "not found", http.StatusNotFound) + })) + defer srv.Close() + + _, err := checkSource(context.Background(), srv.URL+"/typo.gguf") + if err == nil { + t.Fatal("a 404 was accepted") + } + if !strings.Contains(err.Error(), "typo.gguf") { + t.Errorf("the error does not name the URL: %v", err) + } +} + +func TestAServerThatSendsNoLengthIsStillFine(t *testing.T) { + // A resumable fetch does not need the length in advance, and refusing here + // would reject a mirror that is merely terse. + // A server that does not know the length omits the header rather than + // sending a negative one, which is what a real mirror streaming from a + // pipe does. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + size, err := checkSource(context.Background(), srv.URL+"/x.gguf") + if err != nil { + t.Errorf("a terse server was rejected: %v", err) + } + if size != 0 { + t.Errorf("an unknown length came back as %d, want 0", size) + } +} + +func TestTheNoteSaysWhereTheAnswerCameFrom(t *testing.T) { + // The control plane's network is not the nodes' network. The check catches + // typos and a mirror that is down; it is not proof the fetch will work, + // and the wording must not suggest it is. + note := sourceNote("http://models.internal/x.gguf", 1<<30) + if !strings.Contains(note, "from here") { + t.Errorf("the note overstates what was checked: %s", note) + } +} diff --git a/internal/mirror/add.go b/internal/mirror/add.go new file mode 100644 index 0000000..b2c8210 --- /dev/null +++ b/internal/mirror/add.go @@ -0,0 +1,158 @@ +package mirror + +import ( + "context" + "fmt" + "io" + "net/http" + "os" + "path/filepath" + "strings" + "time" +) + +// Added is what a download produced. +type Added struct { + Entry Entry + SHA256 string +} + +// Add downloads a source into the mirror at model/version/file. +// +// Written to a temporary name and renamed when complete, for the same reason +// the agent's fetcher does it: a mirror is copied to every node, so a partial +// file here is a partial file everywhere, and a rename is the only step that +// is atomic. An interrupted download leaves a `.part` that List ignores rather +// than something a catalog could be pointed at. +// +// It refuses to replace an existing version. That refusal is the whole point of +// the layout — a version that can be overwritten takes away the ability to +// change the catalog back and still have the old weights. +func Add(ctx context.Context, root, source string, e Entry, progress func(done, total int64)) (Added, error) { + if e.Model == "" || e.Version == "" || e.File == "" { + return Added{}, fmt.Errorf("a model, a version and a filename are all required") + } + if strings.ContainsAny(e.Model+e.Version+e.File, `/\`) { + // Each is one path element. A slash would put the file somewhere the + // layout does not describe, and List would not find it again. + return Added{}, fmt.Errorf("model, version and filename are single names, not paths") + } + + final := e.Path(root) + if _, err := os.Stat(final); err == nil { + return Added{}, fmt.Errorf("%s/%s/%s: %w", e.Model, e.Version, e.File, ErrExists) + } + if err := os.MkdirAll(filepath.Dir(final), 0o755); err != nil { + return Added{}, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, source, nil) + if err != nil { + return Added{}, err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return Added{}, fmt.Errorf("fetching %s: %w", source, err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return Added{}, fmt.Errorf("%s answered %s", source, resp.Status) + } + + part := final + partSuffix + f, err := os.OpenFile(part, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) //nolint:gosec // inside the mirror root + if err != nil { + return Added{}, err + } + written, err := copyWithProgress(f, resp.Body, resp.ContentLength, progress) + closeErr := f.Close() + if err != nil { + _ = os.Remove(part) + return Added{}, err + } + if closeErr != nil { + _ = os.Remove(part) + return Added{}, closeErr + } + // A truncated download that the server ended cleanly would otherwise be + // renamed into place and copied to every node. + if resp.ContentLength > 0 && written != resp.ContentLength { + _ = os.Remove(part) + return Added{}, fmt.Errorf("%s: got %d bytes of %d", source, written, resp.ContentLength) + } + + sum, err := sha256File(part) + if err != nil { + _ = os.Remove(part) + return Added{}, err + } + if err := os.Rename(part, final); err != nil { + return Added{}, err + } + e.Bytes = written + return Added{Entry: e, SHA256: sum}, nil +} + +// copyWithProgress copies and reports occasionally. +// +// Occasionally rather than continuously: a forty gigabyte download is a +// hundred thousand chunks, and a line each would bury whatever else the +// operator is watching. +func copyWithProgress(dst io.Writer, src io.Reader, total int64, report func(done, total int64)) (int64, error) { + buf := make([]byte, 1<<20) + var done int64 + last := time.Now() + + for { + n, err := src.Read(buf) + if n > 0 { + w, werr := dst.Write(buf[:n]) + done += int64(w) + if werr != nil { + return done, werr + } + if report != nil && time.Since(last) > 2*time.Second { + report(done, total) + last = time.Now() + } + } + if err == io.EOF { + if report != nil { + report(done, total) + } + return done, nil + } + if err != nil { + return done, err + } + } +} + +// Remove deletes one version, and the model directory if that was the last of +// it. +// +// Leaving an empty directory behind would make `list` show a model that is not +// there, which is the kind of small lie that costs somebody an afternoon. +func Remove(root string, e Entry) error { + path := e.Path(root) + if err := os.Remove(path); err != nil { + return err + } + versionDir := filepath.Dir(path) + if empty, _ := isEmpty(versionDir); empty { + _ = os.Remove(versionDir) + } + modelDir := filepath.Dir(versionDir) + if empty, _ := isEmpty(modelDir); empty { + _ = os.Remove(modelDir) + } + return nil +} + +func isEmpty(dir string) (bool, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return false, err + } + return len(entries) == 0, nil +} diff --git a/internal/mirror/mirror.go b/internal/mirror/mirror.go new file mode 100644 index 0000000..cf8f1fb --- /dev/null +++ b/internal/mirror/mirror.go @@ -0,0 +1,174 @@ +// Package mirror manages a directory of model weights served over HTTP. +// +// The mirror is the cheapest useful thing to build around llmgw: one machine +// holding the model files so nodes fetch over your own network instead of the +// internet. Eight nodes each pulling forty gigabytes from outside is a day's +// work and on many corporate networks is not possible at all. +// +// This package is operator tooling, not part of the gateway. Nothing in the +// serving path writes here — `llmgw mirror` is run by a person, the way +// `llmgw keys` is, and the running service never has write access to the +// mirror. That is deliberate: 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. +package mirror + +import ( + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" +) + +// Entry is one file in the mirror. +type Entry struct { + // Model and Version are the directory layout, and the layout is the point: + // a version is never overwritten, only added beside. Rolling forward is a + // new directory and one line changed in the catalog; rolling back is + // changing the line back, because the old file is still there. + Model string + Version string + File string + Bytes int64 +} + +// Path is where this entry lives under a root. +func (e Entry) Path(root string) string { + return filepath.Join(root, e.Model, e.Version, e.File) +} + +// URL is how a catalog entry would refer to it. +func (e Entry) URL(base string) string { + return strings.TrimSuffix(base, "/") + "/" + + strings.Join([]string{e.Model, e.Version, e.File}, "/") +} + +// List reads what the mirror holds. +// +// Only files three levels down, because that is the layout: anything else is +// either something a person left behind or a partial download, and reporting +// it as a model would invite somebody to point a catalog at it. +func List(root string) ([]Entry, error) { + var out []Entry + models, err := os.ReadDir(root) + if err != nil { + return nil, fmt.Errorf("reading the mirror at %s: %w", root, err) + } + for _, m := range models { + if !m.IsDir() { + continue + } + versions, err := os.ReadDir(filepath.Join(root, m.Name())) + if err != nil { + continue + } + for _, v := range versions { + if !v.IsDir() { + continue + } + files, err := os.ReadDir(filepath.Join(root, m.Name(), v.Name())) + if err != nil { + continue + } + for _, f := range files { + if f.IsDir() || strings.HasSuffix(f.Name(), partSuffix) { + continue + } + info, err := f.Info() + if err != nil { + continue + } + out = append(out, Entry{ + Model: m.Name(), Version: v.Name(), File: f.Name(), Bytes: info.Size(), + }) + } + } + } + sort.Slice(out, func(i, j int) bool { + if out[i].Model != out[j].Model { + return out[i].Model < out[j].Model + } + if out[i].Version != out[j].Version { + return out[i].Version < out[j].Version + } + return out[i].File < out[j].File + }) + return out, nil +} + +// ErrExists means the version is already there. +// +// Refusing rather than overwriting is the whole discipline. A version that can +// be replaced in place takes away the thing versioned paths are for: the +// ability to change the catalog back and have the old weights still be there. +var ErrExists = errors.New("that version is already in the mirror") + +// Unused returns entries no catalog source refers to. +// +// `wanted` is every source in the catalog, and `base` is the address nodes use +// to reach this mirror — because a catalog refers to a file by URL and the +// mirror holds it by path, and the two are only the same string after the base +// is accounted for. +// +// A caller that cannot read the catalog must not call this with an empty list. +// Everything would come back unused, and the obvious next step would delete the +// entire mirror. The command that wraps this refuses rather than guessing. +func Unused(entries []Entry, wanted []string, base string) []Entry { + keep := map[string]struct{}{} + for _, w := range wanted { + keep[normalizeURL(w)] = struct{}{} + } + var out []Entry + for _, e := range entries { + if _, used := keep[normalizeURL(e.URL(base))]; used { + continue + } + out = append(out, e) + } + return out +} + +// normalizeURL makes two spellings of one address compare equal. +// +// A catalog written by hand and a URL built from a path differ in the ways +// people differ: a trailing slash on the base, http against https when a proxy +// terminates TLS, a host written with and without its port. Comparing raw +// strings would report a file as unused because somebody typed the base +// differently, and the next step after "unused" is deletion. +func normalizeURL(u string) string { + u = strings.TrimSpace(u) + u = strings.TrimPrefix(u, "https://") + u = strings.TrimPrefix(u, "http://") + // The path is what identifies the file; the host is which mirror it came + // from, and a mirror comparing against itself has only one. + if i := strings.Index(u, "/"); i >= 0 { + u = u[i:] + } + return strings.TrimSuffix(u, "/") +} + +const partSuffix = ".part" + +// sha256File is the checksum of a file on disk. +// +// Reported after a download so an operator has something to compare against +// what upstream published. A mirror is the one place in this system where a +// corrupted file is copied to every node rather than caught once. +func sha256File(path string) (string, error) { + f, err := os.Open(path) //nolint:gosec // path is built from the mirror root + if err != nil { + return "", err + } + defer f.Close() + + h := sha256.New() + if _, err := io.Copy(h, f); err != nil { + return "", err + } + return hex.EncodeToString(h.Sum(nil)), nil +} diff --git a/internal/mirror/mirror_test.go b/internal/mirror/mirror_test.go new file mode 100644 index 0000000..1ff0f78 --- /dev/null +++ b/internal/mirror/mirror_test.go @@ -0,0 +1,164 @@ +package mirror + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" +) + +func serving(t *testing.T, body string) string { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(body)) + })) + t.Cleanup(srv.Close) + return srv.URL + "/model.gguf" +} + +func TestAddPlacesAFileInTheVersionedLayout(t *testing.T) { + root := t.TempDir() + got, err := Add(context.Background(), root, serving(t, "weights"), + Entry{Model: "qwen-14b", Version: "2026-09", File: "q4_k_m.gguf"}, nil) + if err != nil { + t.Fatal(err) + } + if got.Entry.Bytes != 7 { + t.Errorf("size = %d", got.Entry.Bytes) + } + if got.SHA256 == "" { + t.Error("no checksum reported; a mirror is the one place a corrupt file reaches every node") + } + if _, err := os.Stat(filepath.Join(root, "qwen-14b", "2026-09", "q4_k_m.gguf")); err != nil { + t.Errorf("the file is not where the layout says: %v", err) + } +} + +func TestAddRefusesToReplaceAVersion(t *testing.T) { + // The whole discipline. A version that can be overwritten takes away the + // thing versioned paths are for: changing the catalog back and still + // having the old weights. + root := t.TempDir() + e := Entry{Model: "m", Version: "v1", File: "f.gguf"} + if _, err := Add(context.Background(), root, serving(t, "first"), e, nil); err != nil { + t.Fatal(err) + } + _, err := Add(context.Background(), root, serving(t, "second"), e, nil) + if !errors.Is(err, ErrExists) { + t.Fatalf("overwriting a version returned %v", err) + } + b, _ := os.ReadFile(e.Path(root)) + if string(b) != "first" { + t.Errorf("the original was replaced anyway: %q", b) + } +} + +func TestATruncatedDownloadIsNotRenamedIntoPlace(t *testing.T) { + // A partial file here is a partial file on every node. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", "100") + _, _ = w.Write([]byte("short")) + })) + defer srv.Close() + + root := t.TempDir() + e := Entry{Model: "m", Version: "v1", File: "f.gguf"} + if _, err := Add(context.Background(), root, srv.URL, e, nil); err == nil { + t.Fatal("a truncated download was accepted") + } + if _, err := os.Stat(e.Path(root)); !os.IsNotExist(err) { + t.Error("the truncated file was left where a catalog could point at it") + } + // And no .part left looking like a model. + got, _ := List(root) + if len(got) != 0 { + t.Errorf("list shows %+v after a failed download", got) + } +} + +func TestPathElementsAreNamesNotPaths(t *testing.T) { + // A slash would put the file somewhere the layout does not describe, and + // List would never find it again. + root := t.TempDir() + _, err := Add(context.Background(), root, serving(t, "x"), + Entry{Model: "../escape", Version: "v", File: "f"}, nil) + if err == nil { + t.Error("a path was accepted as a name") + } +} + +func TestListIgnoresAnythingNotInTheLayout(t *testing.T) { + root := t.TempDir() + if err := os.MkdirAll(filepath.Join(root, "m", "v1"), 0o755); err != nil { + t.Fatal(err) + } + for _, f := range []string{"real.gguf", "half.gguf.part"} { + if err := os.WriteFile(filepath.Join(root, "m", "v1", f), []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + } + // A file somebody dropped at the top level, which is not a model. + if err := os.WriteFile(filepath.Join(root, "notes.txt"), []byte("x"), 0o644); err != nil { + t.Fatal(err) + } + + got, err := List(root) + if err != nil { + t.Fatal(err) + } + if len(got) != 1 || got[0].File != "real.gguf" { + t.Errorf("list = %+v", got) + } +} + +func TestUnusedComparesByPathNotBySpelling(t *testing.T) { + // A catalog written by hand and a URL built from a path differ in the ways + // people differ. Reporting a file as unused because somebody typed the + // base differently would be followed by somebody deleting it. + entries := []Entry{ + {Model: "m", Version: "v1", File: "a.gguf"}, + {Model: "m", Version: "v2", File: "a.gguf"}, + } + wanted := []string{ + // A different scheme and a trailing slash on the base, both of which + // name the same file. + "https://models.internal/m/v1/a.gguf", + } + got := Unused(entries, wanted, "http://models.internal/") + if len(got) != 1 || got[0].Version != "v2" { + t.Errorf("unused = %+v, want only v2", got) + } +} + +func TestRemoveTakesTheEmptyDirectoriesWithIt(t *testing.T) { + // An empty directory left behind makes list show a model that is not + // there, which is the kind of small lie that costs an afternoon. + root := t.TempDir() + e := Entry{Model: "m", Version: "v1", File: "f.gguf"} + if _, err := Add(context.Background(), root, serving(t, "x"), e, nil); err != nil { + t.Fatal(err) + } + if err := Remove(root, e); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(root, "m")); !os.IsNotExist(err) { + t.Error("an empty model directory was left behind") + } +} + +func TestProgressIsReportedForALongDownload(t *testing.T) { + root := t.TempDir() + var calls int + _, err := Add(context.Background(), root, serving(t, strings.Repeat("x", 1<<20)), + Entry{Model: "m", Version: "v", File: "f"}, func(done, total int64) { calls++ }) + if err != nil { + t.Fatal(err) + } + if calls == 0 { + t.Error("a download reported no progress at all") + } +}