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

fix retry on unauthorized chunk

Browse source
This commit is contained in:
Michael Yang 2023-09-07 12:01:50 -07:00
parent f0f4943577
commit bf146fb072
1 changed files with 25 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
server/upload.go
View file

@ -66,12 +66,19 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
sectionReader := io.NewSectionReader(f, int64(offset), chunk)
for try := 0; try < MaxRetries; try++ {
ch := make(chan error, 1)
r, w := io.Pipe()
defer r.Close()
go func() {
defer w.Close()
for chunked := int64(0); chunked < chunk; {
select {
case err := <-ch:
log.Printf("chunk interrupted: %v", err)
return
default:
n, err := io.CopyN(w, sectionReader, 1024*1024)
if err != nil && !errors.Is(err, io.EOF) {
fn(api.ProgressResponse{
@ -92,6 +99,7 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
Completed: int(offset) + int(chunked),
})
}
}
}()
headers := make(http.Header)
@ -113,6 +121,8 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
switch {
case resp.StatusCode == http.StatusUnauthorized:
ch <- errors.New("unauthorized")
auth := resp.Header.Get("www-authenticate")
authRedir := ParseAuthRedirectString(auth)
token, err := getAuthToken(ctx, authRedir, regOpts)
@ -121,10 +131,7 @@ func uploadBlobChunked(ctx context.Context, requestURL *url.URL, layer *Layer, r
}
regOpts.Token = token
if _, err := sectionReader.Seek(0, io.SeekStart); err != nil {
return err
}
sectionReader = io.NewSectionReader(f, int64(offset), chunk)
continue
case resp.StatusCode >= http.StatusBadRequest:
body, _ := io.ReadAll(resp.Body)