llmgw/internal/bundle/bundle.go

// Package bundle carries what a cluster would otherwise fetch from the
// internet to a site that has no way to it.
//
// A bundle is made on a machine with internet access: llmgw's own release
// binaries with their signed checksums, the engine builds the nodes install
// with the checksums their publisher gave, and model files with the
// checksums their registry gave — each checked as it arrives, exactly as a
// node online would check it. Its manifest lists every file with its sha256,
// and the manifest's own sha256 is what the person carrying the bundle
// compares at the other end.
//
// Imported, the files go into the cluster's mirror: releases and engines
// under _llmgw/, which no model can be called, models where the mirror keeps
// models. The engines' checksums go to the control plane, which tells each
// node what every file it fetches from the mirror must be. The mirror serves
// plain HTTP to every node, so nothing a node runs is taken on the mirror's
// word: releases carry the project's signature, engines the checksum the
// control plane holds.
package bundle

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"path"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
)

// ManifestName is the manifest's file name in a bundle, and the first entry
// of a bundle packed as a tar.
const ManifestName = "manifest.json"

// Format is the manifest's layout version. A bundle of a later format is
// refused rather than half understood.
const Format = 1

// Prefix is where a mirror keeps what is not a model: _llmgw/releases and
// _llmgw/engines. A model name may not start with an underscore here, so
// none can be called this.
const Prefix = "_llmgw"

// Manifest says what a bundle holds.
type Manifest struct {
	Format  int       `json:"format"`
	Created time.Time `json:"created"`
	// MadeBy is the llmgw version that made it.
	MadeBy   string    `json:"made_by"`
	Releases []Release `json:"releases,omitempty"`
	Engines  []Engine  `json:"engines,omitempty"`
	Models   []Model   `json:"models,omitempty"`
	// Files are every file in the bundle but the manifest, with their
	// checksums: what an import checks each one against.
	Files []File `json:"files"`
}

// Release is llmgw's own release, for some architectures.
type Release struct {
	Version string   `json:"version"`
	Arches  []string `json:"arches"`
}

// Engine is an engine version's builds, for some architectures.
type Engine struct {
	Engine   string   `json:"engine"`
	Version  string   `json:"version"`
	Arches   []string `json:"arches"`
	Variants []string `json:"variants"`
}

// Model is a model's files, as the mirror will hold them.
type Model struct {
	Name    string `json:"name"`
	Version string `json:"version"`
	// Source is what it was asked for as: hf:owner/repo/file, or a URL.
	Source string `json:"source"`
	// Revision is the registry's commit it was taken at, when it has one.
	Revision string `json:"revision,omitempty"`
}

// File is one file of a bundle.
type File struct {
	// Path is where it is in the bundle: releases/<version>/<file>,
	// engines/<engine>/<version>/<file> or models/<model>/<version>/<file>.
	Path   string `json:"path"`
	Bytes  int64  `json:"bytes"`
	SHA256 string `json:"sha256"`
	// Source is where it was downloaded from.
	Source string `json:"source,omitempty"`
}

// Kinds of file, by the first element of their path.
const (
	kindRelease = "releases"
	kindEngine  = "engines"
	kindModel   = "models"
)

// Digest is a manifest's sha256: what `bundle create` prints and an import
// is given to check the bundle by.
func Digest(raw []byte) string {
	sum := sha256.Sum256(raw)
	return hex.EncodeToString(sum[:])
}

// parseManifest reads a manifest and checks every path in it, before any
// file is read: a path is where the file is written.
func parseManifest(raw []byte) (Manifest, error) {
	var m Manifest
	if err := json.Unmarshal(raw, &m); err != nil {
		return Manifest{}, fmt.Errorf("the bundle's manifest cannot be read: %w", err)
	}
	if m.Format != Format {
		return Manifest{}, fmt.Errorf("the bundle is of format %d and this llmgw reads format %d: "+
			"import it with the llmgw that made it", m.Format, Format)
	}
	seen := map[string]bool{}
	for _, f := range m.Files {
		if _, err := Destination(f.Path); err != nil {
			return Manifest{}, err
		}
		if len(f.SHA256) != 64 {
			return Manifest{}, fmt.Errorf("the manifest gives %s no sha256", f.Path)
		}
		if _, err := hex.DecodeString(f.SHA256); err != nil {
			return Manifest{}, fmt.Errorf("the manifest's sha256 of %s is not one", f.Path)
		}
		if seen[f.Path] {
			return Manifest{}, fmt.Errorf("the manifest lists %s twice", f.Path)
		}
		seen[f.Path] = true
	}
	return m, nil
}

// Destination is where a bundle's file goes in a mirror, relative to its
// root, or why it goes nowhere.
//
// Every path in a bundle has a fixed shape — three or four plain elements
// under one of three directories — because a path is where the import
// writes, as root. Anything else is refused, not cleaned.
func Destination(p string) (string, error) {
	parts := strings.Split(p, "/")
	bad := func(why string) (string, error) {
		return "", fmt.Errorf("the bundle names a file %q, which %s", p, why)
	}
	if path.Clean(p) != p || strings.HasPrefix(p, "/") {
		return bad("is not a plain path")
	}
	for _, e := range parts {
		if err := mirror.ValidName(e); err != nil {
			return bad("has an element that is not a plain name")
		}
	}
	switch parts[0] {
	case kindRelease:
		if len(parts) != 3 || !release.Tag.MatchString(parts[1]) {
			return bad("is not releases/<version>/<file>")
		}
		return Prefix + "/" + p, nil
	case kindEngine:
		if len(parts) != 4 {
			return bad("is not engines/<engine>/<version>/<file>")
		}
		return Prefix + "/" + p, nil
	case kindModel:
		if len(parts) != 4 {
			return bad("is not models/<model>/<version>/<file>")
		}
		if strings.HasPrefix(parts[1], "_") {
			return bad("names a model starting with an underscore, which the mirror keeps for itself")
		}
		return strings.Join(parts[1:], "/"), nil
	}
	return bad("is not under releases/, engines/ or models/")
}

// EngineBase is where an engine version's files are on a mirror whose files
// are served at mirrorURL.
func EngineBase(mirrorURL, engine, version string) string {
	return strings.TrimSuffix(mirrorURL, "/") + "/" + Prefix + "/" + kindEngine + "/" + engine + "/" + version
}

// ReleasesBase is where llmgw's releases are on a mirror whose files are
// served at mirrorURL: what `releases` is on an offline cluster.
func ReleasesBase(mirrorURL string) string {
	return strings.TrimSuffix(mirrorURL, "/") + "/" + Prefix + "/" + kindRelease
}

// ErrNoDigest is an import or trust given no sha256 to check the bundle by.
var ErrNoDigest = errors.New("give the bundle's sha256, as 'llmgw bundle create' printed it, " +
	"so what was carried is known to be what was made; or say it is to be taken unchecked")

// checkDigest compares a manifest with the digest it was given. Unchecked
// takes it as it is, said by the caller.
func checkDigest(raw []byte, want string, unchecked bool) error {
	want = strings.ToLower(strings.TrimSpace(want))
	switch {
	case want == "" && unchecked:
		return nil
	case want == "":
		return ErrNoDigest
	case Digest(raw) != want:
		return fmt.Errorf("the bundle's manifest is %s, not the %s it was given as: "+
			"it is not the bundle that was made, or the digest was mistyped", Digest(raw)[:16], want)
	}
	return nil
}