upload.go 9.02 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
package server

import (
	"context"
5
	"crypto/md5"
Michael Yang's avatar
Michael Yang committed
6
7
	"errors"
	"fmt"
8
	"hash"
Michael Yang's avatar
Michael Yang committed
9
	"io"
10
	"log/slog"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
11
	"math"
Michael Yang's avatar
Michael Yang committed
12
13
14
	"net/http"
	"net/url"
	"os"
Michael Yang's avatar
lint  
Michael Yang committed
15
	"strconv"
Michael Yang's avatar
Michael Yang committed
16
	"sync"
Michael Yang's avatar
Michael Yang committed
17
18
	"sync/atomic"
	"time"
Michael Yang's avatar
Michael Yang committed
19

Michael Yang's avatar
lint  
Michael Yang committed
20
21
	"golang.org/x/sync/errgroup"

22
23
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
24
25
)

Michael Yang's avatar
Michael Yang committed
26
27
28
var blobUploadManager sync.Map

type blobUpload struct {
29
	Layer
Michael Yang's avatar
Michael Yang committed
30
31
32
33
34
35
36
37
38
39

	Total     int64
	Completed atomic.Int64

	Parts []blobUploadPart

	nextURL chan *url.URL

	context.CancelFunc

40
41
	file *os.File

Michael Yang's avatar
Michael Yang committed
42
43
44
45
46
	done       bool
	err        error
	references atomic.Int32
}

47
const (
48
	numUploadParts          = 16
Michael Yang's avatar
Michael Yang committed
49
50
	minUploadPartSize int64 = 100 * format.MegaByte
	maxUploadPartSize int64 = 1000 * format.MegaByte
51
52
)

Michael Yang's avatar
Michael Yang committed
53
func (b *blobUpload) Prepare(ctx context.Context, requestURL *url.URL, opts *registryOptions) error {
Michael Yang's avatar
Michael Yang committed
54
55
56
57
58
59
	p, err := GetBlobsPath(b.Digest)
	if err != nil {
		return err
	}

	if b.From != "" {
Michael Yang's avatar
Michael Yang committed
60
		values := requestURL.Query()
Michael Yang's avatar
Michael Yang committed
61
		values.Add("mount", b.Digest)
62
		values.Add("from", ParseModelPath(b.From).GetNamespaceRepository())
Michael Yang's avatar
Michael Yang committed
63
64
65
		requestURL.RawQuery = values.Encode()
	}

Michael Yang's avatar
Michael Yang committed
66
	resp, err := makeRequestWithRetry(ctx, http.MethodPost, requestURL, nil, nil, opts)
Michael Yang's avatar
Michael Yang committed
67
	if err != nil {
Michael Yang's avatar
Michael Yang committed
68
		return err
Michael Yang's avatar
Michael Yang committed
69
70
71
	}
	defer resp.Body.Close()

72
	location := resp.Header.Get("Docker-Upload-Location")
Michael Yang's avatar
Michael Yang committed
73
	if location == "" {
74
75
76
		location = resp.Header.Get("Location")
	}

Michael Yang's avatar
Michael Yang committed
77
	fi, err := os.Stat(p)
78
	if err != nil {
Michael Yang's avatar
Michael Yang committed
79
		return err
Michael Yang's avatar
Michael Yang committed
80
81
	}

Michael Yang's avatar
Michael Yang committed
82
83
	b.Total = fi.Size()

Michael Yang's avatar
Michael Yang committed
84
85
86
87
88
89
90
91
	// http.StatusCreated indicates a blob has been mounted
	// ref: https://distribution.github.io/distribution/spec/api/#cross-repository-blob-mount
	if resp.StatusCode == http.StatusCreated {
		b.Completed.Store(b.Total)
		b.done = true
		return nil
	}

Michael Yang's avatar
Michael Yang committed
92
	size := b.Total / numUploadParts
Michael Yang's avatar
Michael Yang committed
93
94
95
96
97
98
99
100
101
102
103
104
105
	switch {
	case size < minUploadPartSize:
		size = minUploadPartSize
	case size > maxUploadPartSize:
		size = maxUploadPartSize
	}

	var offset int64
	for offset < fi.Size() {
		if offset+size > fi.Size() {
			size = fi.Size() - offset
		}

Michael Yang's avatar
Michael Yang committed
106
		// set part.N to the current number of parts
107
		b.Parts = append(b.Parts, blobUploadPart{N: len(b.Parts), Offset: offset, Size: size})
Michael Yang's avatar
Michael Yang committed
108
109
		offset += size
	}
Michael Yang's avatar
Michael Yang committed
110

111
112
113
	if len(b.Parts) > 0 {
		slog.Info(fmt.Sprintf("uploading %s in %d %s part(s)", b.Digest[7:19], len(b.Parts), format.HumanBytes(b.Parts[0].Size)))
	}
Michael Yang's avatar
Michael Yang committed
114

Michael Yang's avatar
Michael Yang committed
115
	requestURL, err = url.Parse(location)
Michael Yang's avatar
Michael Yang committed
116
117
118
119
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
120
121
122
123
124
	b.nextURL = make(chan *url.URL, 1)
	b.nextURL <- requestURL
	return nil
}

