Skip to content

Commit 883a4e7

Browse files
committed
fix(modelartifacts): stage each writer's artifact in its own partial tree
Every writer used to stage into the same `.artifacts/.partial/<cacheKey>`. 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 <mudler@localai.io> Assisted-by: Claude:claude-opus-4-8 [Claude Code]
1 parent 2f33ad6 commit 883a4e7

7 files changed

Lines changed: 678 additions & 16 deletions

File tree

core/application/startup.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
coreStartup "github.com/mudler/LocalAI/core/startup"
2727
"github.com/mudler/LocalAI/internal"
2828
"github.com/mudler/LocalAI/pkg/downloader"
29+
"github.com/mudler/LocalAI/pkg/modelartifacts"
2930
"github.com/mudler/LocalAI/pkg/signals"
3031
"github.com/mudler/LocalAI/pkg/vram"
3132

@@ -100,6 +101,15 @@ func New(opts ...config.AppOption) (*Application, error) {
100101
} else if removed > 0 {
101102
xlog.Info("Reaped stale partial downloads", "count", removed)
102103
}
104+
// Managed artifacts stage into a per-writer tree, which a crashed writer's
105+
// successor no longer overwrites for it, so the tree itself needs reaping
106+
// too. Sweeping here as well as on the materialization path is what
107+
// reclaims a volume whose abandoned artifact is never requested again.
108+
if removed, cErr := modelartifacts.SweepStalePartialTrees(options.SystemState.Model.ModelsPath, modelartifacts.PartialOrphanTTL, ""); cErr != nil {
109+
xlog.Warn("Failed to reap abandoned artifact partials", "error", cErr)
110+
} else if removed > 0 {
111+
xlog.Info("Reaped abandoned artifact partials", "count", removed)
112+
}
103113
if options.GeneratedContentDir != "" {
104114
err := os.MkdirAll(options.GeneratedContentDir, 0o750)
105115
if err != nil {

pkg/modelartifacts/materializer.go

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,11 @@ type Manager struct {
6060
huggingFaceToken string
6161
newLocker func(string) Locker
6262
lockWait time.Duration
63+
// writerID names this manager's staging trees. It is drawn once, at
64+
// construction, and deliberately never persisted: a partial tree belongs to
65+
// the process run that created it, and outliving that run is precisely what
66+
// it must not do.
67+
writerID string
6368
}
6469

6570
type ManagerOption func(*Manager)
@@ -100,6 +105,7 @@ func NewManager(resolver SnapshotResolver, options ...ManagerOption) *Manager {
100105
resolver: resolver,
101106
newLocker: func(path string) Locker { return flock.New(path) },
102107
lockWait: DefaultLockWait,
108+
writerID: newWriterID(),
103109
}
104110
for _, option := range options {
105111
option(manager)
@@ -310,6 +316,22 @@ func removeInvalidFinal(layout Layout) error {
310316
}
311317

312318
func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec Spec, snapshot hfapi.Snapshot, token string, layout Layout) (Result, error) {
319+
layout, err := layout.WithWriter(m.writerID)
320+
if err != nil {
321+
return Result{}, err
322+
}
323+
if err := os.MkdirAll(layout.PartialRoot, 0o750); err != nil {
324+
return Result{}, err
325+
}
326+
// Reclaim before staging, while the lock is held, so the two operations
327+
// that touch foreign trees only ever run when a peer that respects the lock
328+
// cannot be writing.
329+
if removed, err := SweepStalePartialTrees(modelsPath, PartialOrphanTTL, layout.Partial); err != nil {
330+
xlog.Warn("failed to sweep abandoned artifact partials", "error", err)
331+
} else if removed > 0 {
332+
xlog.Info("reclaimed abandoned artifact partials", "count", removed)
333+
}
334+
adoptOrphanPartial(layout)
313335
if err := os.MkdirAll(layout.Partial, 0o750); err != nil {
314336
return Result{}, err
315337
}
@@ -434,8 +456,29 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec
434456
return Result{}, err
435457
}
436458
ReportProgress(ctx, ProgressEvent{Phase: PhaseCommitting, Artifact: spec.Name, CurrentBytes: totalBytes, TotalBytes: totalBytes, CompletedFiles: len(snapshot.Files), TotalFiles: len(snapshot.Files)})
459+
return m.commit(modelsPath, spec, layout, manifest)
460+
}
461+
462+
// commit publishes this writer's staging tree under the artifact's final path.
463+
//
464+
// The rename is atomic and refuses to land on a populated destination, so a
465+
// peer either published before us or has not published at all. Losing that race
466+
// is not an error worth surfacing: the artifact is content-addressed, so the
467+
// peer's tree holds the same bytes we just downloaded and verified. Adopting it
468+
// and dropping ours is what keeps a broken lock cheap - without this, two
469+
// writers racing to commit would hand one caller a bare ENOTEMPTY for work that
470+
// actually succeeded.
471+
func (m *Manager) commit(modelsPath string, spec Spec, layout Layout, manifest Manifest) (Result, error) {
437472
if err := os.Rename(layout.Partial, layout.Final); err != nil {
438-
return Result{}, err
473+
cached, ok := committedResult(modelsPath, spec)
474+
if !ok {
475+
return Result{}, err
476+
}
477+
xlog.Info("another writer published this artifact first; discarding the duplicate", "partial", layout.Partial)
478+
if rmErr := removePartialTree(layout.PartialRoot, layout.Partial); rmErr != nil {
479+
xlog.Warn("failed to discard a duplicate artifact partial", "partial", layout.Partial, "error", rmErr)
480+
}
481+
return cached, nil
439482
}
440483
relative, err := RelativeSnapshotPath(spec.Resolved.CacheKey)
441484
if err != nil {
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
package modelartifacts_test
2+
3+
import (
4+
"context"
5+
"crypto/sha256"
6+
"encoding/hex"
7+
"net/http"
8+
"net/http/httptest"
9+
"os"
10+
"path/filepath"
11+
"sync"
12+
"sync/atomic"
13+
"time"
14+
15+
. "github.com/onsi/ginkgo/v2"
16+
. "github.com/onsi/gomega"
17+
18+
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
19+
"github.com/mudler/LocalAI/pkg/modelartifacts"
20+
)
21+
22+
// bypassedLocker grants the artifact lock to everyone who asks. It reproduces
23+
// the state a real cluster was in during #10981: flock(2) on CIFS returned
24+
// EACCES, both controller replicas concluded the lock was unusable, and both
25+
// proceeded to materialize the same artifact at the same time. The lock is
26+
// meant to prevent that, but a correctness property must not depend on a
27+
// primitive that a network filesystem can silently take away.
28+
type bypassedLocker struct{}
29+
30+
func (bypassedLocker) TryLock() (bool, error) { return true, nil }
31+
func (bypassedLocker) Unlock() error { return nil }
32+
33+
// rendezvousServer serves payload to every GET, but holds each response until
34+
// concurrent GETs have arrived, then dribbles the body out in two chunks. That
35+
// makes the two writers provably overlap inside the download rather than
36+
// relying on scheduling luck, so the spec is deterministic rather than flaky.
37+
func rendezvousServer(payload []byte, concurrent int32) *httptest.Server {
38+
var arrived atomic.Int32
39+
gate := make(chan struct{})
40+
var once sync.Once
41+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
42+
if r.Method != http.MethodGet {
43+
// The resume probe is a HEAD. Answering without Accept-Ranges keeps
44+
// this server non-resumable, which is the harsher case: a writer
45+
// that finds a foreign partial discards it outright.
46+
w.WriteHeader(http.StatusOK)
47+
return
48+
}
49+
if arrived.Add(1) >= concurrent {
50+
once.Do(func() { close(gate) })
51+
}
52+
select {
53+
case <-gate:
54+
case <-time.After(10 * time.Second):
55+
}
56+
half := len(payload) / 2
57+
_, _ = w.Write(payload[:half])
58+
w.(http.Flusher).Flush()
59+
time.Sleep(50 * time.Millisecond)
60+
_, _ = w.Write(payload[half:])
61+
}))
62+
DeferCleanup(server.Close)
63+
return server
64+
}
65+
66+
var _ = Describe("concurrent materialization without a working lock", func() {
67+
// The partial tree is shared mutable state. Before writer-unique partial
68+
// paths it was safe only because the lock held: two writers opened the same
69+
// blob with O_APPEND and interleaved into one file, and the resume probe
70+
// read the other writer's in-flight size. SHA verification caught the
71+
// damage only after both had burned the whole download.
72+
It("lets two writers materialize the same artifact concurrently without corrupting each other", func() {
73+
payload := make([]byte, 128*1024)
74+
for i := range payload {
75+
payload[i] = byte(i % 251)
76+
}
77+
sum := sha256.Sum256(payload)
78+
server := rendezvousServer(payload, 2)
79+
80+
snapshot := hfapi.Snapshot{
81+
Endpoint: "https://huggingface.co", Repo: "owner/repo",
82+
RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
83+
Files: []hfapi.SnapshotFile{{
84+
Path: "model.safetensors", Size: int64(len(payload)),
85+
LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model",
86+
}},
87+
}
88+
modelsPath := GinkgoT().TempDir()
89+
spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}
90+
91+
type outcome struct {
92+
result modelartifacts.Result
93+
err error
94+
}
95+
outcomes := make([]outcome, 2)
96+
var group sync.WaitGroup
97+
for writer := range outcomes {
98+
group.Add(1)
99+
go func() {
100+
defer GinkgoRecover()
101+
defer group.Done()
102+
// Separate managers stand in for separate processes: each one
103+
// draws its own writer identity at construction.
104+
manager := modelartifacts.NewManager(&fakeSnapshotResolver{snapshot: snapshot},
105+
modelartifacts.WithLocker(func(string) modelartifacts.Locker { return bypassedLocker{} }))
106+
result, err := manager.Ensure(context.Background(), modelsPath, spec)
107+
outcomes[writer] = outcome{result: result, err: err}
108+
}()
109+
}
110+
group.Wait()
111+
112+
for writer, got := range outcomes {
113+
Expect(got.err).NotTo(HaveOccurred(),
114+
"writer %d must not be poisoned by its peer; two writers sharing a partial path is exactly the corruption this guards against", writer)
115+
Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(got.result.RelativePath), "model.safetensors"))).
116+
To(Equal(payload), "writer %d resolved to a snapshot whose bytes are not the artifact", writer)
117+
}
118+
119+
// Exactly one snapshot is committed: the loser of the commit race
120+
// reconciles onto the winner's tree rather than publishing a rival.
121+
committed, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", "huggingface", "*"))
122+
Expect(err).NotTo(HaveOccurred())
123+
Expect(committed).To(HaveLen(1))
124+
125+
// Both writers finished, so neither may leave a partial tree behind.
126+
leftovers, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", ".partial", "*"))
127+
Expect(err).NotTo(HaveOccurred())
128+
Expect(leftovers).To(BeEmpty(), "a completed writer must not leak its partial tree onto a shared volume")
129+
})
130+
})

0 commit comments

Comments
 (0)