From 1829fb61bd7a4186881714618f09b2877d0bc9a3 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Mon, 5 Aug 2024 17:13:52 -0700 Subject: [PATCH] manifest: Fix crash on startup when trying to clean up unused files (#5840) Currently if the config field is missing in the manifest file (or corrupted), Ollama will crash when it tries to read it. This can happen at startup or when pulling new models. This data is mostly just used for showing model information so we can be tolerant of it not being present - it is not required to run the models. Besides avoiding crashing, this also gives us the ability to restructure the config in the future by pulling it into the main manifest file. --- server/images.go | 40 ++++++++++++++++++++++++---------------- server/layer.go | 15 ++++++++++++++- server/manifest.go | 18 ++++++++++-------- server/routes.go | 23 +++++++++++++---------- 4 files changed, 61 insertions(+), 35 deletions(-) diff --git a/server/images.go b/server/images.go index 05875a88..7ed35995 100644 --- a/server/images.go +++ b/server/images.go @@ -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 { @@ -781,7 +783,7 @@ func PruneLayers() error { err = deleteUnusedLayers(nil, deleteMap) if err != nil { - slog.Info(fmt.Sprintf("couldn't remove unused layers: %v", err)) + slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err)) return nil } @@ -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,7 @@ 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 { - slog.Info(fmt.Sprintf("couldn't remove unused layers: %v", 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)}) } } diff --git a/server/layer.go b/server/layer.go index cc6709d2..a2b66782 100644 --- a/server/layer.go +++ b/server/layer.go @@ -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 diff --git a/server/manifest.go b/server/manifest.go index b8df11ef..b966ddbe 100644 --- a/server/manifest.go +++ b/server/manifest.go @@ -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, } diff --git a/server/routes.go b/server/routes.go index b9c66b65..e55eaa9d 100644 --- a/server/routes.go +++ b/server/routes.go @@ -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