baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

add autoprune to remove unused layers (#491)

Browse source
This commit is contained in:
Patrick Devine 2023-09-11 11:46:35 -07:00 • committed by GitHub
parent 3920e15386
commit e7e91cd71c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 136 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cmd/cmd.go
View file

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

148
server/images.go
View file

@ -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

1
server/routes.go
View file

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