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

pick chunksize based on location

Browse source
This commit is contained in:
Michael Yang 2023-09-19 14:22:54 -07:00
parent aa45d7c1df
commit 499e9007a5
2 changed files with 20 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
server/images.go
View file

@ -1154,7 +1154,7 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu
Total: layer.Size, Total: layer.Size,
}) })
location, err := startUpload(ctx, mp, layer, regOpts) location, chunkSize, err := startUpload(ctx, mp, layer, regOpts)
if err != nil { if err != nil {
log.Printf("couldn't start upload: %v", err) log.Printf("couldn't start upload: %v", err)
return err return err
@ -1171,7 +1171,7 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu
continue continue
} }
if err := uploadBlobChunked(ctx, location, layer, regOpts, fn); err != nil { if err := uploadBlob(ctx, location, layer, chunkSize, regOpts, fn); err != nil {
log.Printf("error uploading blob: %v", err) log.Printf("error uploading blob: %v", err)
return err return err
} }

27
server/upload.go
View file

@ -14,7 +14,12 @@ import (
"github.com/jmorganca/ollama/api" "github.com/jmorganca/ollama/api"
) )
func startUpload(ctx context.Context, mp ModelPath, layer *Layer, regOpts *RegistryOptions) (*url.URL, error) { const (
redirectChunkSize = 1024 * 1024 * 1024
regularChunkSize = 95 * 1024 * 1024
)
func startUpload(ctx context.Context, mp ModelPath, layer *Layer, regOpts *RegistryOptions) (*url.URL, int64, error) {
requestURL := mp.BaseURL() requestURL := mp.BaseURL()
requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs/uploads/") requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs/uploads/")
if layer.From != "" { if layer.From != "" {
@ -27,20 +32,26 @@ func startUpload(ctx context.Context, mp ModelPath, layer *Layer, regOpts *Regis
resp, err := makeRequestWithRetry(ctx, "POST", requestURL, nil, nil, regOpts) resp, err := makeRequestWithRetry(ctx, "POST", requestURL, nil, nil, regOpts)
if err != nil { if err != nil {
log.Printf("couldn't start upload: %v", err) log.Printf("couldn't start upload: %v", err)
return nil, err return nil, 0, err
} }
defer resp.Body.Close() defer resp.Body.Close()
// Extract UUID location from header location := resp.Header.Get("Docker-Upload-Location")
location := resp.Header.Get("Location") chunkSize := redirectChunkSize
if location == "" { if location == "" {
return nil, fmt.Errorf("location header is missing in response") location = resp.Header.Get("Location")
chunkSize = regularChunkSize
} }
return url.Parse(location) locationURL, err := url.Parse(location)
if err != nil {
return nil, 0, err
}
return locationURL, int64(chunkSize), nil
} }
func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error { func uploadBlob(ctx context.Context, requestURL *url.URL, layer *Layer, chunkSize int64, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
// TODO allow resumability // TODO allow resumability
// TODO allow canceling uploads via DELETE // TODO allow canceling uploads via DELETE
@ -55,8 +66,6 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
} }
defer f.Close() defer f.Close()
// 95MiB chunk size
chunkSize := 95 * 1024 * 1024
pw := ProgressWriter{ pw := ProgressWriter{
status: fmt.Sprintf("uploading %s", layer.Digest), status: fmt.Sprintf("uploading %s", layer.Digest),
digest: layer.Digest, digest: layer.Digest,