llmgw/internal/mirror/add.go

package mirror

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"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) {
	return addWith(ctx, http.DefaultClient, root, source, e, progress)
}

// addWith is Add through a given client. The management API passes one that
// will not reach private addresses; `llmgw mirror add`, run by somebody on
// the mirror, passes the ordinary one.
func addWith(ctx context.Context, client *http.Client, 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")
	}
	for _, n := range []string{e.Model, e.Version, e.File} {
		if err := ValidName(n); err != nil {
			return Added{}, err
		}
	}

	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 := client.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
	}
	// The checksum beside the file, before the file is in place: a manifest
	// listing the file then always has its checksum to hand on, and a node
	// can tell a file damaged on the way from the one that was mirrored.
	if err := writeChecksum(final, sum); err != nil {
		_ = os.Remove(part)
		return Added{}, err
	}
	if err := writeSource(final, source); 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
	}
	// Its checksum goes with it; left, it would keep the directory from
	// being empty, and so from going.
	_ = os.Remove(checksumPath(path))
	_ = os.Remove(sourcePath(path))
	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
}