llmgw/internal/bundle/create.go

package bundle

import (
	"archive/tar"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"slices"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

// EngineSpec is an engine version to carry, and which of its builds.
type EngineSpec struct {
	Engine, Version string
	// Variants are the builds to carry ("cpu", "cuda-13.3"); empty is every
	// one, which is what a site whose nodes are not all known wants.
	Variants []string
}

// ModelSpec is a model file to carry, and where the mirror will hold it.
type ModelSpec struct {
	Name, Version string
	// Source is hf:owner/repo/path/in/repo, or an http(s) URL.
	Source string
}

// CreateOptions say what goes into a bundle.
type CreateOptions struct {
	// Out is the directory the bundle is made in. A bundle made there
	// before is added to: a file already there with the right checksum is
	// not downloaded again.
	Out string
	// Release is llmgw's release to carry, v0.10.4; empty carries none.
	Release string
	// Arches are the architectures, of the releases and of the engines:
	// amd64, arm64.
	Arches []string
	// Releases is where release files are published.
	Releases string
	Engines  []EngineSpec
	Models   []ModelSpec

	// HTTP fetches everything; nil is a client with no overall timeout,
	// because a model file is tens of gigabytes.
	HTTP *http.Client
	// GitHubAPI, GitHubDownload and Registry default to GitHub and Hugging
	// Face; tests point them at fakes.
	GitHubAPI, GitHubDownload, Registry string
	// Log hears what is happening, a line at a time.
	Log func(format string, args ...any)
}

// Create downloads everything a bundle is asked to carry, checking each
// file as a node online would, and writes its manifest. It returns the
// manifest and its digest: the value to carry to the site separately and
// give to the import.
func Create(ctx context.Context, o CreateOptions) (Manifest, string, error) {
	c := &creator{o: o, client: o.HTTP}
	if c.client == nil {
		c.client = &http.Client{}
	}
	if len(o.Arches) == 0 {
		return Manifest{}, "", errors.New("give at least one architecture: amd64, arm64")
	}
	for _, a := range o.Arches {
		if a != "amd64" && a != "arm64" {
			return Manifest{}, "", fmt.Errorf("%q is not an architecture llmgw is built for: amd64 or arm64", a)
		}
	}
	if err := os.MkdirAll(o.Out, 0o755); err != nil {
		return Manifest{}, "", err
	}
	c.m = Manifest{Format: Format, Created: time.Now().UTC().Truncate(time.Second), MadeBy: version.Version}

	if o.Release != "" {
		if err := c.release(ctx); err != nil {
			return Manifest{}, "", err
		}
	}
	for _, e := range o.Engines {
		if err := c.engine(ctx, e); err != nil {
			return Manifest{}, "", err
		}
	}
	for _, m := range o.Models {
		if err := c.model(ctx, m); err != nil {
			return Manifest{}, "", err
		}
	}
	raw, err := json.MarshalIndent(c.m, "", "  ")
	if err != nil {
		return Manifest{}, "", err
	}
	raw = append(raw, '\n')
	if err := os.WriteFile(filepath.Join(o.Out, ManifestName), raw, 0o644); err != nil { //nolint:gosec // a list of checksums
		return Manifest{}, "", err
	}
	return c.m, Digest(raw), nil
}

type creator struct {
	o      CreateOptions
	client *http.Client
	m      Manifest
}

func (c *creator) log(format string, args ...any) {
	if c.o.Log != nil {
		c.o.Log(format, args...)
	}
}

// release carries llmgw's release: its checksums and their signature,
// checked against the key this build trusts, and a binary per
// architecture checked against them.
func (c *creator) release(ctx context.Context) error {
	v := c.o.Release
	base := c.o.Releases
	if base == "" {
		base = release.DefaultBase
	}
	c.log("llmgw %s: its checksums and their signature", v)
	sums, err := release.FetchChecksums(ctx, c.client, base, v)
	if err != nil {
		return fmt.Errorf("llmgw %s: %w", v, err)
	}
	dir := kindRelease + "/" + v + "/"
	if err := c.writeSmall(dir+"checksums.txt", sums.Text); err != nil {
		return err
	}
	if sums.Signature != "" {
		if err := c.writeSmall(dir+release.SignatureName, []byte(sums.Signature)); err != nil {
			return err
		}
	}
	root := strings.TrimSuffix(base, "/") + "/" + v + "/"
	for _, arch := range c.o.Arches {
		asset := release.AssetName(v, "linux", arch)
		want, err := sums.For(asset)
		if err != nil {
			return fmt.Errorf("llmgw %s: %w", v, err)
		}
		if err := c.fetch(ctx, root+asset, dir+asset, want); err != nil {
			return err
		}
	}
	c.m.Releases = append(c.m.Releases, Release{Version: v, Arches: slices.Clone(c.o.Arches)})
	return nil
}

// engine carries an engine version's builds for each architecture, checked
// against the checksums its publisher gives — the same the nodes ask for
// online.
func (c *creator) engine(ctx context.Context, e EngineSpec) error {
	in := &provision.Installer{HTTP: c.client, APIBase: c.o.GitHubAPI, DownloadBase: c.o.GitHubDownload}
	c.log("%s %s: the checksums its publisher gives", e.Engine, e.Version)
	sums, url, err := in.Published(ctx, e.Engine, e.Version)
	if err != nil {
		return fmt.Errorf("%s %s: %w", e.Engine, e.Version, err)
	}
	var variants []string
	done := map[string]bool{}
	for _, arch := range c.o.Arches {
		builds, err := provision.Builds(e.Engine, e.Version, arch)
		if err != nil {
			return err
		}
		var names []string
		for _, b := range builds {
			names = append(names, b.Variant)
			if len(e.Variants) > 0 && !slices.Contains(e.Variants, b.Variant) {
				continue
			}
			if !slices.Contains(variants, b.Variant) {
				variants = append(variants, b.Variant)
			}
			for _, asset := range b.Assets {
				if done[asset] {
					continue
				}
				want, ok := sums[asset]
				if !ok {
					return fmt.Errorf("%s %s publishes no checksum for %s, so it cannot be carried checked", e.Engine, e.Version, asset)
				}
				if err := c.fetch(ctx, url(asset), kindEngine+"/"+e.Engine+"/"+e.Version+"/"+asset, want); err != nil {
					return err
				}
				done[asset] = true
			}
		}
		for _, v := range e.Variants {
			if !slices.Contains(names, v) && !anyArchHas(e, v, c.o.Arches) {
				return fmt.Errorf("%s has no build %q; it has %s", e.Engine, v, strings.Join(names, ", "))
			}
		}
	}
	c.m.Engines = append(c.m.Engines, Engine{Engine: e.Engine, Version: e.Version,
		Arches: slices.Clone(c.o.Arches), Variants: variants})
	return nil
}

// anyArchHas says whether a variant is a build on any of the architectures:
// CUDA 12.8 is published for x64 alone, and asking for it with arm64 too is
// not a mistake.
func anyArchHas(e EngineSpec, variant string, arches []string) bool {
	for _, a := range arches {
		bs, _ := provision.Builds(e.Engine, e.Version, a)
		for _, b := range bs {
			if b.Variant == variant {
				return true
			}
		}
	}
	return false
}

// model carries a model file, checked against the checksum its registry
// publishes when it publishes one, at the commit the registry had when it
// was asked — so the bundle says exactly which weights it carries.
func (c *creator) model(ctx context.Context, s ModelSpec) error {
	for _, n := range []string{s.Name, s.Version} {
		if err := mirror.ValidName(n); err != nil {
			return fmt.Errorf("model %s/%s: %w", s.Name, s.Version, err)
		}
	}
	if strings.HasPrefix(s.Name, "_") {
		return fmt.Errorf("model %q: a name starting with an underscore is the mirror's own", s.Name)
	}
	var url, want, rev string
	switch {
	case strings.HasPrefix(s.Source, "hf:"):
		parts := strings.SplitN(strings.TrimPrefix(s.Source, "hf:"), "/", 3)
		if len(parts) != 3 || parts[2] == "" {
			return fmt.Errorf("%s names a repository, not a file: hf:owner/repo/file", s.Source)
		}
		repoID, file := parts[0]+"/"+parts[1], parts[2]
		reg := hf.New(c.o.Registry, time.Minute)
		repo, err := reg.Repo(ctx, repoID)
		if err != nil {
			return fmt.Errorf("%s: %w", s.Source, err)
		}
		found := false
		for _, f := range repo.Files {
			if f.Path == file {
				want, found = f.SHA256, true
			}
		}
		if !found {
			return fmt.Errorf("%s has no file %s at its latest commit", repoID, file)
		}
		url, rev = repo.DownloadURL(reg.Base(), file), repo.Revision
	case strings.HasPrefix(s.Source, "https://"), strings.HasPrefix(s.Source, "http://"):
		url = s.Source
	default:
		return fmt.Errorf("model source %q: hf:owner/repo/file or an http(s) URL", s.Source)
	}
	name := path.Base(strings.SplitN(url, "?", 2)[0])
	if err := mirror.ValidName(name); err != nil {
		return fmt.Errorf("%s: %w", s.Source, err)
	}
	if err := c.fetch(ctx, url, kindModel+"/"+s.Name+"/"+s.Version+"/"+name, want); err != nil {
		return err
	}
	c.m.Models = append(c.m.Models, Model{Name: s.Name, Version: s.Version, Source: s.Source, Revision: rev})
	return nil
}

// writeSmall puts a file the bundle makes itself, rather than downloads.
func (c *creator) writeSmall(rel string, b []byte) error {
	dst := filepath.Join(c.o.Out, filepath.FromSlash(rel))
	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return err
	}
	if err := os.WriteFile(dst, b, 0o644); err != nil { //nolint:gosec // published files
		return err
	}
	sum := sha256.Sum256(b)
	c.add(File{Path: rel, Bytes: int64(len(b)), SHA256: hex.EncodeToString(sum[:])})
	return nil
}

