LocalAI/pkg/modelartifacts/materializer_test.go

package modelartifacts_test

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
	"github.com/mudler/LocalAI/pkg/modelartifacts"
)

type fakeSnapshotResolver struct {
	mu       sync.Mutex
	snapshot hfapi.Snapshot
	err      error
	calls    int
}

func (f *fakeSnapshotResolver) ResolveSnapshot(context.Context, hfapi.SnapshotRequest) (hfapi.Snapshot, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls++
	return f.snapshot, f.err
}

func (f *fakeSnapshotResolver) callCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.calls
}

var _ = Describe("controller artifact materializer", func() {
	It("requires an injected token before resolving a private artifact", func() {
		resolver := &fakeSnapshotResolver{}
		_, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), GinkgoT().TempDir(),
			modelartifacts.Spec{Source: modelartifacts.Source{
				Type: modelartifacts.SourceTypeHuggingFace, Repo: "owner/private", TokenEnv: modelartifacts.HuggingFaceTokenEnv,
			}})
		Expect(err).To(MatchError(ContainSubstring("non-empty HF_TOKEN")))
		Expect(resolver.callCount()).To(BeZero())
	})

	It("downloads, commits, and reuses a pinned snapshot", func() {
		weight := []byte("weight-bytes")
		sum := sha256.Sum256(weight)
		var requests atomic.Int32
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			requests.Add(1)
			Expect(r.Header.Get("Authorization")).To(Equal("Bearer hf-secret"))
			w.Header().Set("Content-Length", "12")
			_, _ = w.Write(weight)
		}))
		DeferCleanup(server.Close)

		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
			Files: []hfapi.SnapshotFile{{Path: "nested/model.safetensors", Size: 12, LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model"}},
		}}
		manager := modelartifacts.NewManager(resolver, modelartifacts.WithHuggingFaceToken("hf-secret"))
		modelsPath := GinkgoT().TempDir()
		spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", TokenEnv: "HF_TOKEN"}}

		result, err := manager.Ensure(context.Background(), modelsPath, spec)
		Expect(err).NotTo(HaveOccurred())
		Expect(result.CacheHit).To(BeFalse())
		Expect(result.Spec.Resolved.Revision).To(Equal("0123456789abcdef0123456789abcdef01234567"))
		// A snapshot with a single file records it as the primary file so the
		// load target is the file, not the snapshot directory.
		Expect(result.Spec.Resolved.PrimaryFile).To(Equal("nested/model.safetensors"))
		Expect(result.RelativePath).To(HavePrefix(".artifacts/huggingface/"))
		Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(result.RelativePath), "nested", "model.safetensors"))).To(Equal(weight))

		manifestBytes, err := os.ReadFile(filepath.Join(modelsPath, filepath.Dir(filepath.FromSlash(result.RelativePath)), "manifest.json"))
		Expect(err).NotTo(HaveOccurred())
		Expect(string(manifestBytes)).NotTo(ContainSubstring("hf-secret"))

		cached, err := manager.Ensure(context.Background(), modelsPath, result.Spec)
		Expect(err).NotTo(HaveOccurred())
		Expect(cached.CacheHit).To(BeTrue())
		Expect(requests.Load()).To(Equal(int32(1)))
		Expect(resolver.callCount()).To(Equal(1))
	})

	It("leaves the primary file unset for a multi-file snapshot", func() {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte("x")) }))
		DeferCleanup(server.Close)
		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
			Files: []hfapi.SnapshotFile{
				{Path: "config.json", Size: 1, URL: server.URL + "/config"},
				{Path: "model.safetensors", Size: 1, URL: server.URL + "/model"},
			},
		}}
		result, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), GinkgoT().TempDir(),
			modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
		Expect(err).NotTo(HaveOccurred())
		// Multi-file snapshots are consumed as a directory (e.g. transformers), so
		// no single file is promoted.
		Expect(result.Spec.Resolved.PrimaryFile).To(BeEmpty())
	})

	It("materializes a snapshot concurrently with ordered manifests and monotonic aggregate progress", func() {
		contents := map[string][]byte{
			"slow.bin": bytes.Repeat([]byte("s"), 64*1024),
			"fast.bin": bytes.Repeat([]byte("f"), 8*1024),
			"mid.bin":  bytes.Repeat([]byte("m"), 24*1024),
		}
		order := []string{"slow.bin", "fast.bin", "mid.bin"}
		var active atomic.Int32
		var maximum atomic.Int32
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			current := active.Add(1)
			defer active.Add(-1)
			for {
				previous := maximum.Load()
				if current <= previous || maximum.CompareAndSwap(previous, current) {
					break
				}
			}
			name := strings.TrimPrefix(r.URL.Path, "/")
			body := contents[name]
			w.Header().Set("Content-Length", strconv.Itoa(len(body)))
			for offset := 0; offset < len(body); offset += 4096 {
				end := min(offset+4096, len(body))
				_, _ = w.Write(body[offset:end])
				w.(http.Flusher).Flush()
				time.Sleep(time.Millisecond)
			}
		}))
		DeferCleanup(server.Close)

		files := make([]hfapi.SnapshotFile, 0, len(order))
		var total int64
		for _, name := range order {
			sum := sha256.Sum256(contents[name])
			files = append(files, hfapi.SnapshotFile{Path: name, Size: int64(len(contents[name])), LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/" + name})
			total += int64(len(contents[name]))
		}
		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", Files: files,
		}}
		var progressMu sync.Mutex
		var progress []int64
		ctx := modelartifacts.WithProgressSink(context.Background(), func(event modelartifacts.ProgressEvent) {
			if event.Phase == modelartifacts.PhaseDownloading || event.Phase == modelartifacts.PhaseVerifying {
				progressMu.Lock()
				progress = append(progress, event.CurrentBytes)
				progressMu.Unlock()
			}
		})
		result, err := modelartifacts.NewManager(resolver, modelartifacts.WithDownloadConcurrency(2)).Ensure(ctx, GinkgoT().TempDir(),
			modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
		Expect(err).NotTo(HaveOccurred())
		Expect(maximum.Load()).To(Equal(int32(2)))
		Expect(result.Manifest.Files).To(HaveLen(len(order)))
		for index, entry := range result.Manifest.Files {
			Expect(entry.Path).To(Equal(order[index]))
			sum := sha256.Sum256(contents[entry.Path])
			Expect(entry.SHA256).To(Equal(hex.EncodeToString(sum[:])))
		}
		progressMu.Lock()
		defer progressMu.Unlock()
		Expect(progress).NotTo(BeEmpty())
		for index, current := range progress {
			Expect(current).To(BeNumerically("<=", total))
			if index > 0 {
				Expect(current).To(BeNumerically(">=", progress[index-1]))
			}
		}
	})

	It("rejects a path escape before opening a destination", func() {
		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
			Files:            []hfapi.SnapshotFile{{Path: "../escape", Size: 1, URL: "https://example.invalid/file"}},
		}}
		modelsPath := GinkgoT().TempDir()
		_, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), modelsPath,
			modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
		Expect(err).To(MatchError(ContainSubstring("unsafe Hub path")))
		Expect(filepath.Join(modelsPath, "escape")).NotTo(BeAnExistingFile())
	})

	It("does not commit after cancellation", func() {
		started := make(chan struct{})
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			close(started)
			<-r.Context().Done()
		}))
		DeferCleanup(server.Close)
		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
			Files:            []hfapi.SnapshotFile{{Path: "model.bin", Size: 10, URL: server.URL}},
		}}
		ctx, cancel := context.WithCancel(context.Background())
		done := make(chan error, 1)
		modelsPath := GinkgoT().TempDir()
		go func() {
			_, err := modelartifacts.NewManager(resolver).Ensure(ctx, modelsPath,
				modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
			done <- err
		}()
		<-started
		cancel()
		Expect(<-done).To(MatchError(context.Canceled))
		entries, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", "huggingface", "*"))
		Expect(err).NotTo(HaveOccurred())
		Expect(entries).To(BeEmpty())
	})

	It("resumes an interrupted materialization without re-downloading completed files", func() {
		contents := map[string][]byte{
			"a/first.bin":  []byte("first-file-bytes"),
			"b/second.bin": []byte("second-file-bytes-longer"),
			"third.bin":    []byte("third"),
		}
		// order fixes both the download sequence and the expected manifest order.
		order := []string{"a/first.bin", "b/second.bin", "third.bin"}
		oid := func(b []byte) string { sum := sha256.Sum256(b); return hex.EncodeToString(sum[:]) }

		var mu sync.Mutex
		fetched := map[string]int{}
		// During run 1 this file 404s, interrupting the pass after the first file
		// has already completed and been promoted into the staging snapshot.
		failing := "b/second.bin"
		armed := true

		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			name := strings.TrimPrefix(r.URL.Path, "/file/")
			body, ok := contents[name]
			if !ok {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			mu.Lock()
			arm := armed
			mu.Unlock()
			if arm && name == failing {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			mu.Lock()
			fetched[name]++
			mu.Unlock()
			w.Header().Set("Content-Length", strconv.Itoa(len(body)))
			_, _ = w.Write(body)
		}))
		DeferCleanup(server.Close)

		files := make([]hfapi.SnapshotFile, 0, len(order))
		for _, p := range order {
			files = append(files, hfapi.SnapshotFile{
				Path: p, Size: int64(len(contents[p])), LFSOID: oid(contents[p]),
				URL: server.URL + "/file/" + p,
			})
		}
		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
			Files:            files,
		}}
		// A single manager (one writer identity) re-enters the same staging tree on
		// the second call, exactly as a resubmit against a still-idle partial does.
		manager := modelartifacts.NewManager(resolver)
		modelsPath := GinkgoT().TempDir()
		spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}

		_, err := manager.Ensure(context.Background(), modelsPath, spec)
		Expect(err).To(HaveOccurred())
		mu.Lock()
		Expect(fetched["a/first.bin"]).To(Equal(1), "the first file should have completed in run 1")
		armed = false
		fetched = map[string]int{}
		mu.Unlock()

		result, err := manager.Ensure(context.Background(), modelsPath, spec)
		Expect(err).NotTo(HaveOccurred())

		mu.Lock()
		defer mu.Unlock()
		// The decisive assertion: the already-completed file is NOT re-downloaded.
		Expect(fetched).NotTo(HaveKey("a/first.bin"))
		// Only the files missing after the interruption are fetched on resume.
		Expect(fetched).To(HaveKey("b/second.bin"))
		Expect(fetched).To(HaveKey("third.bin"))

		// The manifest lists every file, in order, with the correct size and hash.
		paths := make([]string, 0, len(result.Manifest.Files))
		byPath := map[string]modelartifacts.ManifestFile{}
		for _, f := range result.Manifest.Files {
			paths = append(paths, f.Path)
			byPath[f.Path] = f
		}
		Expect(paths).To(Equal(order))
		for _, p := range order {
			Expect(byPath[p].Size).To(Equal(int64(len(contents[p]))))
			Expect(byPath[p].SHA256).To(Equal(oid(contents[p])))
		}
		// Every file, skipped or freshly downloaded, is present on disk after commit.
		for _, p := range order {
			Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(result.RelativePath), filepath.FromSlash(p)))).To(Equal(contents[p]))
		}
	})

	It("emits resolving, downloading, verifying, and committing phases", func() {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte("x")) }))
		DeferCleanup(server.Close)
		resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
			Endpoint: "https://huggingface.co", Repo: "owner/repo",
			ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
			Files:            []hfapi.SnapshotFile{{Path: "config.json", Size: 1, URL: server.URL}},
		}}
		var phases []modelartifacts.Phase
		ctx := modelartifacts.WithProgressSink(context.Background(), func(event modelartifacts.ProgressEvent) {
			if len(phases) == 0 || phases[len(phases)-1] != event.Phase {
				phases = append(phases, event.Phase)
			}
		})
		_, err := modelartifacts.NewManager(resolver).Ensure(ctx, GinkgoT().TempDir(),
			modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
		Expect(err).NotTo(HaveOccurred())
		Expect(phases).To(Equal([]modelartifacts.Phase{
			modelartifacts.PhaseResolving, modelartifacts.PhaseDownloading, modelartifacts.PhaseVerifying, modelartifacts.PhaseCommitting,
		}))
	})
})