diff --git a/manifest/layer.go b/manifest/layer.go index 82d44953a..9698edb82 100644 --- a/manifest/layer.go +++ b/manifest/layer.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "os" + "time" ) type Layer struct { @@ -60,6 +61,9 @@ func NewLayer(r io.Reader, mediatype string) (Layer, error) { return Layer{}, err } } + if err := touchLayer(blob); err != nil { + return Layer{}, err + } return Layer{ MediaType: mediatype, @@ -83,6 +87,9 @@ func NewLayerFromLayer(digest, mediatype, from string) (Layer, error) { if err != nil { return Layer{}, err } + if err := touchLayer(blob); err != nil { + return Layer{}, err + } return Layer{ MediaType: mediatype, @@ -93,6 +100,11 @@ func NewLayerFromLayer(digest, mediatype, from string) (Layer, error) { }, nil } +func touchLayer(path string) error { + now := time.Now() + return os.Chtimes(path, now, now) +} + func (l *Layer) Open() (io.ReadSeekCloser, error) { if l.Digest == "" { return nil, errors.New("opening layer with empty digest") diff --git a/server/images.go b/server/images.go index 110a7dc6e..35094fa55 100644 --- a/server/images.go +++ b/server/images.go @@ -19,6 +19,7 @@ import ( "slices" "strconv" "strings" + "time" "github.com/ollama/ollama/api" "github.com/ollama/ollama/envconfig" @@ -33,6 +34,10 @@ import ( "github.com/ollama/ollama/x/imagegen/transfer" ) +// Blobs newer than this may belong to another process that has not written its +// manifest yet. They become eligible for the normal mark-and-sweep pass later. +const layerPruneGracePeriod = time.Hour + var ( errCapabilities = errors.New("does not support") errCapabilityCompletion = errors.New("completion") @@ -478,10 +483,23 @@ func PruneLayers() error { } for _, blob := range blobs { + if blob.IsDir() { + continue + } + + info, err := blob.Info() + if err != nil { + slog.Error("couldn't stat blob", "blob", blob.Name(), "error", err) + continue + } + if time.Since(info.ModTime()) < layerPruneGracePeriod { + continue + } + name := blob.Name() name = strings.ReplaceAll(name, "-", ":") - _, err := manifest.BlobsPath(name) + _, err = manifest.BlobsPath(name) if err != nil { if errors.Is(err, manifest.ErrInvalidDigestFormat) { // remove invalid blobs (e.g. partial downloads) diff --git a/server/images_test.go b/server/images_test.go index d7f99afc8..19a479838 100644 --- a/server/images_test.go +++ b/server/images_test.go @@ -5,14 +5,58 @@ import ( "fmt" "net/http" "net/http/httptest" + "os" "strings" "testing" + "time" "github.com/ollama/ollama/fs/ggml" + "github.com/ollama/ollama/manifest" "github.com/ollama/ollama/template" "github.com/ollama/ollama/types/model" ) +func TestPruneLayersSkipsRecentOrphans(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + + recentDigest := "sha256:0000000000000000000000000000000000000000000000000000000000000001" + oldDigest := "sha256:0000000000000000000000000000000000000000000000000000000000000002" + + for _, digest := range []string{recentDigest, oldDigest} { + p, err := manifest.BlobsPath(digest) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, nil, 0o644); err != nil { + t.Fatal(err) + } + } + + oldPath, err := manifest.BlobsPath(oldDigest) + if err != nil { + t.Fatal(err) + } + oldTime := time.Now().Add(-layerPruneGracePeriod - time.Hour) + if err := os.Chtimes(oldPath, oldTime, oldTime); err != nil { + t.Fatal(err) + } + + if err := PruneLayers(); err != nil { + t.Fatal(err) + } + + recentPath, err := manifest.BlobsPath(recentDigest) + if err != nil { + t.Fatal(err) + } + if _, err := os.Stat(recentPath); err != nil { + t.Fatalf("recent orphan was pruned: %v", err) + } + if _, err := os.Stat(oldPath); !os.IsNotExist(err) { + t.Fatalf("old orphan still exists: %v", err) + } +} + func TestModelCapabilities(t *testing.T) { // Create completion model (llama architecture without vision) completionModelPath, _ := createBinFile(t, ggml.KV{