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

tidy intermediate blobs

Browse source
This commit is contained in:
Michael Yang 2024-05-20 14:58:27 -07:00
parent 8800c8a59b
commit f36f1d6be9
3 changed files with 11 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
server/images.go
View file

@ -341,8 +341,8 @@ func CreateModel(ctx context.Context, name, modelFileDir, quantization string, m
} }
} else if strings.HasPrefix(c.Args, "@") { } else if strings.HasPrefix(c.Args, "@") {
digest := strings.TrimPrefix(c.Args, "@") digest := strings.TrimPrefix(c.Args, "@")
if ib, ok := intermediateBlobs.Load(digest); ok { if ib, ok := intermediateBlobs[digest]; ok {
p, err := GetBlobsPath(ib.(string)) p, err := GetBlobsPath(ib)
if err != nil { if err != nil {
return err return err
} }
@ -352,8 +352,8 @@ func CreateModel(ctx context.Context, name, modelFileDir, quantization string, m
} else if err != nil { } else if err != nil {
return err return err
} else { } else {
fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib.(string))}) fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib)})
digest = ib.(string) digest = ib
} }
} }
@ -415,14 +415,10 @@ func CreateModel(ctx context.Context, name, modelFileDir, quantization string, m
return err return err
} }
f16digest := baseLayer.Layer.Digest
baseLayer.Layer, err = NewLayer(temp, baseLayer.Layer.MediaType) baseLayer.Layer, err = NewLayer(temp, baseLayer.Layer.MediaType)
if err != nil { if err != nil {
return err return err
} }
intermediateBlobs.Store(f16digest, baseLayer.Layer.Digest)
} }
} }

7
server/model.go
View file

@ -10,7 +10,6 @@ import (
"net/http" "net/http"
"os" "os"
"path/filepath" "path/filepath"
"sync"
"github.com/ollama/ollama/api" "github.com/ollama/ollama/api"
"github.com/ollama/ollama/convert" "github.com/ollama/ollama/convert"
@ -18,7 +17,7 @@ import (
"github.com/ollama/ollama/types/model" "github.com/ollama/ollama/types/model"
) )
var intermediateBlobs sync.Map var intermediateBlobs map[string]string = make(map[string]string)
type layerWithGGML struct { type layerWithGGML struct {
*Layer *Layer
@ -169,7 +168,7 @@ func parseFromZipFile(_ context.Context, file *os.File, digest string, fn func(a
layer, err := NewLayer(temp, "application/vnd.ollama.image.model") layer, err := NewLayer(temp, "application/vnd.ollama.image.model")
if err != nil { if err != nil {
return nil, fmt.Errorf("aaa: %w", err) return nil, err
} }
bin, err := layer.Open() bin, err := layer.Open()
@ -185,7 +184,7 @@ func parseFromZipFile(_ context.Context, file *os.File, digest string, fn func(a
layers = append(layers, &layerWithGGML{layer, ggml}) layers = append(layers, &layerWithGGML{layer, ggml})
intermediateBlobs.Store(digest, layer.Digest) intermediateBlobs[digest] = layer.Digest
return layers, nil return layers, nil
} }

8
server/routes.go
View file

@ -841,16 +841,16 @@ func (s *Server) HeadBlobHandler(c *gin.Context) {
} }
func (s *Server) CreateBlobHandler(c *gin.Context) { func (s *Server) CreateBlobHandler(c *gin.Context) {
ib, ok := intermediateBlobs.Load(c.Param("digest")) if ib, ok := intermediateBlobs[c.Param("digest")]; ok {
if ok { p, err := GetBlobsPath(ib)
p, err := GetBlobsPath(ib.(string))
if err != nil { if err != nil {
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return return
} }
if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) { if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
intermediateBlobs.Delete(c.Param("digest")) slog.Info("evicting intermediate blob which no longer exists", "digest", ib)
delete(intermediateBlobs, c.Param("digest"))
} else if err != nil { } else if err != nil {
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return return