Merge pull request #6145 from ollama/jessegross/bug5840

Fix crash on startup when trying to clean up unused files (#5840)
This commit is contained in:
Jesse Gross 2024-08-07 11:24:15 -07:00 committed by GitHub
commit 69eb06c40e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 64 additions and 37 deletions

View file

@ -250,19 +250,21 @@ func GetModel(name string) (*Model, error) {
Template: template.DefaultTemplate,
}
filename, err := GetBlobsPath(manifest.Config.Digest)
if err != nil {
return nil, err
}
if manifest.Config.Digest != "" {
filename, err := GetBlobsPath(manifest.Config.Digest)
if err != nil {
return nil, err
}
configFile, err := os.Open(filename)
if err != nil {
return nil, err
}
defer configFile.Close()
configFile, err := os.Open(filename)
if err != nil {
return nil, err
}
defer configFile.Close()
if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
return nil, err
if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
return nil, err
}
}
for _, layer := range manifest.Layers {
@ -714,8 +716,7 @@ func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{})
// save (i.e. delete from the deleteMap) any files used in other manifests
manifest, _, err := GetManifest(fmp)
if err != nil {
//nolint:nilerr
return nil
return err
}
for _, layer := range manifest.Layers {
@ -782,7 +783,8 @@ func PruneLayers() error {
err = deleteUnusedLayers(nil, deleteMap)
if err != nil {
return err
slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
return nil
}
slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
@ -839,7 +841,9 @@ func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
var layers []*Layer
layers = append(layers, manifest.Layers...)
layers = append(layers, manifest.Config)
if manifest.Config.Digest != "" {
layers = append(layers, &manifest.Config)
}
for _, layer := range layers {
if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
@ -890,7 +894,9 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
for _, l := range manifest.Layers {
deleteMap[l.Digest] = struct{}{}
}
deleteMap[manifest.Config.Digest] = struct{}{}
if manifest.Config.Digest != "" {
deleteMap[manifest.Config.Digest] = struct{}{}
}
}
}
@ -907,7 +913,9 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
var layers []*Layer
layers = append(layers, manifest.Layers...)
layers = append(layers, manifest.Config)
if manifest.Config.Digest != "" {
layers = append(layers, &manifest.Config)
}
skipVerify := make(map[string]bool)
for _, layer := range layers {
@ -971,7 +979,8 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
fn(api.ProgressResponse{Status: "removing any unused layers"})
err = deleteUnusedLayers(nil, deleteMap)
if err != nil {
return err
slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
}
}

View file

@ -2,6 +2,7 @@ package server
import (
"crypto/sha256"
"errors"
"fmt"
"io"
"os"
@ -61,6 +62,10 @@ func NewLayer(r io.Reader, mediatype string) (*Layer, error) {
}
func NewLayerFromLayer(digest, mediatype, from string) (*Layer, error) {
if digest == "" {
return nil, errors.New("creating new layer from layer with empty digest")
}
blob, err := GetBlobsPath(digest)
if err != nil {
return nil, err
@ -81,6 +86,10 @@ func NewLayerFromLayer(digest, mediatype, from string) (*Layer, error) {
}
func (l *Layer) Open() (io.ReadSeekCloser, error) {
if l.Digest == "" {
return nil, errors.New("opening layer with empty digest")
}
blob, err := GetBlobsPath(l.Digest)
if err != nil {
return nil, err
@ -90,13 +99,17 @@ func (l *Layer) Open() (io.ReadSeekCloser, error) {
}
func (l *Layer) Remove() error {
if l.Digest == "" {
return nil
}
ms, err := Manifests()
if err != nil {
return err
}
for _, m := range ms {
for _, layer := range append(m.Layers, m.Config) {
for _, layer := range append(m.Layers, &m.Config) {
if layer.Digest == l.Digest {
// something is using this layer
return nil

View file

@ -16,7 +16,7 @@ import (
type Manifest struct {
SchemaVersion int `json:"schemaVersion"`
MediaType string `json:"mediaType"`
Config *Layer `json:"config"`
Config Layer `json:"config"`
Layers []*Layer `json:"layers"`
filepath string
@ -25,7 +25,7 @@ type Manifest struct {
}
func (m *Manifest) Size() (size int64) {
for _, layer := range append(m.Layers, m.Config) {
for _, layer := range append(m.Layers, &m.Config) {
size += layer.Size
}
@ -46,11 +46,13 @@ func (m *Manifest) Remove() error {
}
func (m *Manifest) RemoveLayers() error {
for _, layer := range append(m.Layers, m.Config) {
if err := layer.Remove(); errors.Is(err, os.ErrNotExist) {
slog.Debug("layer does not exist", "digest", layer.Digest)
} else if err != nil {
return err
for _, layer := range append(m.Layers, &m.Config) {
if layer.Digest != "" {
if err := layer.Remove(); errors.Is(err, os.ErrNotExist) {
slog.Debug("layer does not exist", "digest", layer.Digest)
} else if err != nil {
return err
}
}
}
@ -113,7 +115,7 @@ func WriteManifest(name model.Name, config *Layer, layers []*Layer) error {
m := Manifest{
SchemaVersion: 2,
MediaType: "application/vnd.docker.distribution.manifest.v2+json",
Config: config,
Config: *config,
Layers: layers,
}

View file

@ -824,17 +824,20 @@ func (s *Server) ListModelsHandler(c *gin.Context) {
models := []api.ListModelResponse{}
for n, m := range ms {
f, err := m.Config.Open()
if err != nil {
slog.Warn("bad manifest filepath", "name", n, "error", err)
continue
}
defer f.Close()
var cf ConfigV2
if err := json.NewDecoder(f).Decode(&cf); err != nil {
slog.Warn("bad manifest config", "name", n, "error", err)
continue
if m.Config.Digest != "" {
f, err := m.Config.Open()
if err != nil {
slog.Warn("bad manifest filepath", "name", n, "error", err)
continue
}
defer f.Close()
if err := json.NewDecoder(f).Decode(&cf); err != nil {
slog.Warn("bad manifest config", "name", n, "error", err)
continue
}
}
// tag should never be masked