upload.go 9.13 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"
15
	"strings"
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
20

	"github.com/jmorganca/ollama/api"
21
	"github.com/jmorganca/ollama/auth"
Michael Yang's avatar
Michael Yang committed
22
23
	"github.com/jmorganca/ollama/format"
	"golang.org/x/sync/errgroup"
Michael Yang's avatar
Michael Yang committed
24
25
)

Michael Yang's avatar
Michael Yang committed
26
27
28
29
30
31
32
33
34
35
36
37
38
39
var blobUploadManager sync.Map

type blobUpload struct {
	*Layer

	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 (
Michael Yang's avatar
Michael Yang committed
48
	numUploadParts          = 64
Michael Yang's avatar
Michael Yang committed
49
50
	minUploadPartSize int64 = 100 * format.MegaByte
	maxUploadPartSize int64 = 1000 * format.MegaByte
51
52
)

53
func (b *blobUpload) Prepare(ctx context.Context, requestURL *url.URL, opts *auth.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
	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
112

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

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

Michael Yang's avatar
Michael Yang committed
123
124
// 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.
125
func (b *blobUpload) Run(ctx context.Context, opts *auth.RegistryOptions) {
Michael Yang's avatar
Michael Yang committed
126
127
128
	defer blobUploadManager.Delete(b.Digest)
	ctx, b.CancelFunc = context.WithCancel(ctx)

129
130
131
132
133
134
135
136
137
138
139
140
141
	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()

Michael Yang's avatar
Michael Yang committed
142
143
144
145
	g, inner := errgroup.WithContext(ctx)
	g.SetLimit(numUploadParts)
	for i := range b.Parts {
		part := &b.Parts[i]
Michael Yang's avatar
Michael Yang committed
146
147
148
149
		select {
		case <-inner.Done():
		case requestURL := <-b.nextURL:
			g.Go(func() error {
Michael Yang's avatar
Michael Yang committed
150
				var err error
Michael Yang's avatar
Michael Yang committed
151
				for try := 0; try < maxRetries; try++ {
152
					err = b.uploadPart(inner, http.MethodPatch, requestURL, part, opts)
Michael Yang's avatar
Michael Yang committed
153
154
155
156
157
158
					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
159
						sleep := time.Second * time.Duration(math.Pow(2, float64(try)))
160
						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
161
						time.Sleep(sleep)
Michael Yang's avatar
Michael Yang committed
162
163
164
165
						continue
					}

					return nil
Michael Yang's avatar
Michael Yang committed
166
167
				}

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

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

Michael Yang's avatar
Michael Yang committed
178
179
	requestURL := <-b.nextURL

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

	md5sum := md5.Sum([]byte(sb.String()))

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, 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
Michael Yang committed
197
	for try := 0; try < maxRetries; try++ {
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

216
func (b *blobUpload) uploadPart(ctx context.Context, method string, requestURL *url.URL, part *blobUploadPart, opts *auth.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
Michael Yang committed
219
	headers.Set("Content-Length", fmt.Sprintf("%d", part.Size))
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}

231
	resp, err := auth.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
Michael Yang committed
260
		for try := 0; try < maxRetries; try++ {
261
			err = b.uploadPart(ctx, http.MethodPut, redirectURL, part, nil)
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()
281
282
283
		authenticate := resp.Header.Get("www-authenticate")
		authRedir := ParseAuthRedirectString(authenticate)
		token, err := auth.GetAuthToken(ctx, authRedir)
Michael Yang's avatar
Michael Yang committed
284
285
286
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
287

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
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
331
			Status:    fmt.Sprintf("pushing %s", b.Digest[7:19]),
Michael Yang's avatar
Michael Yang committed
332
333
334
			Digest:    b.Digest,
			Total:     b.Total,
			Completed: b.Completed.Load(),
Michael Yang's avatar
Michael Yang committed
335
336
		})

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

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
372
373
374
375
	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
376
		return err
Michael Yang's avatar
Michael Yang committed
377
378
	default:
		defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
379
		fn(api.ProgressResponse{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
380
			Status:    fmt.Sprintf("pushing %s", layer.Digest[7:19]),
Michael Yang's avatar
Michael Yang committed
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
			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
399
		// nolint: contextcheck
Michael Yang's avatar
Michael Yang committed
400
401
402
403
404
		go upload.Run(context.Background(), opts)
	}

	return upload.Wait(ctx, fn)
}