llmgw/internal/mirror/mirror.go

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

// Manifest is what one version of a model consists of.
//
// A version is a directory and a model is routinely several files — sharded
// safetensors, or a GGUF split across parts. A catalog entry names one source,
// so without something that says "these files together are the model" a sharded
// model can be mirrored and never served.
type Manifest struct {
	Model   string         `json:"model"`
	Version string         `json:"version"`
	Files   []ManifestFile `json:"files"`
	Bytes   int64          `json:"bytes"`
}

// ManifestFile is one file of a version.
type ManifestFile struct {
	Name  string `json:"name"`
	Bytes int64  `json:"bytes"`
}

// ManifestFor lists the files of one model version.
//
// Built from the directory rather than from a file written alongside it. A
// manifest on disk is a second record of what is there, and the two disagree
// the first time somebody copies a file in by hand — which is a thing the
// mirror is explicitly meant to allow.
func ManifestFor(root, model, version string) (Manifest, error) {
	entries, err := List(root)
	if err != nil {
		return Manifest{}, err
	}
	m := Manifest{Model: model, Version: version}
	for _, e := range entries {
		if e.Model != model || e.Version != version {
			continue
		}
		m.Files = append(m.Files, ManifestFile{Name: e.File, Bytes: e.Bytes})
		m.Bytes += e.Bytes
	}
	if len(m.Files) == 0 {
		return Manifest{}, fmt.Errorf("%s/%s: %w", model, version, os.ErrNotExist)
	}
	return m, nil
}