Michael Yang's avatar
Michael Yang committed
125
126
// Run uploads blob parts to the upstream. If the upstream supports redirection, parts will be uploaded
// in parallel as defined by Prepare. Otherwise, parts will be uploaded serially. Run sets b.err on error.
Michael Yang's avatar
Michael Yang committed
127
func (b *blobUpload) Run(ctx context.Context, opts *registryOptions) {
Michael Yang's avatar
Michael Yang committed
128
129
130
	defer blobUploadManager.Delete(b.Digest)
	ctx, b.CancelFunc = context.WithCancel(ctx)

131
132
133
134
135
136
137
138
139
140
141
142
143
	p, err := GetBlobsPath(b.Digest)
	if err != nil {
		b.err = err
		return
	}

	b.file, err = os.Open(p)
	if err != nil {
		b.err = err
		return
	}
	defer b.file.Close()

144
145
	g, inner := errgroup.WithContext(ctx)
	g.SetLimit(numUploadParts)
Michael Yang's avatar
Michael Yang committed
146
147
	for i := range b.Parts {
		part := &b.Parts[i]
Michael Yang's avatar
Michael Yang committed
148
149
150
		select {
		case <-inner.Done():
		case requestURL := <-b.nextURL:
151
			g.Go(func() error {
Michael Yang's avatar
Michael Yang committed
152
				var err error
Michael Yang's avatar
lint  
Michael Yang committed
153
				for try := range maxRetries {
154
					err = b.uploadPart(inner, http.MethodPatch, requestURL, part, opts)
Michael Yang's avatar
Michael Yang committed
155
156
157
158
159
160
					switch {
					case errors.Is(err, context.Canceled):
						return err
					case errors.Is(err, errMaxRetriesExceeded):
						return err
					case err != nil:
Jeffrey Morgan's avatar
Jeffrey Morgan committed
161
						sleep := time.Second * time.Duration(math.Pow(2, float64(try)))
162
						slog.Info(fmt.Sprintf("%s part %d attempt %d failed: %v, retrying in %s", b.Digest[7:19], part.N, try, err, sleep))
Michael Yang's avatar
Michael Yang committed
163
						time.Sleep(sleep)
Michael Yang's avatar
Michael Yang committed
164
165
166
167
						continue
					}

					return nil
Michael Yang's avatar
Michael Yang committed
168
169
				}

Michael Yang's avatar
Michael Yang committed
170
				return fmt.Errorf("%w: %w", errMaxRetriesExceeded, err)
Michael Yang's avatar
Michael Yang committed
171
172
			})
		}
Michael Yang's avatar
Michael Yang committed
173
	}
174

Michael Yang's avatar
Michael Yang committed
175
	if err := g.Wait(); err != nil {
Michael Yang's avatar
Michael Yang committed
176
177
		b.err = err
		return
Michael Yang's avatar
Michael Yang committed
178
179
	}

Michael Yang's avatar
Michael Yang committed
180
181
	requestURL := <-b.nextURL

182
	// calculate md5 checksum and add it to the commit request
183
	md5sum := md5.New()
184
	for _, part := range b.Parts {
185
		md5sum.Write(part.Sum(nil))
186
187
	}

Michael Yang's avatar
Michael Yang committed
188
	values := requestURL.Query()
Michael Yang's avatar
Michael Yang committed
189
	values.Add("digest", b.Digest)
190
	values.Add("etag", fmt.Sprintf("%x-%d", md5sum.Sum(nil), len(b.Parts)))
Michael Yang's avatar
Michael Yang committed
191
192
193
194
195
196
	requestURL.RawQuery = values.Encode()

	headers := make(http.Header)
	headers.Set("Content-Type", "application/octet-stream")
	headers.Set("Content-Length", "0")

Michael Yang's avatar
lint  
Michael Yang committed
197
	for try := range maxRetries {
Michael Yang's avatar
Michael Yang committed
198
199
200
201
202
		var resp *http.Response
		resp, err = makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, nil, opts)
		if errors.Is(err, context.Canceled) {
			break
		} else if err != nil {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
203
			sleep := time.Second * time.Duration(math.Pow(2, float64(try)))
204
			slog.Info(fmt.Sprintf("%s complete upload attempt %d failed: %v, retrying in %s", b.Digest[7:19], try, err, sleep))
Michael Yang's avatar
Michael Yang committed
205
206
207
208
			time.Sleep(sleep)
			continue
		}
		defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
209
		break
Michael Yang's avatar
Michael Yang committed
210
	}
