add autoprune to remove unused layers (#491)

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)
 }
 

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

server/routes.go

@@ -363,6 +363,7 @@ func DeleteModelHandler(c *gin.Context) {
 		}
 		return
 	}
+	c.JSON(http.StatusOK, nil)
 }
 
 func ShowModelHandler(c *gin.Context) {