From 883a4e787e5432468ef1c59334f0760d7986aea7 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 20 Jul 2026 20:31:11 +0000 Subject: [PATCH] fix(modelartifacts): stage each writer's artifact in its own partial tree Every writer used to stage into the same `.artifacts/.partial/`. That was safe only because the artifact lock held: two writers that both believed they had it opened the same blob with O_APPEND and interleaved their bytes into one file, while the resume probe read the other writer's in-flight size. SHA verification caught the damage only after both had burned the entire download. #10986 restored the lock's precondition on CIFS but left the dependency in place. Suffix the staging tree with a writer identity drawn once per process run, so concurrent writers cannot corrupt each other whatever the lock does. The lock stops being a correctness dependency and becomes a pure efficiency optimisation: a lock failure now costs a duplicated download, not a corrupted one. Commit stays an atomic rename. The loser of a commit race reconciles onto the winner's tree instead of surfacing a bare ENOTEMPTY for work that actually succeeded, since the artifact is content-addressed and both trees hold the same verified bytes. Writer-unique staging means a crashed writer's tree is no longer overwritten by its successor, so two things are added to keep it from becoming a disk leak and a resume regression: - A sweep reclaims trees whose contents have been untouched for 24h, matching the window the startup reaper already uses for stray *.partial files. It reads the newest mtime anywhere inside the tree, because writing a blob never touches an ancestor, and refuses any name this package did not write. A live download writes continuously, and the downloader's stall watchdog aborts a silent one long before it could look abandoned. - Adoption lets a restarted process claim a dead predecessor's tree for the same artifact and resume from its bytes, which a tens-of-gigabytes repo depends on. The claim is an atomic rename, so racing adopters cannot both win. It runs only under the artifact lock - which is released exactly when the owning process dies - and only on a tree idle for 5 minutes as a second line of defence for when the lock does not exclude. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude:claude-opus-4-8 [Claude Code] --- core/application/startup.go | 10 + pkg/modelartifacts/materializer.go | 45 +++- .../materializer_concurrent_test.go | 130 ++++++++++ pkg/modelartifacts/partial.go | 207 ++++++++++++++++ pkg/modelartifacts/partial_test.go | 232 ++++++++++++++++++ pkg/modelartifacts/path.go | 69 ++++-- pkg/modelartifacts/types.go | 1 + 7 files changed, 678 insertions(+), 16 deletions(-) create mode 100644 pkg/modelartifacts/materializer_concurrent_test.go create mode 100644 pkg/modelartifacts/partial.go create mode 100644 pkg/modelartifacts/partial_test.go diff --git a/core/application/startup.go b/core/application/startup.go index 455d78c13744..7dd5a8d61ee3 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -26,6 +26,7 @@ import ( coreStartup "github.com/mudler/LocalAI/core/startup" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/downloader" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/signals" "github.com/mudler/LocalAI/pkg/vram" @@ -100,6 +101,15 @@ func New(opts ...config.AppOption) (*Application, error) { } else if removed > 0 { xlog.Info("Reaped stale partial downloads", "count", removed) } + // Managed artifacts stage into a per-writer tree, which a crashed writer's + // successor no longer overwrites for it, so the tree itself needs reaping + // too. Sweeping here as well as on the materialization path is what + // reclaims a volume whose abandoned artifact is never requested again. + if removed, cErr := modelartifacts.SweepStalePartialTrees(options.SystemState.Model.ModelsPath, modelartifacts.PartialOrphanTTL, ""); cErr != nil { + xlog.Warn("Failed to reap abandoned artifact partials", "error", cErr) + } else if removed > 0 { + xlog.Info("Reaped abandoned artifact partials", "count", removed) + } if options.GeneratedContentDir != "" { err := os.MkdirAll(options.GeneratedContentDir, 0o750) if err != nil { diff --git a/pkg/modelartifacts/materializer.go b/pkg/modelartifacts/materializer.go index 149b96d2ef41..16b10408cc93 100644 --- a/pkg/modelartifacts/materializer.go +++ b/pkg/modelartifacts/materializer.go @@ -60,6 +60,11 @@ type Manager struct { huggingFaceToken string newLocker func(string) Locker lockWait time.Duration + // writerID names this manager's staging trees. It is drawn once, at + // construction, and deliberately never persisted: a partial tree belongs to + // the process run that created it, and outliving that run is precisely what + // it must not do. + writerID string } type ManagerOption func(*Manager) @@ -100,6 +105,7 @@ func NewManager(resolver SnapshotResolver, options ...ManagerOption) *Manager { resolver: resolver, newLocker: func(path string) Locker { return flock.New(path) }, lockWait: DefaultLockWait, + writerID: newWriterID(), } for _, option := range options { option(manager) @@ -310,6 +316,22 @@ func removeInvalidFinal(layout Layout) error { } func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec Spec, snapshot hfapi.Snapshot, token string, layout Layout) (Result, error) { + layout, err := layout.WithWriter(m.writerID) + if err != nil { + return Result{}, err + } + if err := os.MkdirAll(layout.PartialRoot, 0o750); err != nil { + return Result{}, err + } + // Reclaim before staging, while the lock is held, so the two operations + // that touch foreign trees only ever run when a peer that respects the lock + // cannot be writing. + if removed, err := SweepStalePartialTrees(modelsPath, PartialOrphanTTL, layout.Partial); err != nil { + xlog.Warn("failed to sweep abandoned artifact partials", "error", err) + } else if removed > 0 { + xlog.Info("reclaimed abandoned artifact partials", "count", removed) + } + adoptOrphanPartial(layout) if err := os.MkdirAll(layout.Partial, 0o750); err != nil { return Result{}, err } @@ -434,8 +456,29 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec return Result{}, err } ReportProgress(ctx, ProgressEvent{Phase: PhaseCommitting, Artifact: spec.Name, CurrentBytes: totalBytes, TotalBytes: totalBytes, CompletedFiles: len(snapshot.Files), TotalFiles: len(snapshot.Files)}) + return m.commit(modelsPath, spec, layout, manifest) +} + +// commit publishes this writer's staging tree under the artifact's final path. +// +// The rename is atomic and refuses to land on a populated destination, so a +// peer either published before us or has not published at all. Losing that race +// is not an error worth surfacing: the artifact is content-addressed, so the +// peer's tree holds the same bytes we just downloaded and verified. Adopting it +// and dropping ours is what keeps a broken lock cheap - without this, two +// writers racing to commit would hand one caller a bare ENOTEMPTY for work that +// actually succeeded. +func (m *Manager) commit(modelsPath string, spec Spec, layout Layout, manifest Manifest) (Result, error) { if err := os.Rename(layout.Partial, layout.Final); err != nil { - return Result{}, err + cached, ok := committedResult(modelsPath, spec) + if !ok { + return Result{}, err + } + xlog.Info("another writer published this artifact first; discarding the duplicate", "partial", layout.Partial) + if rmErr := removePartialTree(layout.PartialRoot, layout.Partial); rmErr != nil { + xlog.Warn("failed to discard a duplicate artifact partial", "partial", layout.Partial, "error", rmErr) + } + return cached, nil } relative, err := RelativeSnapshotPath(spec.Resolved.CacheKey) if err != nil { diff --git a/pkg/modelartifacts/materializer_concurrent_test.go b/pkg/modelartifacts/materializer_concurrent_test.go new file mode 100644 index 000000000000..962953b1a3d9 --- /dev/null +++ b/pkg/modelartifacts/materializer_concurrent_test.go @@ -0,0 +1,130 @@ +package modelartifacts_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "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" +) + +// bypassedLocker grants the artifact lock to everyone who asks. It reproduces +// the state a real cluster was in during #10981: flock(2) on CIFS returned +// EACCES, both controller replicas concluded the lock was unusable, and both +// proceeded to materialize the same artifact at the same time. The lock is +// meant to prevent that, but a correctness property must not depend on a +// primitive that a network filesystem can silently take away. +type bypassedLocker struct{} + +func (bypassedLocker) TryLock() (bool, error) { return true, nil } +func (bypassedLocker) Unlock() error { return nil } + +// rendezvousServer serves payload to every GET, but holds each response until +// concurrent GETs have arrived, then dribbles the body out in two chunks. That +// makes the two writers provably overlap inside the download rather than +// relying on scheduling luck, so the spec is deterministic rather than flaky. +func rendezvousServer(payload []byte, concurrent int32) *httptest.Server { + var arrived atomic.Int32 + gate := make(chan struct{}) + var once sync.Once + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + // The resume probe is a HEAD. Answering without Accept-Ranges keeps + // this server non-resumable, which is the harsher case: a writer + // that finds a foreign partial discards it outright. + w.WriteHeader(http.StatusOK) + return + } + if arrived.Add(1) >= concurrent { + once.Do(func() { close(gate) }) + } + select { + case <-gate: + case <-time.After(10 * time.Second): + } + half := len(payload) / 2 + _, _ = w.Write(payload[:half]) + w.(http.Flusher).Flush() + time.Sleep(50 * time.Millisecond) + _, _ = w.Write(payload[half:]) + })) + DeferCleanup(server.Close) + return server +} + +var _ = Describe("concurrent materialization without a working lock", func() { + // The partial tree is shared mutable state. Before writer-unique partial + // paths it was safe only because the lock held: two writers opened the same + // blob with O_APPEND and interleaved into one file, and the resume probe + // read the other writer's in-flight size. SHA verification caught the + // damage only after both had burned the whole download. + It("lets two writers materialize the same artifact concurrently without corrupting each other", func() { + payload := make([]byte, 128*1024) + for i := range payload { + payload[i] = byte(i % 251) + } + sum := sha256.Sum256(payload) + server := rendezvousServer(payload, 2) + + snapshot := hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{ + Path: "model.safetensors", Size: int64(len(payload)), + LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model", + }}, + } + modelsPath := GinkgoT().TempDir() + spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}} + + type outcome struct { + result modelartifacts.Result + err error + } + outcomes := make([]outcome, 2) + var group sync.WaitGroup + for writer := range outcomes { + group.Add(1) + go func() { + defer GinkgoRecover() + defer group.Done() + // Separate managers stand in for separate processes: each one + // draws its own writer identity at construction. + manager := modelartifacts.NewManager(&fakeSnapshotResolver{snapshot: snapshot}, + modelartifacts.WithLocker(func(string) modelartifacts.Locker { return bypassedLocker{} })) + result, err := manager.Ensure(context.Background(), modelsPath, spec) + outcomes[writer] = outcome{result: result, err: err} + }() + } + group.Wait() + + for writer, got := range outcomes { + Expect(got.err).NotTo(HaveOccurred(), + "writer %d must not be poisoned by its peer; two writers sharing a partial path is exactly the corruption this guards against", writer) + Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(got.result.RelativePath), "model.safetensors"))). + To(Equal(payload), "writer %d resolved to a snapshot whose bytes are not the artifact", writer) + } + + // Exactly one snapshot is committed: the loser of the commit race + // reconciles onto the winner's tree rather than publishing a rival. + committed, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", "huggingface", "*")) + Expect(err).NotTo(HaveOccurred()) + Expect(committed).To(HaveLen(1)) + + // Both writers finished, so neither may leave a partial tree behind. + leftovers, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", ".partial", "*")) + Expect(err).NotTo(HaveOccurred()) + Expect(leftovers).To(BeEmpty(), "a completed writer must not leak its partial tree onto a shared volume") + }) +}) diff --git a/pkg/modelartifacts/partial.go b/pkg/modelartifacts/partial.go new file mode 100644 index 000000000000..c1bff813c4ca --- /dev/null +++ b/pkg/modelartifacts/partial.go @@ -0,0 +1,207 @@ +package modelartifacts + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "io/fs" + "os" + "path/filepath" + "time" + + "github.com/mudler/xlog" +) + +const ( + // PartialOrphanTTL is how long an unclaimed partial tree survives before a + // sweep reclaims it. Writer-unique staging means a crashed writer's tree is + // never overwritten by its successor, so without reaping, every crash + // during a large download leaks its bytes onto a shared volume forever. + // + // It matches the window the startup reaper already uses for stray + // *.partial files, and it is orders of magnitude beyond any write-free + // interval a live writer can have: the download stall guard aborts a silent + // transfer after DownloadStallTimeout, and the only other gap between + // writes is hashing one downloaded file. + PartialOrphanTTL = 24 * time.Hour + + // partialAdoptionIdleWindow is how long a foreign partial tree for the same + // artifact must have been untouched before this writer takes it over and + // resumes from it. Without adoption, writer-unique staging would silently + // cost the resume that a restarted process relies on to finish a + // multi-tens-of-gigabytes repo. + // + // Adoption is normally justified by the artifact lock alone: it only ever + // runs while this process holds that lock, and flock is released exactly + // when the owning process dies, so a peer that respects the lock is + // provably not writing. The idle window is the second line of defence for + // the case this whole change is about, where the lock fails to exclude. + // + // The window is a deliberate compromise in both directions. Too long and a + // restarted replica re-downloads from zero; too short and a broken lock + // lets us claim a live peer's tree. Both cost the same thing - one wasted + // download - because every adopted blob is still SHA-verified before it is + // promoted, so a wrongly adopted tree fails verification instead of + // committing corruption. + partialAdoptionIdleWindow = 5 * time.Minute +) + +// newWriterID draws the identity that makes this process run's partial tree its +// own. +// +// It has to be unique across every process that can reach the same models +// volume. That rules out the PID, which repeats freely across containers +// sharing a NAS mount, and the hostname, which a restarted pod can inherit +// while the dead pod's tree is still on disk. 64 bits of crypto/rand depends on +// nothing outside this process; a collision would need two replicas to draw the +// same value in the same 24h window, and even then it degrades to the +// shared-tree behaviour that was the status quo, not to something worse. +func newWriterID() string { + var raw [8]byte + // crypto/rand.Read never returns an error; it panics if the system RNG is + // unusable, which is not a condition worth threading through this API. + _, _ = rand.Read(raw[:]) + return hex.EncodeToString(raw[:]) +} + +// newestModTime reports the most recent modification anywhere in tree. +// +// The directory's own mtime is not enough: writing bytes into +// `.downloads/.partial` does not touch any ancestor, so a busy writer's +// top-level directory can look hours old. Reading the whole subtree is what +// makes "is anyone still working here?" answerable. Unreadable entries are +// skipped rather than aborting the walk, but they count as "now" so a tree we +// cannot fully inspect is never judged stale. +func newestModTime(tree string) (time.Time, error) { + newest := time.Time{} + err := filepath.WalkDir(tree, func(_ string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + newest = time.Now() + return nil + } + info, err := entry.Info() + if err != nil { + newest = time.Now() + return nil + } + if info.ModTime().After(newest) { + newest = info.ModTime() + } + return nil + }) + return newest, err +} + +// idleFor reports whether tree has seen no modification for at least window. +func idleFor(tree string, window time.Duration) bool { + newest, err := newestModTime(tree) + if err != nil { + return false + } + return time.Since(newest) >= window +} + +// removePartialTree deletes one staging tree, refusing any path that is not a +// direct, correctly named child of the partial root. Writer-unique staging +// turned a self-healing overwrite into an explicit delete, and an explicit +// delete on a shared models volume is worth being paranoid about. +func removePartialTree(partialRoot, tree string) error { + name := filepath.Base(tree) + if filepath.Dir(tree) != partialRoot { + return fmt.Errorf("refusing to remove partial tree outside %q", partialRoot) + } + if _, _, ok := splitPartialDirName(name); !ok { + return fmt.Errorf("refusing to remove unrecognised partial tree %q", name) + } + root, err := os.OpenRoot(partialRoot) + if err != nil { + return err + } + defer func() { _ = root.Close() }() + return root.RemoveAll(name) +} + +// adoptOrphanPartial hands this writer a dead peer's staging tree for the same +// artifact, so a download interrupted by a crash or a restart resumes instead +// of starting over. +// +// The claim is a rename, which is atomic: two writers racing for the same +// orphan cannot both win, and the loser simply starts fresh. A failure to adopt +// is never fatal - it costs bytes, not correctness - so every error here is +// logged and swallowed. +func adoptOrphanPartial(layout Layout) { + entries, err := os.ReadDir(layout.PartialRoot) + if err != nil { + return + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + name := entry.Name() + candidate := filepath.Join(layout.PartialRoot, name) + if candidate == layout.Partial { + continue + } + cacheKey, _, ok := splitPartialDirName(name) + if !ok || cacheKey != layout.CacheKey { + continue + } + if !idleFor(candidate, partialAdoptionIdleWindow) { + continue + } + if err := os.Rename(candidate, layout.Partial); err != nil { + xlog.Debug("could not adopt an abandoned artifact partial", "partial", candidate, "error", err) + continue + } + xlog.Info("resuming an abandoned artifact download", "partial", candidate, "adopted-as", layout.Partial) + return + } +} + +// SweepStalePartialTrees reclaims staging trees left behind by writers that +// never finished, returning how many it removed. +// +// A tree is removed only when its name is one this package writes, it is not +// the caller's own, and nothing anywhere inside it has been touched for +// olderThan. That last condition is what makes the sweep safe against a live +// writer: a running download writes continuously, so its newest mtime is +// seconds old, and a stalled one is aborted by the downloader's own watchdog +// long before it could look abandoned. +// +// A missing tree is not an error, and one unremovable entry does not abort the +// sweep. +func SweepStalePartialTrees(modelsPath string, olderThan time.Duration, keep string) (int, error) { + partialRoot := filepath.Join(modelsPath, ".artifacts", ".partial") + entries, err := os.ReadDir(partialRoot) + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, err + } + removed := 0 + for _, entry := range entries { + if !entry.IsDir() { + continue + } + name := entry.Name() + tree := filepath.Join(partialRoot, name) + if tree == keep { + continue + } + if _, _, ok := splitPartialDirName(name); !ok { + continue + } + if !idleFor(tree, olderThan) { + continue + } + if err := removePartialTree(partialRoot, tree); err != nil { + xlog.Warn("failed to remove abandoned artifact partial", "partial", tree, "error", err) + continue + } + removed++ + xlog.Info("removed abandoned artifact partial", "partial", tree) + } + return removed, nil +} diff --git a/pkg/modelartifacts/partial_test.go b/pkg/modelartifacts/partial_test.go new file mode 100644 index 000000000000..f61c764f02db --- /dev/null +++ b/pkg/modelartifacts/partial_test.go @@ -0,0 +1,232 @@ +package modelartifacts_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "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" +) + +const ( + testCacheKey = "b79e23e0b9c50af094d627582df30109eff8637438864172d64be07dfc5a98f9" + testWriterID = "00112233445566aa" + otherWriterID = "ffeeddccbbaa9988" +) + +// backdate rewinds every mtime in tree so it looks abandoned. The sweep and the +// adoption scan both read the newest mtime anywhere inside the tree, not the +// directory's own, so every entry has to move. +func backdate(tree string, age time.Duration) { + when := time.Now().Add(-age) + Expect(filepath.Walk(tree, func(path string, _ os.FileInfo, err error) error { + if err != nil { + return err + } + return os.Chtimes(path, when, when) + })).To(Succeed()) +} + +func partialRoot(modelsPath string) string { + return filepath.Join(modelsPath, ".artifacts", ".partial") +} + +// stageOrphan writes a partial tree as if a writer had died holding it, with +// blob carrying however many bytes of file it had managed to fetch. +func stageOrphan(modelsPath, writerID, hubPath string, blob []byte) string { + tree := filepath.Join(partialRoot(modelsPath), testCacheKey+"."+writerID) + nameSum := sha256.Sum256([]byte(hubPath)) + downloads := filepath.Join(tree, ".downloads") + Expect(os.MkdirAll(downloads, 0o750)).To(Succeed()) + Expect(os.MkdirAll(filepath.Join(tree, "snapshot"), 0o750)).To(Succeed()) + Expect(os.WriteFile(filepath.Join(downloads, hex.EncodeToString(nameSum[:])+".partial"), blob, 0o600)).To(Succeed()) + return tree +} + +var _ = Describe("writer-unique artifact partial trees", func() { + It("gives each writer its own staging path and rejects a malformed identity", func() { + spec := modelartifacts.Spec{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + Resolved: &modelartifacts.Resolved{Endpoint: "https://huggingface.co", Revision: "0123456789abcdef0123456789abcdef01234567", CacheKey: testCacheKey}, + } + layout, err := modelartifacts.LayoutFor("/models", spec) + Expect(err).NotTo(HaveOccurred()) + Expect(layout.Partial).To(BeEmpty(), "a layout without a writer identity must not name a staging tree") + + mine, err := layout.WithWriter(testWriterID) + Expect(err).NotTo(HaveOccurred()) + theirs, err := layout.WithWriter(otherWriterID) + Expect(err).NotTo(HaveOccurred()) + Expect(mine.Partial).To(Equal(filepath.Join("/models", ".artifacts", ".partial", testCacheKey+"."+testWriterID))) + Expect(mine.Partial).NotTo(Equal(theirs.Partial)) + + // A writer identity reaches the filesystem as a path component, so + // anything that is not the exact drawn shape is refused rather than + // sanitised. + for _, bad := range []string{"", "../escape", "not-hex", strings.Repeat("a", 17)} { + _, err := layout.WithWriter(bad) + Expect(err).To(MatchError(ContainSubstring("invalid artifact writer id")), "accepted writer id %q", bad) + } + }) + + Describe("reclaiming abandoned trees", func() { + var modelsPath string + + BeforeEach(func() { + modelsPath = GinkgoT().TempDir() + Expect(os.MkdirAll(partialRoot(modelsPath), 0o750)).To(Succeed()) + }) + + It("removes a tree nothing has touched for the whole window", func() { + tree := stageOrphan(modelsPath, otherWriterID, "model.safetensors", []byte("half")) + backdate(tree, 48*time.Hour) + + removed, err := modelartifacts.SweepStalePartialTrees(modelsPath, modelartifacts.PartialOrphanTTL, "") + Expect(err).NotTo(HaveOccurred()) + Expect(removed).To(Equal(1)) + Expect(tree).NotTo(BeADirectory()) + }) + + // This is the property that keeps the fix from becoming a data-loss + // bug: a running download writes continuously, so the freshness of the + // bytes inside the tree - not the directory's own mtime, which nothing + // updates - is what proves the writer is still alive. + It("never removes a tree whose contents are still being written", func() { + tree := stageOrphan(modelsPath, otherWriterID, "model.safetensors", []byte("half")) + // An old directory with a fresh blob inside is exactly what a + // long-running download looks like. + backdate(tree, 48*time.Hour) + now := time.Now() + blob := filepath.Join(tree, ".downloads", + hex.EncodeToString(func() []byte { sum := sha256.Sum256([]byte("model.safetensors")); return sum[:] }())) + Expect(os.Chtimes(blob+".partial", now, now)).To(Succeed()) + + removed, err := modelartifacts.SweepStalePartialTrees(modelsPath, modelartifacts.PartialOrphanTTL, "") + Expect(err).NotTo(HaveOccurred()) + Expect(removed).To(BeZero()) + Expect(tree).To(BeADirectory()) + }) + + It("never removes the caller's own tree, however old it looks", func() { + tree := stageOrphan(modelsPath, testWriterID, "model.safetensors", []byte("half")) + backdate(tree, 48*time.Hour) + + removed, err := modelartifacts.SweepStalePartialTrees(modelsPath, modelartifacts.PartialOrphanTTL, tree) + Expect(err).NotTo(HaveOccurred()) + Expect(removed).To(BeZero()) + Expect(tree).To(BeADirectory()) + }) + + It("never removes a directory it did not create", func() { + stranger := filepath.Join(partialRoot(modelsPath), "not-an-artifact") + Expect(os.MkdirAll(stranger, 0o750)).To(Succeed()) + backdate(stranger, 48*time.Hour) + + removed, err := modelartifacts.SweepStalePartialTrees(modelsPath, modelartifacts.PartialOrphanTTL, "") + Expect(err).NotTo(HaveOccurred()) + Expect(removed).To(BeZero()) + Expect(stranger).To(BeADirectory()) + }) + + It("treats a missing partial root as nothing to do", func() { + removed, err := modelartifacts.SweepStalePartialTrees(GinkgoT().TempDir(), modelartifacts.PartialOrphanTTL, "") + Expect(err).NotTo(HaveOccurred()) + Expect(removed).To(BeZero()) + }) + }) + + // Writer-unique staging would otherwise cost the resume a restarted replica + // depends on to ever finish a tens-of-gigabytes repo: its predecessor's + // bytes sit under a writer id it will never generate again. + It("adopts an abandoned tree for the same artifact and resumes from its bytes", func() { + payload := []byte(strings.Repeat("artifact-bytes-", 512)) + sum := sha256.Sum256(payload) + already := len(payload) / 2 + + var servedFrom atomic.Int64 + servedFrom.Store(-1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Accept-Ranges", "bytes") + if r.Method == http.MethodHead { + w.WriteHeader(http.StatusOK) + return + } + start := int64(0) + if _, err := fmt.Sscanf(r.Header.Get("Range"), "bytes=%d-", &start); err != nil { + start = 0 + } + servedFrom.Store(start) + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, len(payload)-1, len(payload))) + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write(payload[start:]) + })) + DeferCleanup(server.Close) + + modelsPath := GinkgoT().TempDir() + spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}} + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{ + Path: "model.safetensors", Size: int64(len(payload)), + LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model", + }}, + }} + key, err := modelartifacts.CacheKey(modelartifacts.Spec{ + Source: spec.Source, + Resolved: &modelartifacts.Resolved{Endpoint: "https://huggingface.co", Revision: "0123456789abcdef0123456789abcdef01234567"}, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal(testCacheKey), "the staged orphan must belong to the artifact under test") + + orphan := stageOrphan(modelsPath, otherWriterID, "model.safetensors", payload[:already]) + // Only a tree idle for the adoption window is claimable: the artifact + // lock already proves a peer that respects it is not writing, and this + // is the second line of defence for when the lock does not exclude. + backdate(orphan, time.Hour) + + result, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), modelsPath, spec) + Expect(err).NotTo(HaveOccurred()) + Expect(servedFrom.Load()).To(Equal(int64(already)), "the adopted partial's bytes were re-downloaded instead of resumed") + Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(result.RelativePath), "model.safetensors"))).To(Equal(payload)) + Expect(orphan).NotTo(BeADirectory(), "adoption must move the tree, not copy it") + }) + + It("leaves a freshly written tree for its owner instead of adopting it", func() { + payload := []byte("weight-bytes") + sum := sha256.Sum256(payload) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write(payload) + })) + DeferCleanup(server.Close) + + modelsPath := GinkgoT().TempDir() + // No backdating: this stands for a peer that is mid-download right now. + orphan := stageOrphan(modelsPath, otherWriterID, "model.safetensors", []byte("partial")) + + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", + Files: []hfapi.SnapshotFile{{ + Path: "model.safetensors", Size: int64(len(payload)), + LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model", + }}, + }} + _, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), modelsPath, + modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}) + Expect(err).NotTo(HaveOccurred()) + Expect(orphan).To(BeADirectory(), "a live peer's staging tree must survive another writer's materialization") + }) +}) diff --git a/pkg/modelartifacts/path.go b/pkg/modelartifacts/path.go index 5ac99a088040..052f124993d5 100644 --- a/pkg/modelartifacts/path.go +++ b/pkg/modelartifacts/path.go @@ -10,13 +10,53 @@ import ( ) type Layout struct { - Root string - Lock string - Partial string - PartialSnapshot string - Final string - Snapshot string - Manifest string + Root string + CacheKey string + Lock string + // PartialRoot holds every in-flight materialization for every artifact. + // Individual writers never share a subtree of it. + PartialRoot string + // Partial is this writer's private staging tree. It is empty until + // WithWriter stamps a writer identity onto the layout, because a partial + // tree belongs to one process run and the layout alone cannot name it. + Partial string + Final string + Snapshot string + Manifest string +} + +// WithWriter returns a copy of the layout whose partial tree belongs +// exclusively to writerID. +// +// Every writer used to stage into the same `.partial/`, which made +// the artifact lock a correctness dependency: two writers that both believed +// they held it opened the same blob with O_APPEND, interleaved their bytes into +// one file, and read each other's in-flight size when probing for a resume +// point (#10981 put a real cluster in exactly that state, because flock(2) on +// CIFS reports contention as EACCES). Suffixing the writer identity makes that +// impossible by construction, so a lock failure now costs a duplicated download +// rather than a corrupted one. +func (l Layout) WithWriter(writerID string) (Layout, error) { + if !writerIDPattern.MatchString(writerID) { + return Layout{}, fmt.Errorf("invalid artifact writer id") + } + if !cacheKeyPattern.MatchString(l.CacheKey) || l.PartialRoot == "" { + return Layout{}, fmt.Errorf("writer layout requires a resolved cache key") + } + l.Partial = filepath.Join(l.PartialRoot, l.CacheKey+"."+writerID) + return l, nil +} + +// splitPartialDirName recovers the (cacheKey, writerID) pair a partial tree +// name encodes, reporting false for anything this package did not create. The +// sweep and the adoption scan both refuse to touch a name they cannot parse, so +// an unrelated directory that happens to sit under `.partial` is never removed. +func splitPartialDirName(name string) (cacheKey string, writerID string, ok bool) { + key, writer, found := strings.Cut(name, ".") + if !found || !cacheKeyPattern.MatchString(key) || !writerIDPattern.MatchString(writer) { + return "", "", false + } + return key, writer, true } func CacheKey(spec Spec) (string, error) { @@ -59,16 +99,15 @@ func LayoutFor(modelsPath string, spec Spec) (Layout, error) { return Layout{}, fmt.Errorf("layout requires a resolved cache key") } root := filepath.Join(modelsPath, ".artifacts") - partial := filepath.Join(root, ".partial", spec.Resolved.CacheKey) final := filepath.Join(root, "huggingface", spec.Resolved.CacheKey) return Layout{ - Root: root, - Lock: filepath.Join(root, ".locks", spec.Resolved.CacheKey+".lock"), - Partial: partial, - PartialSnapshot: filepath.Join(partial, "snapshot"), - Final: final, - Snapshot: filepath.Join(final, "snapshot"), - Manifest: filepath.Join(final, "manifest.json"), + Root: root, + CacheKey: spec.Resolved.CacheKey, + Lock: filepath.Join(root, ".locks", spec.Resolved.CacheKey+".lock"), + PartialRoot: filepath.Join(root, ".partial"), + Final: final, + Snapshot: filepath.Join(final, "snapshot"), + Manifest: filepath.Join(final, "manifest.json"), }, nil } diff --git a/pkg/modelartifacts/types.go b/pkg/modelartifacts/types.go index 850ad414b952..e018d91fc2e6 100644 --- a/pkg/modelartifacts/types.go +++ b/pkg/modelartifacts/types.go @@ -23,6 +23,7 @@ const ( var ( commitPattern = regexp.MustCompile(`^[0-9a-f]{40}$`) cacheKeyPattern = regexp.MustCompile(`^[0-9a-f]{64}$`) + writerIDPattern = regexp.MustCompile(`^[0-9a-f]{16}$`) artifactNamePattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,63}$`) )