Michael Yang's avatar
Michael Yang committed
211
212
213

	b.err = err
	b.done = true
Michael Yang's avatar
Michael Yang committed
214
}
Michael Yang's avatar
Michael Yang committed
215

Michael Yang's avatar
Michael Yang committed
216
func (b *blobUpload) uploadPart(ctx context.Context, method string, requestURL *url.URL, part *blobUploadPart, opts *registryOptions) error {
Michael Yang's avatar
Michael Yang committed
217
218
	headers := make(http.Header)
	headers.Set("Content-Type", "application/octet-stream")
Michael Yang's avatar
lint  
Michael Yang committed
219
	headers.Set("Content-Length", strconv.FormatInt(part.Size, 10))
220
221

	if method == http.MethodPatch {
222
		headers.Set("X-Redirect-Uploads", "1")
Michael Yang's avatar
Michael Yang committed
223
		headers.Set("Content-Range", fmt.Sprintf("%d-%d", part.Offset, part.Offset+part.Size-1))
224
	}
Michael Yang's avatar
Michael Yang committed
225

226
	sr := io.NewSectionReader(b.file, part.Offset, part.Size)
227
228
229
230

	md5sum := md5.New()
	w := &progressWriter{blobUpload: b}

Michael Yang's avatar
Michael Yang committed
231
	resp, err := makeRequest(ctx, method, requestURL, headers, io.TeeReader(sr, io.MultiWriter(w, md5sum)), opts)
Michael Yang's avatar
Michael Yang committed
232
	if err != nil {
233
		w.Rollback()
Michael Yang's avatar
Michael Yang committed
234
235
236
		return err
	}
	defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
237

Michael Yang's avatar
Michael Yang committed
238
239
240
241
242
243
244
	location := resp.Header.Get("Docker-Upload-Location")
	if location == "" {
		location = resp.Header.Get("Location")
	}

	nextURL, err := url.Parse(location)
	if err != nil {
245
		w.Rollback()
Michael Yang's avatar
Michael Yang committed
246
247
248
249
250
		return err
	}

	switch {
	case resp.StatusCode == http.StatusTemporaryRedirect:
251
		w.Rollback()
Michael Yang's avatar
Michael Yang committed
252
253
254
255
256
257
		b.nextURL <- nextURL

		redirectURL, err := resp.Location()
		if err != nil {
			return err
		}
258

Jeffrey Morgan's avatar
Jeffrey Morgan committed
259
		// retry uploading to the redirect URL
Michael Yang's avatar
lint  
Michael Yang committed
260
		for try := range maxRetries {
Michael Yang's avatar
Michael Yang committed
261
			err = b.uploadPart(ctx, http.MethodPut, redirectURL, part, &registryOptions{})
Michael Yang's avatar
Michael Yang committed
262
263
264
265
266
267
			switch {
			case errors.Is(err, context.Canceled):
				return err
			case errors.Is(err, errMaxRetriesExceeded):
				return err
			case err != nil:
Jeffrey Morgan's avatar
Jeffrey Morgan committed
268
				sleep := time.Second * time.Duration(math.Pow(2, float64(try)))
269
				slog.Info(fmt.Sprintf("%s part %d attempt %d failed: %v, retrying in %s", b.Digest[7:19], part.N, try, err, sleep))
Jeffrey Morgan's avatar
Jeffrey Morgan committed
270
				time.Sleep(sleep)
271
272
273
				continue
			}

Michael Yang's avatar
Michael Yang committed
274
275
276
			return nil
		}

Michael Yang's avatar
Michael Yang committed
277
		return fmt.Errorf("%w: %w", errMaxRetriesExceeded, err)
Michael Yang's avatar
Michael Yang committed
278

Michael Yang's avatar
Michael Yang committed
279
	case resp.StatusCode == http.StatusUnauthorized:
280
		w.Rollback()
Michael Yang's avatar
Michael Yang committed
281
282
		challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
		token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
283
284
285
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
286

Michael Yang's avatar
Michael Yang committed
287
288
289
		opts.Token = token
		fallthrough
	case resp.StatusCode >= http.StatusBadRequest:
290
		w.Rollback()
Michael Yang's avatar
Michael Yang committed
291
292
293
		body, err := io.ReadAll(resp.Body)
		if err != nil {
			return err
Michael Yang's avatar
Michael Yang committed
294
295
		}

296
		return fmt.Errorf("http status %s: %s", resp.Status, body)
Michael Yang's avatar
Michael Yang committed
297
298
	}

Michael Yang's avatar
Michael Yang committed
299
300
301
302
	if method == http.MethodPatch {
		b.nextURL <- nextURL
	}

Michael Yang's avatar
Michael Yang committed
303
	part.Hash = md5sum
Michael Yang's avatar
Michael Yang committed
304
	return nil
Michael Yang's avatar
Michael Yang committed
305
306
}

