llmgw/internal/agent/weights/set.go

package weights

import (
	"context"
	"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"`
	} `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) (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 {
		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)
	}

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

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