llmgw/internal/mirror/add.go

package mirror

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"syscall"
	"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
	}

	part := final + partSuffix
	written, err := downloadResuming(ctx, client, source, part, progress)
	if err != nil {
		return Added{}, err
	}

	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
}

// downloadAttempts is how many times a download is tried before it fails:
// a registry's CDN ends a long transfer now and then ("stream error …
// CANCEL"), and a 60 GB file failed at 23 GB and started again from nothing.
var downloadAttempts = 6

// downloadPause is the wait before the next attempt, doubled each time.
var downloadPause = 2 * time.Second

// downloadResuming fetches source into part, carrying on from what part
// already holds — a download cut short, or one a restarted mirror began —
// and trying again, from where it got to, when the transfer breaks. The part
// is kept when it fails, for the next fetch of the same file to carry on.
func downloadResuming(ctx context.Context, client *http.Client, source, part string,
	progress func(done, total int64),
) (int64, error) {
	var have int64
	if st, err := os.Stat(part); err == nil {
		have = st.Size()
	}
	pause := downloadPause
	var last error
	for attempt := 0; attempt < downloadAttempts; attempt++ {
		if attempt > 0 {
			select {
			case <-ctx.Done():
				return have, ctx.Err()
			case <-time.After(pause):
			}
			pause = min(pause*2, 30*time.Second)
		}
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, source, nil)
		if err != nil {
			return have, err
		}
		if have > 0 {
			req.Header.Set("Range", fmt.Sprintf("bytes=%d-", have))
		}
		resp, err := client.Do(req)
		if err != nil {
			if ctx.Err() != nil {
				return have, ctx.Err()
			}
			// A refused address, a name that does not resolve: asking
			// again changes nothing.
			if !transient(err) {
				return have, fmt.Errorf("fetching %s: %w", source, err)
			}
			last = fmt.Errorf("fetching %s: %w", source, err)
			continue
		}
		flags := os.O_CREATE | os.O_WRONLY
		switch {
		case resp.StatusCode == http.StatusPartialContent && have > 0:
			flags |= os.O_APPEND
		case resp.StatusCode == http.StatusOK:
			have = 0 // the server ignored the range: from the start
			flags |= os.O_TRUNC
		case resp.StatusCode == http.StatusRequestedRangeNotSatisfiable && have > 0:
			// Everything was there already.
			resp.Body.Close()
			return have, nil
		case resp.StatusCode >= 500 || resp.StatusCode == http.StatusTooManyRequests:
			resp.Body.Close()
			last = fmt.Errorf("%s answered %s", source, resp.Status)
			continue
		default:
			resp.Body.Close()
			return have, fmt.Errorf("%s answered %s", source, resp.Status)
		}
		want := resp.ContentLength
		f, err := os.OpenFile(part, flags, 0o644) //nolint:gosec // inside the mirror root
		if err != nil {
			resp.Body.Close()
			return have, err
		}
		offset := have
		var report func(done, total int64)
		if progress != nil {
			report = func(done, total int64) {
				if total > 0 {
					total += offset
				}
				progress(offset+done, total)
			}
		}
		n, err := copyWithProgress(f, resp.Body, want, report)
		resp.Body.Close()
		closeErr := f.Close()
		have += n
		if err == nil && closeErr != nil {
			return have, closeErr
		}
		if err != nil {
			if ctx.Err() != nil {
				return have, ctx.Err()
			}
			last = fmt.Errorf("fetching %s after %d bytes: %w", source, have, err)
			continue
		}
		// Ended cleanly but short: the rest is asked for again.
		if want > 0 && n != want {
			last = fmt.Errorf("%s: got %d bytes of %d", source, n, want)
			continue
		}
		return have, nil
	}
	return have, fmt.Errorf("%w (tried %d times; what arrived is kept and the next fetch carries on from it)",
		last, downloadAttempts)
}

// transient says whether a failure to reach a server may pass if asked
// again: a timeout, a connection the far end reset or ended, an HTTP/2
// stream it cancelled.
func transient(err error) bool {
	var ne net.Error
	if errors.As(err, &ne) && ne.Timeout() {
		return true
	}
	if errors.Is(err, syscall.ECONNRESET) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) {
		return true
	}
	msg := err.Error()
	return strings.Contains(msg, "stream error") || strings.Contains(msg, "connection reset") ||
		strings.Contains(msg, "GOAWAY")
}

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