From 499e9007a5d402484ed53c4e0c65e1a66a60a522 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Tue, 19 Sep 2023 14:22:54 -0700
Subject: [PATCH] pick chunksize based on location

---
 server/images.go |  4 ++--
 server/upload.go | 27 ++++++++++++++++++---------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/server/images.go b/server/images.go
index e60546d1..75ff1aab 100644
--- a/server/images.go
+++ b/server/images.go
@@ -1154,7 +1154,7 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu
 			Total:  layer.Size,
 		})
 
-		location, err := startUpload(ctx, mp, layer, regOpts)
+		location, chunkSize, err := startUpload(ctx, mp, layer, regOpts)
 		if err != nil {
 			log.Printf("couldn't start upload: %v", err)
 			return err
@@ -1171,7 +1171,7 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu
 			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)
 			return err
 		}
diff --git a/server/upload.go b/server/upload.go
index 161709b5..618195f7 100644
--- a/server/upload.go
+++ b/server/upload.go
@@ -14,7 +14,12 @@ import (
 	"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 = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs/uploads/")
 	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)
 	if err != nil {
 		log.Printf("couldn't start upload: %v", err)
-		return nil, err
+		return nil, 0, err
 	}
 	defer resp.Body.Close()
 
-	// Extract UUID location from header
-	location := resp.Header.Get("Location")
+	location := resp.Header.Get("Docker-Upload-Location")
+	chunkSize := redirectChunkSize
 	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 canceling uploads via DELETE
 
@@ -55,8 +66,6 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
 	}
 	defer f.Close()
 
-	// 95MiB chunk size
-	chunkSize := 95 * 1024 * 1024
 	pw := ProgressWriter{
 		status: fmt.Sprintf("uploading %s", layer.Digest),
 		digest: layer.Digest,