Merge pull request #428 from jmorganca/mxyng/upload-chunks

update upload chunks

api/client.go

@@ -29,7 +29,7 @@ type Client struct {
 }
 
 func checkError(resp *http.Response, body []byte) error {
-	if resp.StatusCode >= 200 && resp.StatusCode < 400 {
+	if resp.StatusCode < http.StatusBadRequest {
 		return nil
 	}
 
@@ -165,7 +165,7 @@ func (c *Client) stream(ctx context.Context, method, path string, data any, fn f
 			return fmt.Errorf(errorResponse.Error)
 		}
 
-		if response.StatusCode >= 400 {
+		if response.StatusCode >= http.StatusBadRequest {
 			return StatusError{
 				StatusCode:   response.StatusCode,
 				Status:       response.Status,

server/auth.go

@@ -109,7 +109,7 @@ func getAuthToken(ctx context.Context, redirData AuthRedirect, regOpts *Registry
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusOK {
+	if resp.StatusCode >= http.StatusBadRequest {
 		body, _ := io.ReadAll(resp.Body)
 		return "", fmt.Errorf("on pull registry responded with code %d: %s", resp.StatusCode, body)
 	}

server/download.go

@@ -168,7 +168,7 @@ func doDownload(ctx context.Context, opts downloadOpts, f *FileDownload) error {
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
+	if resp.StatusCode >= http.StatusBadRequest {
 		body, _ := io.ReadAll(resp.Body)
 		return fmt.Errorf("%w: on download registry responded with code %d: %v", errDownload, resp.StatusCode, string(body))
 	}

server/images.go

@@ -980,7 +980,7 @@ func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn fu
 			continue
 		}
 
-		if err := uploadBlobChunked(ctx, mp, location, layer, regOpts, fn); err != nil {
+		if err := uploadBlobChunked(ctx, location, layer, regOpts, fn); err != nil {
 			log.Printf("error uploading blob: %v", err)
 			return err
 		}
@@ -1092,11 +1092,11 @@ func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *RegistryOptio
 	}
 	defer resp.Body.Close()
 
-	// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
-	if resp.StatusCode != http.StatusOK {
+	if resp.StatusCode >= http.StatusBadRequest {
 		if resp.StatusCode == http.StatusNotFound {
 			return nil, fmt.Errorf("model not found")
 		}
+
 		body, _ := io.ReadAll(resp.Body)
 		return nil, fmt.Errorf("on pull registry responded with code %d: %s", resp.StatusCode, body)
 	}
@@ -1157,7 +1157,7 @@ func checkBlobExistence(ctx context.Context, mp ModelPath, digest string, regOpt
 	defer resp.Body.Close()
 
 	// Check for success: If the blob exists, the Docker registry will respond with a 200 OK
-	return resp.StatusCode == http.StatusOK, nil
+	return resp.StatusCode < http.StatusBadRequest, nil
 }
 
 func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
@@ -1171,10 +1171,8 @@ func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.UR
 
 		status = resp.Status
 
-		switch resp.StatusCode {
-		case http.StatusAccepted, http.StatusCreated:
-			return resp, nil
-		case http.StatusUnauthorized:
+		switch {
+		case resp.StatusCode == http.StatusUnauthorized:
 			auth := resp.Header.Get("www-authenticate")
 			authRedir := ParseAuthRedirectString(auth)
 			token, err := getAuthToken(ctx, authRedir, regOpts)
@@ -1190,9 +1188,11 @@ func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.UR
 			}
 
 			continue
-		default:
+		case resp.StatusCode >= http.StatusBadRequest:
 			body, _ := io.ReadAll(resp.Body)
 			return nil, fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
+		default:
+			return resp, nil
 		}
 	}
 

server/upload.go

@@ -40,7 +40,7 @@ func startUpload(ctx context.Context, mp ModelPath, layer *Layer, regOpts *Regis
 	return url.Parse(location)
 }
 
-func uploadBlobChunked(ctx context.Context, mp ModelPath, requestURL *url.URL, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
+func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
 	// TODO allow resumability
 	// TODO allow canceling uploads via DELETE
 
@@ -55,48 +55,88 @@ func uploadBlobChunked(ctx context.Context, mp ModelPath, requestURL *url.URL, l
 	}
 	defer f.Close()
 
-	var completed int64
-	chunkSize := 10 * 1024 * 1024
+	// 95MB chunk size
+	chunkSize := 95 * 1024 * 1024
 
-	for {
-		chunk := int64(layer.Size) - completed
+	for offset := int64(0); offset < int64(layer.Size); {
+		chunk := int64(layer.Size) - offset
 		if chunk > int64(chunkSize) {
 			chunk = int64(chunkSize)
 		}
 
-		sectionReader := io.NewSectionReader(f, int64(completed), chunk)
-
-		headers := make(http.Header)
-		headers.Set("Content-Type", "application/octet-stream")
-		headers.Set("Content-Length", strconv.Itoa(int(chunk)))
-		headers.Set("Content-Range", fmt.Sprintf("%d-%d", completed, completed+sectionReader.Size()-1))
-		resp, err := makeRequestWithRetry(ctx, "PATCH", requestURL, headers, sectionReader, regOpts)
-		if err != nil && !errors.Is(err, io.EOF) {
-			fn(api.ProgressResponse{
-				Status:    fmt.Sprintf("error uploading chunk: %v", err),
-				Digest:    layer.Digest,
-				Total:     layer.Size,
-				Completed: int(completed),
-			})
-
-			return err
-		}
-		defer resp.Body.Close()
-
-		completed += sectionReader.Size()
-		fn(api.ProgressResponse{
-			Status:    fmt.Sprintf("uploading %s", layer.Digest),
-			Digest:    layer.Digest,
-			Total:     layer.Size,
-			Completed: int(completed),
-		})
-
-		requestURL, err = url.Parse(resp.Header.Get("Location"))
-		if err != nil {
-			return err
-		}
+		sectionReader := io.NewSectionReader(f, int64(offset), chunk)
+		for try := 0; try < MaxRetries; try++ {
+			r, w := io.Pipe()
+			defer r.Close()
+			go func() {
+				defer w.Close()
+
+				for chunked := int64(0); chunked < chunk; {
+					n, err := io.CopyN(w, sectionReader, 1024*1024)
+					if err != nil && !errors.Is(err, io.EOF) {
+						fn(api.ProgressResponse{
+							Status:    fmt.Sprintf("error reading chunk: %v", err),
+							Digest:    layer.Digest,
+							Total:     layer.Size,
+							Completed: int(offset),
+						})
+
+						return
+					}
+
+					chunked += n
+					fn(api.ProgressResponse{
+						Status:    fmt.Sprintf("uploading %s", layer.Digest),
+						Digest:    layer.Digest,
+						Total:     layer.Size,
+						Completed: int(offset) + int(chunked),
+					})
+				}
+			}()
+
+			headers := make(http.Header)
+			headers.Set("Content-Type", "application/octet-stream")
+			headers.Set("Content-Length", strconv.Itoa(int(chunk)))
+			headers.Set("Content-Range", fmt.Sprintf("%d-%d", offset, offset+sectionReader.Size()-1))
+			resp, err := makeRequest(ctx, "PATCH", requestURL, headers, r, regOpts)
+			if err != nil && !errors.Is(err, io.EOF) {
+				fn(api.ProgressResponse{
+					Status:    fmt.Sprintf("error uploading chunk: %v", err),
+					Digest:    layer.Digest,
+					Total:     layer.Size,
+					Completed: int(offset),
+				})
+
+				return err
+			}
+			defer resp.Body.Close()
+
+			switch {
+			case resp.StatusCode == http.StatusUnauthorized:
+				auth := resp.Header.Get("www-authenticate")
+				authRedir := ParseAuthRedirectString(auth)
+				token, err := getAuthToken(ctx, authRedir, regOpts)
+				if err != nil {
+					return err
+				}
+
+				regOpts.Token = token
+				if _, err := sectionReader.Seek(0, io.SeekStart); err != nil {
+					return err
+				}
+
+				continue
+			case resp.StatusCode >= http.StatusBadRequest:
+				body, _ := io.ReadAll(resp.Body)
+				return fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
+			}
+
+			offset += sectionReader.Size()
+			requestURL, err = url.Parse(resp.Header.Get("Location"))
+			if err != nil {
+				return err
+			}
 
-		if completed >= int64(layer.Size) {
 			break
 		}
 	}
@@ -117,7 +157,7 @@ func uploadBlobChunked(ctx context.Context, mp ModelPath, requestURL *url.URL, l
 	}
 	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusCreated {
+	if resp.StatusCode >= http.StatusBadRequest {
 		body, _ := io.ReadAll(resp.Body)
 		return fmt.Errorf("on finish upload registry responded with code %d: %v", resp.StatusCode, string(body))
 	}