llmgw/internal/agent/weights/set.go

package weights

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

// 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
	}
	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 {
			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); err != nil {
			return "", err
		}
	}

	// 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.
	for _, file := range m.Files {
		if file.SHA256 == "" {
			continue
		}
		target := filepath.Join(dir, file.Name)
		got, err := fileSHA256(target)
		if err != nil {
			return "", err
		}
		if !strings.EqualFold(got, file.SHA256) {
			_ = os.Remove(target)
			return "", fmt.Errorf("%s%s: sha256 %s, the mirror has %s; removed, to be fetched again",
				source, file.Name, got, file.SHA256)
		}
	}

	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
}