diff --git a/pkg/modelartifacts/materializer.go b/pkg/modelartifacts/materializer.go index 16b10408c..8f1f03192 100644 --- a/pkg/modelartifacts/materializer.go +++ b/pkg/modelartifacts/materializer.go @@ -370,18 +370,42 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec totalBytes += file.Size } - manifest := Manifest{Version: ManifestVersion, Artifact: spec, Files: make([]ManifestFile, 0, len(snapshot.Files))} + // Files land in manifest.Files at their snapshot index, not in completion + // order, so a mix of skipped and freshly downloaded files still records the + // manifest in the resolved snapshot's order. committedResult and staging both + // read this manifest, and getting its order or contents wrong would make a + // corrupt tree look valid. + manifest := Manifest{Version: ManifestVersion, Artifact: spec, Files: make([]ManifestFile, len(snapshot.Files))} completedBytes := int64(0) + skippedFiles := 0 + skippedBytes := int64(0) tasks := make([]downloader.FileTask, 0, len(snapshot.Files)) for index, file := range snapshot.Files { if err := ctx.Err(); err != nil { return Result{}, err } + file := file + taskIndex := index + // A file already present and verified in this staging tree survives a + // restart: an interrupted pass promotes each completed file into + // snapshot/ before it moves on, so on re-entry (a resubmit, a controller + // roll, an adopted orphan) we must resume past it rather than re-fetch + // tens of gigabytes from scratch. Verification reuses the exact happy-path + // check so the recorded manifest entry is byte-for-byte identical to the + // one a fresh download would have produced; a file that fails it falls + // through to a normal re-download. + snapshotRel := path.Join("snapshot", file.Path) + snapshotAbs := filepath.Join(layout.Partial, filepath.FromSlash(snapshotRel)) + if entry, ok := reuseMaterializedFile(snapshotAbs, file); ok { + manifest.Files[taskIndex] = entry + completedBytes += file.Size + skippedFiles++ + skippedBytes += file.Size + continue + } nameSum := sha256.Sum256([]byte(file.Path)) blobRel := path.Join(".downloads", hex.EncodeToString(nameSum[:])) blobAbs := filepath.Join(layout.Partial, filepath.FromSlash(blobRel)) - file := file - taskIndex := index task := downloader.FileTask{ URI: downloader.URI(file.URL), Destination: blobAbs, @@ -421,17 +445,32 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec if err := root.MkdirAll(path.Dir(destination), 0o750); err != nil { return err } + // A freshly downloaded file replaces whatever sits at the + // destination (a stale or unverifiable leftover); a file we chose + // to keep never reaches this path, so the removal only ever + // discards bytes we are about to overwrite. _ = root.Remove(destination) if err := root.Rename(blobRel, destination); err != nil { return err } - manifest.Files = append(manifest.Files, entry) + manifest.Files[taskIndex] = entry completedBytes += file.Size return nil }, } tasks = append(tasks, task) } + // Surface resume at INFO: the absence of this signal is part of what made a + // never-converging download invisible in production, where each restart + // silently re-fetched every completed file. + if skippedFiles > 0 { + xlog.Info("resuming artifact materialization; keeping already-completed files", + "artifact", spec.Name, + "skipped_files", skippedFiles, + "skipped_bytes", skippedBytes, + "remaining_files", len(tasks), + "total_files", len(snapshot.Files)) + } if err := downloader.DownloadFilesWithContext(ctx, tasks, nil); err != nil { return Result{}, err } @@ -487,6 +526,34 @@ func (m *Manager) commit(modelsPath string, spec Spec, layout Layout, manifest M return Result{Spec: spec, RelativePath: relative, Manifest: manifest}, nil } +// reuseMaterializedFile reports whether a file already staged in this tree's +// snapshot/ can be kept as-is, returning the manifest entry it should +// contribute. It is the resume counterpart to the download path: a completed +// file is promoted into snapshot/ before the pass moves on, so on re-entry we +// verify what is there and skip the fetch instead of restarting from the first +// shard. +// +// Verification is a full re-hash via the same verifyDownloadedFile the happy +// path uses, not a size-only check. The manifest requires a SHA-256 for every +// file, and a non-LFS file carries no precomputed SHA-256 to borrow, so a hash +// is unavoidable for the manifest's sake; doing it through the shared verifier +// also guarantees the kept entry is byte-for-byte identical to a freshly +// downloaded one and re-checks integrity for free. Reading a large file from +// local disk is still orders of magnitude cheaper than re-downloading it. A +// file that is missing, the wrong size, or fails verification is not reused; the +// caller re-downloads it. +func reuseMaterializedFile(fileName string, source hfapi.SnapshotFile) (ManifestFile, bool) { + info, err := os.Stat(fileName) + if err != nil || !info.Mode().IsRegular() || info.Size() != source.Size { + return ManifestFile{}, false + } + entry, err := verifyDownloadedFile(fileName, source) + if err != nil { + return ManifestFile{}, false + } + return entry, true +} + func verifyDownloadedFile(fileName string, source hfapi.SnapshotFile) (ManifestFile, error) { file, err := os.Open(fileName) if err != nil { diff --git a/pkg/modelartifacts/materializer_test.go b/pkg/modelartifacts/materializer_test.go index ad11d0ff9..b891960ae 100644 --- a/pkg/modelartifacts/materializer_test.go +++ b/pkg/modelartifacts/materializer_test.go @@ -8,6 +8,8 @@ import ( "net/http/httptest" "os" "path/filepath" + "strconv" + "strings" "sync" "sync/atomic" @@ -151,6 +153,100 @@ var _ = Describe("controller artifact materializer", func() { 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)