Repository: ollama/ollama
File: server/upload.go (Go)

Code before:
		headers.Set("Content-Range", fmt.Sprintf("%d-%d", offset, offset+sectionReader.Size()-1))
	}

	for try := 0; try < MaxRetries; try++ {
		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)
			token, err := getAuthToken(ctx, authRedir)
			if err != nil {
				return nil, err
			}

			opts.Token = token

			pw.completed = int(offset)
			sectionReader = io.NewSectionReader(r, offset, limit)
			continue
		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)
		}

		return resp, nil
	}

	return nil, fmt.Errorf("max retries exceeded")
}

type ProgressWriter struct {
	status    string

Review comment (bug):
Just thinking through if there was an error here. I think this will try calling PUT back on the original requestURL which would then give the temp redirect again. It does that dance 3 times then errors out. Is there any error handling we want to do here like auth/bad request stuff? If there is a bad request it might be tricky to debug.

Code after:
	headers.Set("Content-Type", "application/octet-stream")
	headers.Set("Content-Length", strconv.Itoa(int(limit)))
	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, 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)
			token, err := getAuthToken(ctx, authRedir)
			if err != nil {
				return nil, err
			}

			opts.Token = token

			pw.completed = int(offset)
			sectionReader = io.NewSectionReader(r, offset, limit)
			continue
		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)
		}

		return resp, nil
	}