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

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

Browse source 
update upload chunks
This commit is contained in:
Michael Yang 2023-08-30 07:47:17 -07:00 • committed by GitHub
commit 982c535428
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 87 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
api/client.go
View file

@ -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,

2
server/auth.go
View file

@ -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)
}

2
server/download.go
View file

@ -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))
}

18
server/images.go
View file

@ -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
}
}

76
server/upload.go
View file

@ -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)
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", completed, completed+sectionReader.Size()-1))
resp, err := makeRequestWithRetry(ctx, "PATCH", requestURL, headers, sectionReader, regOpts)
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(completed),
Completed: int(offset),
})
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),
})
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))
}