func (c *creator) add(f File) {
	for i := range c.m.Files {
		if c.m.Files[i].Path == f.Path {
			c.m.Files[i] = f
			return
		}
	}
	c.m.Files = append(c.m.Files, f)
}

// fetch downloads url into the bundle at rel, checked against want when
// there is one; a file already there that matches is kept.
func (c *creator) fetch(ctx context.Context, url, rel, want string) error {
	dst := filepath.Join(c.o.Out, filepath.FromSlash(rel))
	if want != "" {
		if got, err := fileSHA256(dst); err == nil && got == want {
			info, _ := os.Stat(dst)
			c.log("  have %s", rel)
			c.add(File{Path: rel, Bytes: info.Size(), SHA256: want, Source: url})
			return nil
		}
	}
	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return err
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return fmt.Errorf("fetching %s: %w", url, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("fetching %s: %s", url, resp.Status)
	}
	if resp.ContentLength > 0 {
		c.log("  fetching %s (%d MB)", rel, resp.ContentLength>>20)
	} else {
		c.log("  fetching %s", rel)
	}
	part := dst + partSuffix
	f, err := os.Create(part) //nolint:gosec // inside the bundle
	if err != nil {
		return err
	}
	h := sha256.New()
	n, err := io.Copy(io.MultiWriter(f, h), &progressReader{r: resp.Body, total: resp.ContentLength, log: c.log, name: rel})
	if cerr := f.Close(); err == nil {
		err = cerr
	}
	if err == nil && resp.ContentLength > 0 && n != resp.ContentLength {
		err = fmt.Errorf("got %d bytes of %d", n, resp.ContentLength)
	}
	if err != nil {
		_ = os.Remove(part)
		return fmt.Errorf("fetching %s: %w", url, err)
	}
	got := hex.EncodeToString(h.Sum(nil))
	if want != "" && got != want {
		_ = os.Remove(part)
		return fmt.Errorf("%s does not match the checksum its publisher gave (got %s, expected %s)", url, got[:12], want[:12])
	}
	if err := os.Rename(part, dst); err != nil {
		return err
	}
	c.add(File{Path: rel, Bytes: n, SHA256: got, Source: url})
	return nil
}

