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"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"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
}

// ValidName checks one element of an entry — a model, version or file name.
//
// Each is one path element, and a name that is not one writes somewhere the
// layout does not describe. Rejecting slashes was not enough: ".." is a
// single element too, and a model called ".." put its file beside the mirror
// root rather than inside it. A name starting with a dot is refused as well,
// because that is where the mirror keeps its own bookkeeping.
func ValidName(n string) error {
	if n == "" || n[0] == '.' || strings.ContainsAny(n, "/\\\x00") {
		return fmt.Errorf("%q is not a usable name: it must be a single path element "+
			"that does not start with a dot", n)
	}
	return nil
}

// 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) || strings.HasPrefix(f.Name(), ".") {
					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"

// checksumPath is where a file's sha256 is kept: beside it, hidden, so no
// listing takes it for a file of the model and no name a person gives a
// file (which may not start with a dot) can be the same.
func checksumPath(file string) string {
	return filepath.Join(filepath.Dir(file), "."+filepath.Base(file)+".sha256")
}

func writeChecksum(file, sum string) error {
	tmp := checksumPath(file) + partSuffix
	if err := os.WriteFile(tmp, []byte(sum+"\n"), 0o644); err != nil { //nolint:gosec // a checksum, for anyone to read
		return err
	}
	return os.Rename(tmp, checksumPath(file))
}

// sourcePath is where the URL a file was fetched from is kept, beside it like
// its checksum: what says, after the mirror restarts, which repository and
// commit a file is, so a newer one there can be noticed.
func sourcePath(file string) string {
	return filepath.Join(filepath.Dir(file), "."+filepath.Base(file)+".source")
}

func writeSource(file, source string) error {
	tmp := sourcePath(file) + partSuffix
	if err := os.WriteFile(tmp, []byte(source+"\n"), 0o644); err != nil { //nolint:gosec // a URL, for anyone to read
		return err
	}
	return os.Rename(tmp, sourcePath(file))
}

// Record keeps a file's checksum and source beside it, as a download into
// the mirror does: for a file put in place another way — an offline bundle's
// — that should be offered to nodes with its checksum all the same.
func Record(file, sha256, source string) error {
	if err := writeChecksum(file, sha256); err != nil {
		return err
	}
	if source == "" {
		return nil
	}
	return writeSource(file, source)
}

// Recorded is a file's checksum as kept beside it, "" for none.
func Recorded(file string) string { return readChecksum(file) }

// readSource is where a file was fetched from, or "" for one fetched before
// this was kept, or put in place by hand.
func readSource(file string) string {
	b, err := os.ReadFile(sourcePath(file)) //nolint:gosec // beside a file in the mirror
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// readChecksum is a file's recorded sha256, or "" for a file put in place by
// hand, which has none.
func readChecksum(file string) string {
	b, err := os.ReadFile(checksumPath(file)) //nolint:gosec // beside a file in the mirror
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// 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"`
	// SHA256 is the file's checksum when the mirror fetched it, for a node
	// to check what it downloaded; empty for a file put there by hand.
	SHA256 string `json:"sha256,omitempty"`
}

// ErrIncomplete is a version not yet whole: a download still running, or a
// shard its own index names missing. Its manifest is not given out, so no
// node starts a model from part of it.
var ErrIncomplete = errors.New("this version is not complete")

// 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}
	dir := filepath.Join(root, model, 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,
			SHA256: readChecksum(filepath.Join(dir, e.File))})
		m.Bytes += e.Bytes
	}
	if len(m.Files) == 0 {
		return Manifest{}, fmt.Errorf("%s/%s: %w", model, version, os.ErrNotExist)
	}
	if err := complete(dir, m.Files); err != nil {
		return Manifest{}, fmt.Errorf("%s/%s: %w: %v", model, version, ErrIncomplete, err)
	}
	return m, nil
}

// complete says whether a version's files are the whole of it: nothing still
// downloading, every shard a safetensors index maps a tensor to, and every
// part of a split ("-00002-of-00004").
//
// From what the files say of each other rather than from a list written when
// the fetch began, so a version filled by hand is held to the same.
func complete(dir string, files []ManifestFile) error {
	if parts, _ := filepath.Glob(filepath.Join(dir, "*"+partSuffix)); len(parts) > 0 {
		var names []string
		for _, p := range parts {
			if b := filepath.Base(p); !strings.HasPrefix(b, ".") {
				names = append(names, strings.TrimSuffix(b, partSuffix))
			}
		}
		if len(names) > 0 {
			return fmt.Errorf("still downloading %s", strings.Join(names, ", "))
		}
	}
	have := map[string]bool{}
	for _, f := range files {
		have[f.Name] = true
	}
	var missing []string
	want := map[string]bool{}
	for _, f := range files {
		if strings.HasSuffix(f.Name, ".safetensors.index.json") {
			var idx struct {
				WeightMap map[string]string `json:"weight_map"`
			}
			b, err := os.ReadFile(filepath.Join(dir, f.Name)) //nolint:gosec // a file in the mirror
			if err != nil || json.Unmarshal(b, &idx) != nil {
				return fmt.Errorf("%s cannot be read", f.Name)
			}
			for _, shard := range idx.WeightMap {
				want[shard] = true
			}
		}
		if m := splitPart.FindStringSubmatch(f.Name); m != nil {
			total, _ := strconv.Atoi(m[3])
			for i := 1; i <= total; i++ {
				want[fmt.Sprintf("%s-%0*d-of-%s%s", m[1], len(m[2]), i, m[3], m[4])] = true
			}
		}
	}
	for name := range want {
		if !have[name] {
			missing = append(missing, name)
		}
	}
	if len(missing) > 0 {
		sort.Strings(missing)
		return fmt.Errorf("missing %s", strings.Join(missing, ", "))
	}
	return nil
}

// splitPart matches one part of a split: prefix, part, total, extension.
var splitPart = regexp.MustCompile(`^(.+)-(\d{5})-of-(\d{5})(\.[A-Za-z0-9]+)$`)