llmgw/internal/agent/weights/set_test.go

package weights

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

// mirrorOf serves a manifest at "/m/v/" and the files beneath it.
func mirrorOf(t *testing.T, files map[string]string) (*Fetcher, *httptest.Server, *int) {
	t.Helper()
	hits := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/") {
			var b strings.Builder
			b.WriteString(`{"model":"m","version":"v","files":[`)
			first := true
			total := 0
			for name, body := range files {
				if !first {
					b.WriteString(",")
				}
				first = false
				fmt.Fprintf(&b, `{"name":%q,"bytes":%d}`, name, len(body))
				total += len(body)
			}
			fmt.Fprintf(&b, `],"bytes":%d}`, total)
			_, _ = w.Write([]byte(b.String()))
			return
		}
		body, ok := files[strings.TrimPrefix(r.URL.Path, "/m/v/")]
		if !ok {
			http.NotFound(w, r)
			return
		}
		hits++
		_, _ = w.Write([]byte(body))
	}))
	t.Cleanup(srv.Close)
	return &Fetcher{CacheDir: t.TempDir()}, srv, &hits
}

func TestASetIsFetchedIntoOneDirectory(t *testing.T) {
	f, srv, _ := mirrorOf(t, map[string]string{
		"model-00001-of-00002.safetensors": "aaaa",
		"model-00002-of-00002.safetensors": "bb",
		"config.json":                      "{}",
	})

	got, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
	if err != nil {
		t.Fatal(err)
	}
	st, err := os.Stat(got)
	if err != nil || !st.IsDir() {
		t.Fatalf("expected a directory, got %q: %v", got, err)
	}
	// A directory, because that is what a multi-file model is: vLLM reads the
	// shards and the config out of the folder, and a catalog entry naming one
	// file could express neither.
	for _, name := range []string{"model-00001-of-00002.safetensors", "config.json"} {
		if _, err := os.Stat(filepath.Join(got, name)); err != nil {
			t.Errorf("%s did not arrive: %v", name, err)
		}
	}
}

func TestAFinishedSetIsNotFetchedAgain(t *testing.T) {
	f, srv, hits := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})

	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
		t.Fatal(err)
	}
	after := *hits
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
		t.Fatal(err)
	}
	if *hits != after {
		t.Errorf("the second call downloaded again: %d then %d", after, *hits)
	}
}

func TestAHalfFetchedSetIsNotTreatedAsFinished(t *testing.T) {
	f, srv, _ := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})
	dir, err := f.cachePath(srv.URL + "/m/v")
	if err != nil {
		t.Fatal(err)
	}
	// One file present, no completion marker — an interrupted download.
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "a.gguf"), []byte("aaaa"), 0o644); err != nil {
		t.Fatal(err)
	}

	got, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
	if err != nil {
		t.Fatal(err)
	}
	// A directory with three of four shards loads as a broken model rather
	// than failing to load, so the marker is what says the set is whole.
	if _, err := os.Stat(filepath.Join(got, "b.gguf")); err != nil {
		t.Errorf("the missing file was not fetched: %v", err)
	}
}

func TestAFileAlreadyThereAtTheRightSizeIsNotRefetched(t *testing.T) {
	f, srv, hits := mirrorOf(t, map[string]string{"a.gguf": "aaaa", "b.gguf": "bb"})
	dir, _ := f.cachePath(srv.URL + "/m/v")
	_ = os.MkdirAll(dir, 0o755)
	_ = os.WriteFile(filepath.Join(dir, "a.gguf"), []byte("aaaa"), 0o644)

	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
		t.Fatal(err)
	}
	// Resuming a set means not paying for the shards that already landed.
	if *hits != 1 {
		t.Errorf("expected only the missing file to be fetched, got %d downloads", *hits)
	}
}