Michael Yang's avatar
Michael Yang committed
307
308
func (b *blobUpload) acquire() {
	b.references.Add(1)
Michael Yang's avatar
Michael Yang committed
309
310
}

Michael Yang's avatar
Michael Yang committed
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
func (b *blobUpload) release() {
	if b.references.Add(-1) == 0 {
		b.CancelFunc()
	}
}

func (b *blobUpload) Wait(ctx context.Context, fn func(api.ProgressResponse)) error {
	b.acquire()
	defer b.release()

	ticker := time.NewTicker(60 * time.Millisecond)
	for {
		select {
		case <-ticker.C:
		case <-ctx.Done():
			return ctx.Err()
		}

		fn(api.ProgressResponse{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
330
			Status:    fmt.Sprintf("pushing %s", b.Digest[7:19]),
Michael Yang's avatar
Michael Yang committed
331
332
333
			Digest:    b.Digest,
			Total:     b.Total,
			Completed: b.Completed.Load(),
Michael Yang's avatar
Michael Yang committed
334
335
		})

Michael Yang's avatar
Michael Yang committed
336
337
338
		if b.done || b.err != nil {
			return b.err
		}
Michael Yang's avatar
Michael Yang committed
339
	}
Michael Yang's avatar
Michael Yang committed
340
}
Michael Yang's avatar
Michael Yang committed
341

Michael Yang's avatar
Michael Yang committed
342
343
type blobUploadPart struct {
	// N is the part number
344
345
346
347
348
	N      int
	Offset int64
	Size   int64
	hash.Hash
}
349

350
351
type progressWriter struct {
	written int64
Michael Yang's avatar
Michael Yang committed
352
353
354
	*blobUpload
}

355
func (p *progressWriter) Write(b []byte) (n int, err error) {
Michael Yang's avatar
Michael Yang committed
356
357
358
	n = len(b)
	p.written += int64(n)
	p.Completed.Add(int64(n))
Michael Yang's avatar
Michael Yang committed
359
360
	return n, nil
}
Michael Yang's avatar
Michael Yang committed
361

362
363
func (p *progressWriter) Rollback() {
	p.Completed.Add(-p.written)
Michael Yang's avatar
Michael Yang committed
364
365
366
	p.written = 0
}

367
func uploadBlob(ctx context.Context, mp ModelPath, layer Layer, opts *registryOptions, fn func(api.ProgressResponse)) error {
Michael Yang's avatar
Michael Yang committed
368
369
370
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs", layer.Digest)

Michael Yang's avatar
Michael Yang committed
371
372
373
374
	resp, err := makeRequestWithRetry(ctx, http.MethodHead, requestURL, nil, nil, opts)
	switch {
	case errors.Is(err, os.ErrNotExist):
	case err != nil:
Michael Yang's avatar
Michael Yang committed
375
		return err
Michael Yang's avatar
Michael Yang committed
376
377
	default:
		defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
378
		fn(api.ProgressResponse{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
379
			Status:    fmt.Sprintf("pushing %s", layer.Digest[7:19]),
Michael Yang's avatar
Michael Yang committed
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
			Digest:    layer.Digest,
			Total:     layer.Size,
			Completed: layer.Size,
		})

		return nil
	}

	data, ok := blobUploadManager.LoadOrStore(layer.Digest, &blobUpload{Layer: layer})
	upload := data.(*blobUpload)
	if !ok {
		requestURL := mp.BaseURL()
		requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs/uploads/")
		if err := upload.Prepare(ctx, requestURL, opts); err != nil {
			blobUploadManager.Delete(layer.Digest)
			return err
		}

Michael Yang's avatar
Michael Yang committed
398
		//nolint:contextcheck
Michael Yang's avatar
Michael Yang committed
399
400
401
402
403
		go upload.Run(context.Background(), opts)
	}

	return upload.Wait(ctx, fn)
}