// progressReader says how far a long download has got, every few seconds.
type progressReader struct {
	r     io.Reader
	total int64
	done  int64
	last  time.Time
	log   func(string, ...any)
	name  string
}

func (p *progressReader) Read(b []byte) (int, error) {
	n, err := p.r.Read(b)
	p.done += int64(n)
	if p.total > 0 && time.Since(p.last) > 10*time.Second {
		if !p.last.IsZero() {
			p.log("    %d%% of %s", p.done*100/p.total, p.name)
		}
		p.last = time.Now()
	}
	return n, err
}

// Pack writes a bundle directory as one tar file, manifest first, so it can
// be carried as one file and imported as a stream. Weights do not compress,
// so it is not compressed.
func Pack(dir, out string) error {
	src, m, err := open(dir)
	if err != nil {
		return err
	}
	defer src.close()                     //nolint:errcheck // read only
	f, err := os.Create(out + partSuffix) //nolint:gosec // the operator's path
	if err != nil {
		return err
	}
	defer os.Remove(out + partSuffix) //nolint:errcheck // gone once renamed
	tw := tar.NewWriter(f)
	now := time.Now()
	raw := src.manifest()
	if err := tw.WriteHeader(&tar.Header{Name: ManifestName, Mode: 0o644, Size: int64(len(raw)), ModTime: now, Typeflag: tar.TypeReg}); err != nil {
		f.Close()
		return err
	}
	if _, err := tw.Write(raw); err != nil {
		f.Close()
		return err
	}
	for _, mf := range m.Files {
		in, err := os.Open(filepath.Join(dir, filepath.FromSlash(mf.Path))) //nolint:gosec // a checked path in the bundle
		if err != nil {
			f.Close()
			return err
		}
		err = tw.WriteHeader(&tar.Header{Name: mf.Path, Mode: 0o644, Size: mf.Bytes, ModTime: now, Typeflag: tar.TypeReg})
		if err == nil {
			_, err = io.Copy(tw, in)
		}
		in.Close()
		if err != nil {
			f.Close()
			return fmt.Errorf("packing %s: %w", mf.Path, err)
		}
	}
	if err := tw.Close(); err != nil {
		f.Close()
		return err
	}
	if err := f.Close(); err != nil {
		return err
	}
	return os.Rename(out+partSuffix, out)
}