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

prune partial downloads (#4272)

Browse source
This commit is contained in:
Jeffrey Morgan 2024-05-09 16:35:20 -07:00 committed by GitHub
parent 0e1ba65855
commit 302d7fdbf3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
server/images.go
View file

@ -565,7 +565,7 @@ func CreateModel(ctx context.Context, name, modelFileDir, quantization string, m
}
if !envconfig.NoPrune {
if err := deleteUnusedLayers(nil, unref, false); err != nil {
if err := deleteUnusedLayers(nil, unref); err != nil {
return err
}
}
@ -613,7 +613,7 @@ func CopyModel(src, dst model.Name) error {
return err
}
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}) error {
fp, err := GetManifestPath()
if err != nil {
return err
@ -660,13 +660,9 @@ func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{},
slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
continue
}
if !dryRun {
if err := os.Remove(fp); err != nil {
slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
continue
}
} else {
slog.Info(fmt.Sprintf("wanted to remove: %s", fp))
if err := os.Remove(fp); err != nil {
slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
continue
}
}
@ -689,14 +685,25 @@ func PruneLayers() error {
for _, blob := range blobs {
name := blob.Name()
name = strings.ReplaceAll(name, "-", ":")
if strings.HasPrefix(name, "sha256:") {
deleteMap[name] = struct{}{}
_, err := GetBlobsPath(name)
if err != nil {
if errors.Is(err, ErrInvalidDigestFormat) {
// remove invalid blobs (e.g. partial downloads)
if err := os.Remove(filepath.Join(p, blob.Name())); err != nil {
slog.Error("couldn't remove blob", "blob", blob.Name(), "error", err)
}
}
continue
}
deleteMap[name] = struct{}{}
}
slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
err = deleteUnusedLayers(nil, deleteMap, false)
err = deleteUnusedLayers(nil, deleteMap)
if err != nil {
return err
}
@ -752,7 +759,7 @@ func DeleteModel(name string) error {
}
deleteMap[manifest.Config.Digest] = struct{}{}
err = deleteUnusedLayers(&mp, deleteMap, false)
err = deleteUnusedLayers(&mp, deleteMap)
if err != nil {
return err
}
@ -912,7 +919,7 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
if noprune == "" {
fn(api.ProgressResponse{Status: "removing any unused layers"})
err = deleteUnusedLayers(nil, deleteMap, false)
err = deleteUnusedLayers(nil, deleteMap)
if err != nil {
return err
}

3
server/modelpath.go
View file

@ -154,9 +154,6 @@ func GetBlobsPath(digest string) (string, error) {
// only accept actual sha256 digests
pattern := "^sha256[:-][0-9a-fA-F]{64}$"
re := regexp.MustCompile(pattern)
if err != nil {
return "", err
}
if digest != "" && !re.MatchString(digest) {
return "", ErrInvalidDigestFormat