func TestAFilenameThatEscapesTheDirectoryIsRefused(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"files":[{"name":"../../etc/passwd","bytes":1}]}`))
	}))
	defer srv.Close()
	f := &Fetcher{CacheDir: t.TempDir()}

	_, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
	// The mirror is a place an operator copies files into by hand, so what it
	// reports is not automatically safe to join onto a path.
	if err == nil || !strings.Contains(err.Error(), "offered a file named") {
		t.Errorf("expected the name to be refused, got %v", err)
	}
}

func TestADirectoryListingIsNotAManifest(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte("<pre><a href=\"a.gguf\">a.gguf</a></pre>"))
	}))
	defer srv.Close()
	f := &Fetcher{CacheDir: t.TempDir()}

	_, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
	// "invalid character '<'" sends somebody looking at the wrong thing.
	if err == nil || !strings.Contains(err.Error(), "has to be a mirror version") {
		t.Errorf("the reason should name the cause, got %v", err)
	}
}

func TestOnlyATrailingSlashMeansASet(t *testing.T) {
	for src, want := range map[string]bool{
		"http://mirror/m/v/":           true,
		"http://mirror/m/v/model.gguf": false,
		"/var/lib/models/":             false, // a local directory is not fetched
		"Qwen/Qwen2.5-7B-Instruct":     false, // a repo id, resolved by the engine
	} {
		if got := IsSet(src); got != want {
			t.Errorf("IsSet(%q) = %v, want %v", src, got, want)
		}
	}
}

func TestASetFileThatDoesNotMatchTheMirrorsChecksumIsRefused(t *testing.T) {
	good := "the shard as mirrored"
	sum := sha256.Sum256([]byte(good))
	served := "the shard damaged on the way"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/") {
			fmt.Fprintf(w, `{"model":"m","version":"v","files":[{"name":"s.safetensors","bytes":%d,"sha256":%q}]}`,
				len(served), hex.EncodeToString(sum[:]))
			return
		}
		_, _ = w.Write([]byte(served))
	}))
	t.Cleanup(srv.Close)
	f := &Fetcher{CacheDir: t.TempDir()}
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err == nil || !strings.Contains(err.Error(), "sha256") {
		t.Fatalf("a damaged shard was accepted: %v", err)
	}
	served = good
	// The size differs, so it is fetched again; this time it matches.
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
		t.Fatalf("the shard as mirrored: %v", err)
	}
}

func TestASetStillDownloadingOnTheMirrorSaysSo(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		http.Error(w, "m/v: this version is not complete: still downloading model-00002-of-00002.safetensors", http.StatusConflict)
	}))
	t.Cleanup(srv.Close)
	f := &Fetcher{CacheDir: t.TempDir()}
	_, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false)
	if err == nil || !strings.Contains(err.Error(), "not complete on the mirror yet") ||
		!strings.Contains(err.Error(), "model-00002-of-00002") {
		t.Fatalf("got %v", err)
	}
}

// A set's download is one number climbing against the whole set, not each
// part's own: a 46 GB part of 188 was shown as the download.
func TestASetsProgressIsTheWholeSets(t *testing.T) {
	f, srv, _ := mirrorOf(t, map[string]string{
		"m-00001-of-00002.gguf": strings.Repeat("a", 600),
		"m-00002-of-00002.gguf": strings.Repeat("b", 400),
	})
	var seen []Progress
	f.OnProgress = func(p Progress) { seen = append(seen, p) }
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
		t.Fatal(err)
	}
	if len(seen) == 0 {
		t.Fatal("no progress reported")
	}
	last := int64(-1)
	for _, p := range seen {
		if p.BytesTotal != 1000 {
			t.Fatalf("reported against %d, not the set's 1000: %+v", p.BytesTotal, seen)
		}
		if p.BytesDone < last {
			t.Fatalf("went back: %+v", seen)
		}
		last = p.BytesDone
	}
	if last != 1000 {
		t.Errorf("ended at %d", last)
	}
}

// A set's files checked after they arrive say so, as far as they have got,
// rather than leaving the download at 100% for minutes.
func TestASetBeingCheckedSaysSo(t *testing.T) {
	parts := map[string]string{"a.gguf": strings.Repeat("a", 300), "b.gguf": strings.Repeat("b", 200)}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/") {
			var files []string
			for n, body := range parts {
				sum := sha256.Sum256([]byte(body))
				files = append(files, fmt.Sprintf(`{"name":%q,"bytes":%d,"sha256":%q}`, n, len(body), hex.EncodeToString(sum[:])))
			}
			fmt.Fprintf(w, `{"model":"m","version":"v","files":[%s]}`, strings.Join(files, ","))
			return
		}
		_, _ = w.Write([]byte(parts[strings.TrimPrefix(r.URL.Path, "/m/v/")]))
	}))
	t.Cleanup(srv.Close)
	var mu sync.Mutex
	var checking []Progress
	f := &Fetcher{CacheDir: t.TempDir(), OnProgress: func(p Progress) {
		if p.Phase == PhaseChecking {
			mu.Lock()
			checking = append(checking, p)
			mu.Unlock()
		}
	}}
	if _, err := f.Ensure(t.Context(), srv.URL+"/m/v/", false); err != nil {
		t.Fatal(err)
	}
	if len(checking) < 2 || checking[0].BytesDone != 0 || checking[len(checking)-1].BytesDone != 500 ||
		checking[len(checking)-1].BytesTotal != 500 {
		t.Errorf("%+v", checking)
	}
}

// Four replicas on one node's four cards start together and want the same
// set: it is fetched once, and every one of them gets it. Each fetching it
// alone wrote the same .part files over one another, and three failed.
func TestASetWantedByFourReplicasAtOnceIsFetchedOnce(t *testing.T) {
	body := strings.Repeat("a", 1<<20)
	var hits atomic.Int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasSuffix(r.URL.Path, "/") {
			fmt.Fprintf(w, `{"model":"m","version":"v","files":[{"name":"a.safetensors","bytes":%d},{"name":"b.safetensors","bytes":%d}]}`,
				len(body), len(body))
			return
		}
		hits.Add(1)
		time.Sleep(100 * time.Millisecond) // long enough for all four to ask
		_, _ = w.Write([]byte(body))
	}))
	t.Cleanup(srv.Close)
	f := &Fetcher{CacheDir: t.TempDir()}

	var wg sync.WaitGroup
	errs := make([]error, 4)
	for i := range errs {
		wg.Go(func() { _, errs[i] = f.Ensure(t.Context(), srv.URL+"/m/v/", false) })
	}
	wg.Wait()
	for i, err := range errs {
		if err != nil {
			t.Errorf("replica %d: %v", i, err)
		}
	}
	if n := hits.Load(); n != 2 {
		t.Errorf("%d downloads of a set of two files", n)
	}
}