From e7e91cd71c2b5f9440e995087c7d7524c2cfed1c Mon Sep 17 00:00:00 2001 From: Patrick Devine Date: Mon, 11 Sep 2023 11:46:35 -0700 Subject: [PATCH] add autoprune to remove unused layers (#491) --- cmd/cmd.go | 6 ++ server/images.go | 148 +++++++++++++++++++++++++++++++++++++++++------ server/routes.go | 1 + 3 files changed, 136 insertions(+), 19 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 80ce6822..41276f8f 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -672,6 +672,12 @@ func RunServer(cmd *cobra.Command, _ []string) error { origins = strings.Split(o, ",") } + if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" { + if err := server.PruneLayers(); err != nil { + return err + } + } + return server.Serve(ln, origins) } diff --git a/server/images.go b/server/images.go index 21b91843..f4181f43 100644 --- a/server/images.go +++ b/server/images.go @@ -269,6 +269,29 @@ func filenameWithPath(path, f string) (string, error) { } func CreateModel(ctx context.Context, name string, path string, fn func(resp api.ProgressResponse)) error { + mp := ParseModelPath(name) + + var manifest *ManifestV2 + var err error + var noprune string + + // build deleteMap to prune unused layers + deleteMap := make(map[string]bool) + + if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" { + manifest, _, err = GetManifest(mp) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + + if manifest != nil { + for _, l := range manifest.Layers { + deleteMap[l.Digest] = true + } + deleteMap[manifest.Config.Digest] = true + } + } + mf, err := os.Open(path) if err != nil { fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't open modelfile '%s'", path)}) @@ -506,6 +529,7 @@ func CreateModel(ctx context.Context, name string, path string, fn func(resp api var manifestLayers []*Layer for _, l := range layers { manifestLayers = append(manifestLayers, &l.Layer) + delete(deleteMap, l.Layer.Digest) } // Create a layer for the config object @@ -515,6 +539,7 @@ func CreateModel(ctx context.Context, name string, path string, fn func(resp api return err } layers = append(layers, cfg) + delete(deleteMap, cfg.Layer.Digest) if err := SaveLayers(layers, fn, false); err != nil { return err @@ -527,6 +552,14 @@ func CreateModel(ctx context.Context, name string, path string, fn func(resp api return err } + if noprune == "" { + fn(api.ProgressResponse{Status: "removing any unused layers"}) + err = deleteUnusedLayers(nil, deleteMap, false) + if err != nil { + return err + } + } + fn(api.ProgressResponse{Status: "success"}) return nil } @@ -869,18 +902,7 @@ func CopyModel(src, dest string) error { return nil } -func DeleteModel(name string) error { - mp := ParseModelPath(name) - manifest, _, err := GetManifest(mp) - if err != nil { - return err - } - deleteMap := make(map[string]bool) - for _, layer := range manifest.Layers { - deleteMap[layer.Digest] = true - } - deleteMap[manifest.Config.Digest] = true - +func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]bool, dryRun bool) error { fp, err := GetManifestPath() if err != nil { return err @@ -897,14 +919,13 @@ func DeleteModel(name string) error { fmp := ParseModelPath(tag) // skip the manifest we're trying to delete - if mp.GetFullTagname() == fmp.GetFullTagname() { + if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() { return nil } // save (i.e. delete from the deleteMap) any files used in other manifests manifest, _, err := GetManifest(fmp) if err != nil { - log.Printf("skipping file: %s", fp) return nil } @@ -928,14 +949,72 @@ func DeleteModel(name string) error { log.Printf("couldn't get file path for '%s': %v", k, err) continue } - if err := os.Remove(fp); err != nil { - log.Printf("couldn't remove file '%s': %v", fp, err) - continue + if !dryRun { + if err := os.Remove(fp); err != nil { + log.Printf("couldn't remove file '%s': %v", fp, err) + continue + } + } else { + log.Printf("wanted to remove: %s", fp) } } } - fp, err = mp.GetManifestPath(false) + return nil +} + +func PruneLayers() error { + deleteMap := make(map[string]bool) + p, err := GetBlobsPath("") + if err != nil { + return err + } + + blobs, err := os.ReadDir(p) + if err != nil { + log.Printf("couldn't read dir '%s': %v", p, err) + return err + } + + for _, blob := range blobs { + name := blob.Name() + if runtime.GOOS == "windows" { + name = strings.ReplaceAll(name, "-", ":") + } + deleteMap[name] = true + } + + log.Printf("total blobs: %d", len(deleteMap)) + + err = deleteUnusedLayers(nil, deleteMap, false) + if err != nil { + return err + } + + log.Printf("total unused blobs removed: %d", len(deleteMap)) + + return nil +} + +func DeleteModel(name string) error { + mp := ParseModelPath(name) + manifest, _, err := GetManifest(mp) + if err != nil { + return err + } + + deleteMap := make(map[string]bool) + for _, layer := range manifest.Layers { + deleteMap[layer.Digest] = true + } + deleteMap[manifest.Config.Digest] = true + + err = deleteUnusedLayers(&mp, deleteMap, false) + if err != nil { + return err + } + + fp, err := mp.GetManifestPath(false) if err != nil { return err } @@ -1114,13 +1193,34 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error { mp := ParseModelPath(name) + var manifest *ManifestV2 + var err error + var noprune string + + // build deleteMap to prune unused layers + deleteMap := make(map[string]bool) + + if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" { + manifest, _, err = GetManifest(mp) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + + if manifest != nil { + for _, l := range manifest.Layers { + deleteMap[l.Digest] = true + } + deleteMap[manifest.Config.Digest] = true + } + } + if mp.ProtocolScheme == "http" && !regOpts.Insecure { return fmt.Errorf("insecure protocol http") } fn(api.ProgressResponse{Status: "pulling manifest"}) - manifest, err := pullModelManifest(ctx, mp, regOpts) + manifest, err = pullModelManifest(ctx, mp, regOpts) if err != nil { return fmt.Errorf("pull model manifest: %s", err) } @@ -1140,7 +1240,9 @@ func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu }); err != nil { return err } + delete(deleteMap, layer.Digest) } + delete(deleteMap, manifest.Config.Digest) fn(api.ProgressResponse{Status: "verifying sha256 digest"}) for _, layer := range layers { @@ -1178,6 +1280,14 @@ func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu return err } + if noprune == "" { + fn(api.ProgressResponse{Status: "removing any unused layers"}) + err = deleteUnusedLayers(nil, deleteMap, false) + if err != nil { + return err + } + } + fn(api.ProgressResponse{Status: "success"}) return nil diff --git a/server/routes.go b/server/routes.go index 44871bba..472585ed 100644 --- a/server/routes.go +++ b/server/routes.go @@ -363,6 +363,7 @@ func DeleteModelHandler(c *gin.Context) { } return } + c.JSON(http.StatusOK, nil) } func ShowModelHandler(c *gin.Context) {