From aa45d7c1df895c444aacefdec7e5f13d09fff1a4 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Thu, 14 Sep 2023 15:42:50 -0700 Subject: [PATCH 1/2] draft: explicitly follow upload redirects --- server/upload.go | 39 ++++++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/server/upload.go b/server/upload.go index 1f02cef3..161709b5 100644 --- a/server/upload.go +++ b/server/upload.go @@ -55,7 +55,7 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r } defer f.Close() - // 95MB chunk size + // 95MiB chunk size chunkSize := 95 * 1024 * 1024 pw := ProgressWriter{ status: fmt.Sprintf("uploading %s", layer.Digest), @@ -70,10 +70,10 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r chunk = int64(chunkSize) } - resp, err := uploadBlobChunk(ctx, requestURL, f, offset, chunk, regOpts, &pw) + resp, err := uploadBlobChunk(ctx, http.MethodPatch, requestURL, f, offset, chunk, regOpts, &pw) if err != nil { fn(api.ProgressResponse{ - Status: fmt.Sprintf("error uploading limit: %v", err), + Status: fmt.Sprintf("error uploading chunk: %v", err), Digest: layer.Digest, Total: layer.Size, Completed: int(offset), @@ -83,12 +83,15 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r } offset += chunk - location, err := resp.Location() + location := resp.Header.Get("Docker-Upload-Location") + if location == "" { + location = resp.Header.Get("Location") + } + + requestURL, err = url.Parse(location) if err != nil { return err } - - requestURL = location } values := requestURL.Query() @@ -114,22 +117,40 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r return nil } -func uploadBlobChunk(ctx context.Context, requestURL *url.URL, r io.ReaderAt, offset, limit int64, opts *RegistryOptions, pw *ProgressWriter) (*http.Response, error) { +func uploadBlobChunk(ctx context.Context, method string, requestURL *url.URL, r io.ReaderAt, offset, limit int64, opts *RegistryOptions, pw *ProgressWriter) (*http.Response, error) { sectionReader := io.NewSectionReader(r, int64(offset), limit) headers := make(http.Header) headers.Set("Content-Type", "application/octet-stream") headers.Set("Content-Length", strconv.Itoa(int(limit))) - headers.Set("Content-Range", fmt.Sprintf("%d-%d", offset, offset+sectionReader.Size()-1)) + headers.Set("X-Redirect-Uploads", "1") + + if method == http.MethodPatch { + headers.Set("Content-Range", fmt.Sprintf("%d-%d", offset, offset+sectionReader.Size()-1)) + } for try := 0; try < MaxRetries; try++ { - resp, err := makeRequest(ctx, "PATCH", requestURL, headers, io.TeeReader(sectionReader, pw), opts) + resp, err := makeRequest(ctx, method, requestURL, headers, io.TeeReader(sectionReader, pw), opts) if err != nil && !errors.Is(err, io.EOF) { return nil, err } defer resp.Body.Close() switch { + case resp.StatusCode == http.StatusTemporaryRedirect: + location, err := resp.Location() + if err != nil { + return nil, err + } + + pw.completed = int(offset) + if _, err := uploadBlobChunk(ctx, http.MethodPut, location, r, offset, limit, nil, pw); err != nil { + // retry + log.Printf("retrying redirected upload: %v", err) + continue + } + + return resp, nil case resp.StatusCode == http.StatusUnauthorized: auth := resp.Header.Get("www-authenticate") authRedir := ParseAuthRedirectString(auth) From 499e9007a5d402484ed53c4e0c65e1a66a60a522 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Tue, 19 Sep 2023 14:22:54 -0700 Subject: [PATCH 2/2] 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,