llmgw/internal/agent/weights/set.go

package weights

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// IsSet reports whether a source names a set of files rather than one.
//
// By the trailing slash, which is the whole signal: "…/v1/model.gguf" is a file
// and "…/v1/" is the version it belongs to. A flag beside the source would be a
// second place for the two to disagree, and the URL already distinguishes a
// directory from a file everywhere else on the web.
func IsSet(source string) bool {
	return isRemote(source) && strings.HasSuffix(source, "/")
}

// manifest is what the mirror answers for a version.
type manifest struct {
	Model   string `json:"model"`
	Version string `json:"version"`
	Files   []struct {
		Name   string `json:"name"`
		Bytes  int64  `json:"bytes"`
		SHA256 string `json:"sha256"`
	} `json:"files"`
	Bytes int64 `json:"bytes"`
}

// ensureSet downloads every file of a version into one directory.
//
// A directory, because that is what a multi-file model is: vLLM is handed the
// folder and reads the shards and the config out of it, and llama.cpp is handed
// the first part and finds the rest beside it. Downloading them into one place
// is what makes either possible, and a catalog entry naming one file could
// express neither.
//
// The whole set or nothing. A directory with three of four shards in it loads
// as a broken model rather than failing to load, so a partial set is left
// unfinished under its temporary name and the next attempt starts again.
func (f *Fetcher) ensureSet(ctx context.Context, source string, allowPickle bool) (string, error) {
	dir, err := f.cachePath(strings.TrimSuffix(source, "/"))
	if err != nil {
		return "", err
	}
	// A complete set is marked by a file written last. Checking for the
	// directory alone would treat an interrupted download as finished, which is
	// exactly the failure the marker exists to prevent.
	done := filepath.Join(dir, ".complete")
	if _, err := os.Stat(done); err == nil {
		// Already here, maybe fetched for a model that allowed pickles, or
		// before they were refused: looked at again, not trusted.
		if !allowPickle {
			entries, _ := os.ReadDir(dir)
			for _, e := range entries {
				if Pickled(e.Name()) {
					return "", refusePickle(source + e.Name())
				}
			}
		}
		return dir, nil
	}

	m, err := f.manifestFor(ctx, source)
	if err != nil {
		return "", err
	}
	if len(m.Files) == 0 {
		return "", fmt.Errorf("%s holds no files", source)
	}
	// Refused whole rather than skipped: an engine given the directory
	// loads the pickle when there is no safetensors beside it.
	if !allowPickle {
		for _, file := range m.Files {
			if Pickled(file.Name) {
				return "", refusePickle(source + file.Name)
			}
		}
	}

	if err := os.MkdirAll(dir, 0o755); err != nil {
		return "", err
	}
	// One number for the whole set, as the shards waiting on it see it: the
	// file being fetched was shown alone, a 46 GB part of 188.
	var setTotal, before int64
	for _, file := range m.Files {
		setTotal += file.Bytes
	}
	for _, file := range m.Files {
		// Rejected rather than joined: a name with a separator in it would
		// write outside the directory the mirror said it was in, and the
		// mirror is a place an operator copies files into by hand.
		if strings.ContainsAny(file.Name, `/\`) || file.Name == ".." {
			return "", fmt.Errorf("%s: the mirror offered a file named %q", source, file.Name)
		}
		target := filepath.Join(dir, file.Name)
		if st, err := os.Stat(target); err == nil && st.Size() == file.Bytes {
			before += file.Bytes
			continue
		}
		// Reported under the set's own address, because that is what the
		// shards are waiting on. Reusing fetch rather than writing a second
		// downloader keeps resume, the .part rename and the progress callback
		// in one place.
		if _, err := f.fetchAs(ctx, source+file.Name, source, target, setShare{before: before, total: setTotal}); err != nil {
			return "", err
		}
		before += file.Bytes
	}

	// Each file against the checksum the mirror took when it fetched it,
	// before the set is called complete: a shard damaged on the way loads as
	// a model that answers nonsense, not as an error.
	if err := f.checkSet(ctx, source, dir, m); err != nil {
		return "", err
	}

	if err := os.WriteFile(done, []byte(source+"\n"), 0o644); err != nil {
		return "", err
	}
	return dir, nil
}

// manifestFor asks the mirror what a version consists of.
func (f *Fetcher) manifestFor(ctx context.Context, source string) (manifest, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, source, nil)
	if err != nil {
		return manifest{}, err
	}
	client := f.Client
	if client == nil {
		client = &http.Client{}
	}
	resp, err := client.Do(req)
	if err != nil {
		return manifest{}, fmt.Errorf("asking %s what it holds: %w", source, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode == http.StatusConflict {
		// Still downloading on the mirror, or a shard missing: said in the
		// mirror's words, which name what is not there yet.
		why, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10))
		return manifest{}, fmt.Errorf("%s is not complete on the mirror yet: %s", source, strings.TrimSpace(string(why)))
	}
	if resp.StatusCode != http.StatusOK {
		return manifest{}, fmt.Errorf("%s answered %d", source, resp.StatusCode)
	}

	var m manifest
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return manifest{}, err
	}
	if err := json.Unmarshal(body, &m); err != nil {
		// The likely cause is a plain file server with a directory index, which
		// answers HTML. Said plainly, because "invalid character '<'" sends
		// somebody looking at the wrong thing.
		return manifest{}, fmt.Errorf(
			"%s did not answer with a manifest; a source ending in / has to be a "+
				"mirror version, not a directory listing", source)
	}
	return m, nil
}

func fileSHA256(path string) (string, error) {
	f, err := os.Open(path) //nolint:gosec // a file this fetcher wrote
	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
}

// checkSetAtOnce is how many of a set's files are checked at once: a set of
// six 50 GB files took eight minutes on one core, with nothing said.
const checkSetAtOnce = 4

// checkSet checks each file of a set against the checksum the mirror took
// when it fetched it, several at once, saying how far it has got.
func (f *Fetcher) checkSet(ctx context.Context, source, dir string, m manifest) error {
	var total int64
	for _, file := range m.Files {
		if file.SHA256 != "" {
			total += file.Bytes
		}
	}
	if total == 0 {
		return nil // nothing the mirror took a checksum of
	}
	var done atomic.Int64
	stop := make(chan struct{})
	defer close(stop)
	go func() {
		t := time.NewTicker(2 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-stop:
				return
			case <-t.C:
				f.report(Progress{Source: source, BytesDone: done.Load(), BytesTotal: total, Phase: PhaseChecking})
			}
		}
	}()
	f.report(Progress{Source: source, BytesTotal: total, Phase: PhaseChecking})

	var (
		mu    sync.Mutex
		first error
		wg    sync.WaitGroup
	)
	slots := make(chan struct{}, checkSetAtOnce)
	for _, file := range m.Files {
		if file.SHA256 == "" {
			continue
		}
		wg.Add(1)
		slots <- struct{}{}
		go func() {
			defer wg.Done()
			defer func() { <-slots }()
			target := filepath.Join(dir, file.Name)
			got, err := countedSHA256(ctx, target, &done)
			if err == nil && !strings.EqualFold(got, file.SHA256) {
				_ = os.Remove(target)
				err = fmt.Errorf("%s%s: sha256 %s, the mirror has %s; removed, to be fetched again",
					source, file.Name, got, file.SHA256)
			}
			if err != nil {
				mu.Lock()
				if first == nil {
					first = err
				}
				mu.Unlock()
			}
		}()
	}
	wg.Wait()
	if first == nil {
		f.report(Progress{Source: source, BytesDone: total, BytesTotal: total, Phase: PhaseChecking})
	}
	return first
}

// countedSHA256 is a file's checksum, adding what it has read to done.
func countedSHA256(ctx context.Context, path string, done *atomic.Int64) (string, error) {
	file, err := os.Open(path) //nolint:gosec // a file this fetcher wrote
	if err != nil {
		return "", err
	}
	defer file.Close()
	h := sha256.New()
	buf := make([]byte, 4<<20)
	for {
		if err := ctx.Err(); err != nil {
			return "", err
		}
		n, err := file.Read(buf)
		if n > 0 {
			h.Write(buf[:n])
			done.Add(int64(n))
		}
		if errors.Is(err, io.EOF) {
			return hex.EncodeToString(h.Sum(nil)), nil
		}
		if err != nil {
			return "", err
		}
	}
}