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.
This commit is contained in:
Jesse Gross 2024-08-05 17:13:52 -07:00
parent 685a53534b
commit 1829fb61bd
4 changed files with 61 additions and 35 deletions

View file

@ -250,6 +250,7 @@ func GetModel(name string) (*Model, error) {
Template: template.DefaultTemplate, Template: template.DefaultTemplate,
} }
if manifest.Config.Digest != "" {
filename, err := GetBlobsPath(manifest.Config.Digest) filename, err := GetBlobsPath(manifest.Config.Digest)
if err != nil { if err != nil {
return nil, err return nil, err
@ -264,6 +265,7 @@ func GetModel(name string) (*Model, error) {
if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil { if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
return nil, err return nil, err
} }
}
for _, layer := range manifest.Layers { for _, layer := range manifest.Layers {
filename, err := GetBlobsPath(layer.Digest) filename, err := GetBlobsPath(layer.Digest)
@ -781,7 +783,7 @@ func PruneLayers() error {
err = deleteUnusedLayers(nil, deleteMap) err = deleteUnusedLayers(nil, deleteMap)
if err != nil { 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 return nil
} }
@ -839,7 +841,9 @@ func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
var layers []*Layer var layers []*Layer
layers = append(layers, manifest.Layers...) layers = append(layers, manifest.Layers...)
layers = append(layers, manifest.Config) if manifest.Config.Digest != "" {
layers = append(layers, &manifest.Config)
}
for _, layer := range layers { for _, layer := range layers {
if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil { if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
@ -890,9 +894,11 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
for _, l := range manifest.Layers { for _, l := range manifest.Layers {
deleteMap[l.Digest] = struct{}{} deleteMap[l.Digest] = struct{}{}
} }
if manifest.Config.Digest != "" {
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")
@ -907,7 +913,9 @@ func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn fu
var layers []*Layer var layers []*Layer
layers = append(layers, manifest.Layers...) layers = append(layers, manifest.Layers...)
layers = append(layers, manifest.Config) if manifest.Config.Digest != "" {
layers = append(layers, &manifest.Config)
}
skipVerify := make(map[string]bool) skipVerify := make(map[string]bool)
for _, layer := range layers { 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"}) fn(api.ProgressResponse{Status: "removing any unused layers"})
err = deleteUnusedLayers(nil, deleteMap) err = deleteUnusedLayers(nil, deleteMap)
if err != nil { 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)}) fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
} }
} }

View file

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

View file

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

View file

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