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

fix noprune

Browse source
This commit is contained in:
Michael Yang 2024-08-14 14:37:51 -07:00
parent 0a8d6ea86d
commit b3f75fc812
1 changed files with 24 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
server/images.go
View file

@ -215,25 +215,20 @@ func GetManifest(mp ModelPath) (*Manifest, string, error) {
return nil, "", err return nil, "", err
} }
if _, err = os.Stat(fp); err != nil { f, err := os.Open(fp)
return nil, "", err
}
var manifest *Manifest
bts, err := os.ReadFile(fp)
if err != nil { if err != nil {
return nil, "", fmt.Errorf("couldn't open file '%s'", fp) return nil, "", err
} }
defer f.Close()
shaSum := sha256.Sum256(bts) sha256sum := sha256.New()
shaStr := hex.EncodeToString(shaSum[:])
if err := json.Unmarshal(bts, &manifest); err != nil { var manifest Manifest
if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&manifest); err != nil {
return nil, "", err return nil, "", err
} }
return manifest, shaStr, nil return &manifest, hex.EncodeToString(sha256sum.Sum(nil)), nil
} }
func GetModel(name string) (*Model, error) { func GetModel(name string) (*Model, error) {
@ -716,7 +711,7 @@ func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{})
// save (i.e. delete from the deleteMap) any files used in other manifests // save (i.e. delete from the deleteMap) any files used in other manifests
manifest, _, err := GetManifest(fmp) manifest, _, err := GetManifest(fmp)
if err != nil { if err != nil {
return err return fmt.Errorf("error reading manifest %s: %w", path, err)
} }
for _, layer := range manifest.Layers { for _, layer := range manifest.Layers {
@ -781,8 +776,7 @@ func PruneLayers() error {
slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap))) slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
err = deleteUnusedLayers(nil, deleteMap) if err := deleteUnusedLayers(nil, deleteMap); err != nil {
if err != nil {
slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err)) slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
return nil return nil
} }
@ -877,20 +871,14 @@ 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 { func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name) mp := ParseModelPath(name)
var manifest *Manifest
var err error
var noprune string
// build deleteMap to prune unused layers // build deleteMap to prune unused layers
deleteMap := make(map[string]struct{}) deleteMap := make(map[string]struct{})
manifest, _, err := GetManifest(mp)
if !envconfig.NoPrune() { if errors.Is(err, os.ErrNotExist) {
manifest, _, err = GetManifest(mp) // noop
if err != nil && !errors.Is(err, os.ErrNotExist) { } else if err != nil && !errors.Is(err, os.ErrNotExist) {
return err return err
} } else {
if manifest != nil {
for _, l := range manifest.Layers { for _, l := range manifest.Layers {
deleteMap[l.Digest] = struct{}{} deleteMap[l.Digest] = struct{}{}
} }
@ -898,7 +886,6 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
deleteMap[manifest.Config.Digest] = struct{}{} deleteMap[manifest.Config.Digest] = struct{}{}
} }
} }
}
if mp.ProtocolScheme == "http" && !regOpts.Insecure { if mp.ProtocolScheme == "http" && !regOpts.Insecure {
return errors.New("insecure protocol http") return errors.New("insecure protocol http")
@ -975,11 +962,9 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
return err return err
} }
if noprune == "" { if !envconfig.NoPrune() && len(deleteMap) > 0 {
fn(api.ProgressResponse{Status: "removing any unused layers"}) fn(api.ProgressResponse{Status: "removing unused layers"})
err = deleteUnusedLayers(nil, deleteMap) if err := deleteUnusedLayers(nil, deleteMap); err != nil {
if err != nil {
slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)}) fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
} }
} }
@ -1000,12 +985,12 @@ func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *registryOptio
} }
defer resp.Body.Close() defer resp.Body.Close()
var m *Manifest var m Manifest
if err := json.NewDecoder(resp.Body).Decode(&m); err != nil { if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
return nil, err return nil, err
} }
return m, err return &m, err
} }
// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer