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

Merge pull request #536 from jmorganca/mxyng/redirect-uploads

Browse source 
explicitly follow upload redirects
This commit is contained in:
Michael Yang 2023-09-20 11:27:03 -07:00 committed by GitHub
commit 637f8bc6a5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 49 additions and 19 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,
})
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
}

64
server/upload.go
View file

@ -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()
// 95MB chunk size
chunkSize := 95 * 1024 * 1024
pw := ProgressWriter{
status: fmt.Sprintf("uploading %s", layer.Digest),
digest: layer.Digest,
@ -70,10 +79,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 +92,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 +126,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)