images.go 22.8 KB
Newer Older
1
2
3
4
package server

import (
	"bytes"
5
	"context"
6
	"crypto/sha256"
Patrick Devine's avatar
Patrick Devine committed
7
	"encoding/hex"
8
9
10
11
12
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
13
	"log/slog"
14
	"net"
15
	"net/http"
Michael Yang's avatar
Michael Yang committed
16
	"net/url"
17
18
	"os"
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
19
	"runtime"
20
	"slices"
Michael Yang's avatar
Michael Yang committed
21
	"strconv"
22
23
	"strings"

24
	"github.com/ollama/ollama/api"
Michael Yang's avatar
Michael Yang committed
25
	"github.com/ollama/ollama/envconfig"
26
	"github.com/ollama/ollama/fs/gguf"
Devon Rifkin's avatar
Devon Rifkin committed
27
	"github.com/ollama/ollama/model/parsers"
28
	"github.com/ollama/ollama/parser"
Michael Yang's avatar
Michael Yang committed
29
	"github.com/ollama/ollama/template"
30
	"github.com/ollama/ollama/thinking"
Michael Yang's avatar
Michael Yang committed
31
	"github.com/ollama/ollama/types/model"
32
	"github.com/ollama/ollama/version"
33
34
)

35
36
37
38
39
var (
	errCapabilities         = errors.New("does not support")
	errCapabilityCompletion = errors.New("completion")
	errCapabilityTools      = errors.New("tools")
	errCapabilityInsert     = errors.New("insert")
40
41
	errCapabilityVision     = errors.New("vision")
	errCapabilityEmbedding  = errors.New("embedding")
42
	errCapabilityThinking   = errors.New("thinking")
CYJiang's avatar
CYJiang committed
43
	errInsecureProtocol     = errors.New("insecure protocol http")
44
)
Michael Yang's avatar
Michael Yang committed
45

Michael Yang's avatar
Michael Yang committed
46
47
48
49
50
type registryOptions struct {
	Insecure bool
	Username string
	Password string
	Token    string
51
52

	CheckRedirect func(req *http.Request, via []*http.Request) error
Michael Yang's avatar
Michael Yang committed
53
54
}

55
type Model struct {
Michael Yang's avatar
Michael Yang committed
56
	Name           string `json:"name"`
57
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
58
59
	ShortName      string
	ModelPath      string
60
	ParentModel    string
Michael Yang's avatar
Michael Yang committed
61
62
63
64
65
	AdapterPaths   []string
	ProjectorPaths []string
	System         string
	License        []string
	Digest         string
66
	Options        map[string]any
Michael Yang's avatar
Michael Yang committed
67
	Messages       []api.Message
Michael Yang's avatar
Michael Yang committed
68
69

	Template *template.Template
70
71
}

72
73
74
// Capabilities returns the capabilities that the model supports
func (m *Model) Capabilities() []model.Capability {
	capabilities := []model.Capability{}
75

76
	// Check for completion capability
77
78
79
80
81
82
83
84
85
86
87
88
89
90
	if m.ModelPath != "" {
		f, err := gguf.Open(m.ModelPath)
		if err == nil {
			defer f.Close()

			if f.KeyValue("pooling_type").Valid() {
				capabilities = append(capabilities, model.CapabilityEmbedding)
			} else {
				// If no embedding is specified, we assume the model supports completion
				capabilities = append(capabilities, model.CapabilityCompletion)
			}
			if f.KeyValue("vision.block_count").Valid() {
				capabilities = append(capabilities, model.CapabilityVision)
			}
91
		} else {
92
			slog.Error("couldn't open model file", "error", err)
93
		}
94
95
96
	} else if len(m.Config.Capabilities) > 0 {
		for _, c := range m.Config.Capabilities {
			capabilities = append(capabilities, model.Capability(c))
97
98
		}
	} else {
99
		slog.Warn("unknown capabilities for model", "model", m.Name)
100
101
102
103
104
105
	}

	if m.Template == nil {
		return capabilities
	}

Devon Rifkin's avatar
Devon Rifkin committed
106
	builtinParser := parsers.ParserForName(m.Config.Parser)
107
	// Check for tools capability
Devon Rifkin's avatar
Devon Rifkin committed
108
	if slices.Contains(m.Template.Vars(), "tools") || (builtinParser != nil && builtinParser.HasToolSupport()) {
109
110
111
112
113
114
115
116
		capabilities = append(capabilities, model.CapabilityTools)
	}

	// Check for insert capability
	if slices.Contains(m.Template.Vars(), "suffix") {
		capabilities = append(capabilities, model.CapabilityInsert)
	}

117
118
119
120
121
	// Check for vision capability in projector-based models
	if len(m.ProjectorPaths) > 0 {
		capabilities = append(capabilities, model.CapabilityVision)
	}

122
123
124
125
126
	// Skip the thinking check if it's already set
	if slices.Contains(capabilities, "thinking") {
		return capabilities
	}

127
	// Check for thinking capability
128
	openingTag, closingTag := thinking.InferTags(m.Template.Template)
Michael Yang's avatar
Michael Yang committed
129
	hasTags := openingTag != "" && closingTag != ""
Devon Rifkin's avatar
Devon Rifkin committed
130
131
	isGptoss := slices.Contains([]string{"gptoss", "gpt-oss"}, m.Config.ModelFamily)
	if hasTags || isGptoss || (builtinParser != nil && builtinParser.HasThinkingSupport()) {
132
133
134
		capabilities = append(capabilities, model.CapabilityThinking)
	}

135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
	return capabilities
}

// CheckCapabilities checks if the model has the specified capabilities returning an error describing
// any missing or unknown capabilities
func (m *Model) CheckCapabilities(want ...model.Capability) error {
	available := m.Capabilities()
	var errs []error

	// Map capabilities to their corresponding error
	capToErr := map[model.Capability]error{
		model.CapabilityCompletion: errCapabilityCompletion,
		model.CapabilityTools:      errCapabilityTools,
		model.CapabilityInsert:     errCapabilityInsert,
		model.CapabilityVision:     errCapabilityVision,
		model.CapabilityEmbedding:  errCapabilityEmbedding,
151
		model.CapabilityThinking:   errCapabilityThinking,
152
153
154
155
156
	}

	for _, cap := range want {
		err, ok := capToErr[cap]
		if !ok {
Michael Yang's avatar
Michael Yang committed
157
			slog.Error("unknown capability", "capability", cap)
Michael Yang's avatar
Michael Yang committed
158
			return fmt.Errorf("unknown capability: %s", cap)
Michael Yang's avatar
Michael Yang committed
159
		}
160
161
162
163

		if !slices.Contains(available, cap) {
			errs = append(errs, err)
		}
Michael Yang's avatar
Michael Yang committed
164
165
	}

166
	var err error
167
	if len(errs) > 0 {
168
169
170
171
172
173
174
175
		err = fmt.Errorf("%w %w", errCapabilities, errors.Join(errs...))
	}

	if slices.Contains(errs, errCapabilityThinking) {
		if m.Config.ModelFamily == "qwen3" || model.ParseName(m.Name).Model == "deepseek-r1" {
			// append a message to the existing error
			return fmt.Errorf("%w. Pull the model again to get the latest version with full thinking support", err)
		}
Michael Yang's avatar
Michael Yang committed
176
177
	}

178
	return err
179
180
}

Michael Yang's avatar
Michael Yang committed
181
func (m *Model) String() string {
182
	var modelfile parser.Modelfile
Michael Yang's avatar
Michael Yang committed
183

184
	modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
185
186
187
		Name: "model",
		Args: m.ModelPath,
	})
188

Michael Yang's avatar
Michael Yang committed
189
	for _, adapter := range m.AdapterPaths {
190
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
191
192
			Name: "adapter",
			Args: adapter,
Michael Yang's avatar
Michael Yang committed
193
		})
194
195
	}

Michael Yang's avatar
Michael Yang committed
196
	for _, projector := range m.ProjectorPaths {
197
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
198
199
			Name: "model",
			Args: projector,
Michael Yang's avatar
Michael Yang committed
200
		})
201
202
	}

Michael Yang's avatar
Michael Yang committed
203
	if m.Template != nil {
204
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
205
			Name: "template",
Michael Yang's avatar
Michael Yang committed
206
			Args: m.Template.String(),
Michael Yang's avatar
Michael Yang committed
207
		})
208
209
	}

Michael Yang's avatar
Michael Yang committed
210
	if m.System != "" {
211
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
212
213
			Name: "system",
			Args: m.System,
Michael Yang's avatar
Michael Yang committed
214
		})
215
216
	}

Devon Rifkin's avatar
Devon Rifkin committed
217
218
219
220
221
222
223
224
225
226
227
228
229
230
	if m.Config.Renderer != "" {
		modelfile.Commands = append(modelfile.Commands, parser.Command{
			Name: "renderer",
			Args: m.Config.Renderer,
		})
	}

	if m.Config.Parser != "" {
		modelfile.Commands = append(modelfile.Commands, parser.Command{
			Name: "parser",
			Args: m.Config.Parser,
		})
	}

231
232
233
234
	for k, v := range m.Options {
		switch v := v.(type) {
		case []any:
			for _, s := range v {
235
				modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
236
237
238
					Name: k,
					Args: fmt.Sprintf("%v", s),
				})
239
240
			}
		default:
241
			modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
242
243
244
				Name: k,
				Args: fmt.Sprintf("%v", v),
			})
245
246
247
248
		}
	}

	for _, license := range m.License {
249
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
250
251
252
			Name: "license",
			Args: license,
		})
253
254
255
	}

	for _, msg := range m.Messages {
256
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
257
			Name: "message",
Michael Yang's avatar
Michael Yang committed
258
			Args: fmt.Sprintf("%s: %s", msg.Role, msg.Content),
Michael Yang's avatar
Michael Yang committed
259
		})
260
261
	}

Michael Yang's avatar
Michael Yang committed
262
	return modelfile.String()
263
264
}

265
type ConfigV2 struct {
266
267
268
	ModelFormat   string   `json:"model_format"`
	ModelFamily   string   `json:"model_family"`
	ModelFamilies []string `json:"model_families"`
269
270
	ModelType     string   `json:"model_type"` // shown as Parameter Size
	FileType      string   `json:"file_type"`  // shown as Quantization Level
Devon Rifkin's avatar
Devon Rifkin committed
271
272
	Renderer      string   `json:"renderer,omitempty"`
	Parser        string   `json:"parser,omitempty"`
273

274
275
276
277
278
279
280
281
282
	RemoteHost  string `json:"remote_host,omitempty"`
	RemoteModel string `json:"remote_model,omitempty"`

	// used for remotes
	Capabilities []string `json:"capabilities,omitempty"`
	ContextLen   int      `json:"context_length,omitempty"`
	EmbedLen     int      `json:"embedding_length,omitempty"`
	BaseName     string   `json:"base_name,omitempty"`

283
	// required by spec
284
285
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
286
	RootFS       RootFS `json:"rootfs"`
287
288
289
290
291
292
293
}

type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
294
func GetManifest(mp ModelPath) (*Manifest, string, error) {
295
	fp, err := mp.GetManifestPath()
296
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
297
		return nil, "", err
298
	}
299

Michael Yang's avatar
Michael Yang committed
300
	f, err := os.Open(fp)
301
	if err != nil {
Michael Yang's avatar
Michael Yang committed
302
		return nil, "", err
303
	}
Michael Yang's avatar
Michael Yang committed
304
	defer f.Close()
305

Michael Yang's avatar
Michael Yang committed
306
	sha256sum := sha256.New()
Patrick Devine's avatar
Patrick Devine committed
307

Michael Yang's avatar
Michael Yang committed
308
309
	var manifest Manifest
	if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
310
		return nil, "", err
311
312
	}

Michael Yang's avatar
Michael Yang committed
313
	return &manifest, hex.EncodeToString(sha256sum.Sum(nil)), nil
314
315
316
}

func GetModel(name string) (*Model, error) {
317
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
318
	manifest, digest, err := GetManifest(mp)
319
320
321
322
323
	if err != nil {
		return nil, err
	}

	model := &Model{
324
325
326
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
Michael Yang's avatar
Michael Yang committed
327
		Template:  template.DefaultTemplate,
328
329
	}

330
331
332
333
334
	if manifest.Config.Digest != "" {
		filename, err := GetBlobsPath(manifest.Config.Digest)
		if err != nil {
			return nil, err
		}
335

336
337
338
339
340
		configFile, err := os.Open(filename)
		if err != nil {
			return nil, err
		}
		defer configFile.Close()
341

342
343
344
		if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
			return nil, err
		}
345
346
	}

347
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
348
		filename, err := GetBlobsPath(layer.Digest)
349
350
351
352
		if err != nil {
			return nil, err
		}

353
354
355
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
356
			model.ParentModel = layer.From
357
		case "application/vnd.ollama.image.embed":
358
359
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
360
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
361
362
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
363
364
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
Michael Yang's avatar
Michael Yang committed
365
366
		case "application/vnd.ollama.image.prompt",
			"application/vnd.ollama.image.template":
367
368
369
370
371
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
372
			model.Template, err = template.Parse(string(bts))
373
374
375
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
376
		case "application/vnd.ollama.image.system":
377
378
379
380
381
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
382
			model.System = string(bts)
383
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
384
385
386
387
388
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
389

390
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
391
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
392
393
				return nil, err
			}
394
395
396
397
398
399
400
401
402
403
		case "application/vnd.ollama.image.messages":
			msgs, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer msgs.Close()

			if err = json.NewDecoder(msgs).Decode(&model.Messages); err != nil {
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
404
405
406
407
408
409
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
410
411
412
413
414
415
		}
	}

	return model, nil
}

Michael Yang's avatar
Michael Yang committed
416
func CopyModel(src, dst model.Name) error {
417
418
419
420
421
422
423
	if !dst.IsFullyQualified() {
		return model.Unqualified(dst)
	}
	if !src.IsFullyQualified() {
		return model.Unqualified(src)
	}

424
425
426
427
	if src.Filepath() == dst.Filepath() {
		return nil
	}

Michael Yang's avatar
Michael Yang committed
428
	manifests, err := GetManifestPath()
429
430
431
432
	if err != nil {
		return err
	}

433
	dstpath := filepath.Join(manifests, dst.Filepath())
Michael Yang's avatar
Michael Yang committed
434
	if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
435
436
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
437

438
	srcpath := filepath.Join(manifests, src.Filepath())
Michael Yang's avatar
Michael Yang committed
439
	srcfile, err := os.Open(srcpath)
Patrick Devine's avatar
Patrick Devine committed
440
441
442
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
443
	defer srcfile.Close()
Patrick Devine's avatar
Patrick Devine committed
444

Michael Yang's avatar
Michael Yang committed
445
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
446
447
448
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
449
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
450

Michael Yang's avatar
Michael Yang committed
451
452
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
453
454
}

Michael Yang's avatar
Michael Yang committed
455
func deleteUnusedLayers(deleteMap map[string]struct{}) error {
456
457
	// Ignore corrupt manifests to avoid blocking deletion of layers that are freshly orphaned
	manifests, err := Manifests(true)
458
459
460
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
461

Michael Yang's avatar
Michael Yang committed
462
	for _, manifest := range manifests {
Michael Yang's avatar
Michael Yang committed
463
464
465
466
467
		for _, layer := range manifest.Layers {
			delete(deleteMap, layer.Digest)
		}

		delete(deleteMap, manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
468
	}
469
470

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
471
472
473
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
474
			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
Michael Yang's avatar
Michael Yang committed
475
476
			continue
		}
477
478
479
		if err := os.Remove(fp); err != nil {
			slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
			continue
480
481
482
		}
	}

483
484
485
486
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
487
	deleteMap := make(map[string]struct{})
488
489
490
491
492
493
494
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
495
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
496
497
498
499
500
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
501
		name = strings.ReplaceAll(name, "-", ":")
502
503
504
505
506
507
508
509
510
511
512

		_, err := GetBlobsPath(name)
		if err != nil {
			if errors.Is(err, ErrInvalidDigestFormat) {
				// remove invalid blobs (e.g. partial downloads)
				if err := os.Remove(filepath.Join(p, blob.Name())); err != nil {
					slog.Error("couldn't remove blob", "blob", blob.Name(), "error", err)
				}
			}

			continue
Michael Yang's avatar
Michael Yang committed
513
		}
514
515

		deleteMap[name] = struct{}{}
516
517
	}

518
	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
519

Michael Yang's avatar
Michael Yang committed
520
	if err := deleteUnusedLayers(deleteMap); err != nil {
521
		slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
522
		return nil
523
524
	}

525
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
526
527
528
529

	return nil
}

Michael Yang's avatar
Michael Yang committed
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
func PruneDirectory(path string) error {
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}

	if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
		entries, err := os.ReadDir(path)
		if err != nil {
			return err
		}

		for _, entry := range entries {
			if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
				return err
			}
		}

		entries, err = os.ReadDir(path)
		if err != nil {
			return err
		}

		if len(entries) > 0 {
			return nil
		}

		return os.Remove(path)
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
563
func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
564
	mp := ParseModelPath(name)
565
566
	fn(api.ProgressResponse{Status: "retrieving manifest"})

567
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
CYJiang's avatar
CYJiang committed
568
		return errInsecureProtocol
569
570
	}

Patrick Devine's avatar
Patrick Devine committed
571
	manifest, _, err := GetManifest(mp)
572
	if err != nil {
573
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
574
575
576
		return err
	}

577
	var layers []Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
578
	layers = append(layers, manifest.Layers...)
579
	if manifest.Config.Digest != "" {
580
		layers = append(layers, manifest.Config)
581
	}
582
583

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
584
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
585
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
586
587
			return err
		}
588
589
	}

590
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
591
592
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
593
594
595
596
597
598

	manifestJSON, err := json.Marshal(manifest)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
599
600
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
601
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
602
603
604
605
606
	if err != nil {
		return err
	}
	defer resp.Body.Close()

607
	fn(api.ProgressResponse{Status: "success"})
608
609
610
611

	return nil
}

Michael Yang's avatar
Michael Yang committed
612
func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
613
614
	mp := ParseModelPath(name)

615
	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
616
	deleteMap := make(map[string]struct{})
Michael Yang's avatar
Michael Yang committed
617
618
619
	manifest, _, err := GetManifest(mp)
	if errors.Is(err, os.ErrNotExist) {
		// noop
620
621
	} else if err != nil {
		slog.Warn("pulling model with bad existing manifest", "name", name, "error", err)
Michael Yang's avatar
Michael Yang committed
622
623
624
	} else {
		for _, l := range manifest.Layers {
			deleteMap[l.Digest] = struct{}{}
625
		}
Michael Yang's avatar
Michael Yang committed
626
627
		if manifest.Config.Digest != "" {
			deleteMap[manifest.Config.Digest] = struct{}{}
628
629
630
		}
	}

631
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
CYJiang's avatar
CYJiang committed
632
		return errInsecureProtocol
633
	}
634

635
	fn(api.ProgressResponse{Status: "pulling manifest"})
636

637
	manifest, err = pullModelManifest(ctx, mp, regOpts)
638
	if err != nil {
639
		return fmt.Errorf("pull model manifest: %s", err)
640
641
	}

642
	var layers []Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
643
	layers = append(layers, manifest.Layers...)
644
	if manifest.Config.Digest != "" {
645
		layers = append(layers, manifest.Config)
646
	}
647

648
	skipVerify := make(map[string]bool)
649
	for _, layer := range layers {
650
651
652
653
654
655
656
		cacheHit, err := downloadBlob(ctx, downloadOpts{
			mp:      mp,
			digest:  layer.Digest,
			regOpts: regOpts,
			fn:      fn,
		})
		if err != nil {
657
658
			return err
		}
659
		skipVerify[layer.Digest] = cacheHit
660
		delete(deleteMap, layer.Digest)
661
	}
662
	delete(deleteMap, manifest.Config.Digest)
663

Michael Yang's avatar
Michael Yang committed
664
665
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
666
667
668
		if skipVerify[layer.Digest] {
			continue
		}
Michael Yang's avatar
Michael Yang committed
669
		if err := verifyBlob(layer.Digest); err != nil {
670
671
672
673
674
675
676
677
			if errors.Is(err, errDigestMismatch) {
				// something went wrong, delete the blob
				fp, err := GetBlobsPath(layer.Digest)
				if err != nil {
					return err
				}
				if err := os.Remove(fp); err != nil {
					// log this, but return the original error
678
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
679
680
				}
			}
Michael Yang's avatar
Michael Yang committed
681
682
683
684
			return err
		}
	}

685
	fn(api.ProgressResponse{Status: "writing manifest"})
686

687
	manifestJSON, err := json.Marshal(manifest)
688
689
690
691
	if err != nil {
		return err
	}

692
	fp, err := mp.GetManifestPath()
693
694
695
	if err != nil {
		return err
	}
696
697
698
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
699

Bruce MacDonald's avatar
Bruce MacDonald committed
700
	err = os.WriteFile(fp, manifestJSON, 0o644)
701
	if err != nil {
702
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
703
704
705
		return err
	}

Michael Yang's avatar
Michael Yang committed
706
707
	if !envconfig.NoPrune() && len(deleteMap) > 0 {
		fn(api.ProgressResponse{Status: "removing unused layers"})
Michael Yang's avatar
Michael Yang committed
708
		if err := deleteUnusedLayers(deleteMap); err != nil {
709
			fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
710
711
712
		}
	}

713
	fn(api.ProgressResponse{Status: "success"})
714
715
716
717

	return nil
}

Michael Yang's avatar
Michael Yang committed
718
func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *registryOptions) (*Manifest, error) {
Michael Yang's avatar
Michael Yang committed
719
	requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
720

Michael Yang's avatar
Michael Yang committed
721
722
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
723
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
724
725
726
727
728
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
729
	var m Manifest
730
731
732
733
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
734
	return &m, err
735
736
737
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
738
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
739
740
741
742
743
744
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
745
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
746
747
}

Michael Yang's avatar
lint  
Michael Yang committed
748
var errUnauthorized = errors.New("unauthorized: access denied")
749

Michael Yang's avatar
Michael Yang committed
750
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
Michael Yang's avatar
lint  
Michael Yang committed
751
	for range 2 {
Michael Yang's avatar
Michael Yang committed
752
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
753
		if err != nil {
Michael Yang's avatar
Michael Yang committed
754
			if !errors.Is(err, context.Canceled) {
755
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
756
757
			}

Michael Yang's avatar
Michael Yang committed
758
759
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
760
761
762

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
763
764
			resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
765
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
766
767
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
768
769
770
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
771
772
773
774
775
776
777
778
			regOpts.Token = token
			if body != nil {
				_, err = body.Seek(0, io.SeekStart)
				if err != nil {
					return nil, err
				}
			}
		case resp.StatusCode == http.StatusNotFound:
779
			resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
780
781
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
782
			defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
783
784
785
786
787
788
789
			responseBody, err := io.ReadAll(resp.Body)
			if err != nil {
				return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
			}
			return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
		default:
			return resp, nil
Michael Yang's avatar
Michael Yang committed
790
791
792
		}
	}

Michael Yang's avatar
Michael Yang committed
793
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
794
795
}

796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
// testMakeRequestDialContext specifies the dial function for the http client in
// makeRequest. It can be used to resolve hosts in model names to local
// addresses for testing. For example, the model name ("example.com/my/model")
// can be directed to push/pull from "127.0.0.1:1234".
//
// This is not safe to set across goroutines. It should be set in
// the main test goroutine, and not by tests marked to run in parallel with
// t.Parallel().
//
// It should be cleared after use, otherwise it will affect other tests.
//
// Ideally we would have some set this up the stack, but the code is not
// structured in a way that makes this easy, so this will have to do for now.
var testMakeRequestDialContext func(ctx context.Context, network, addr string) (net.Conn, error)

Michael Yang's avatar
Michael Yang committed
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *registryOptions) (*http.Response, error) {
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
		requestURL.Scheme = "http"
	}

	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
	if err != nil {
		return nil, err
	}

	if headers != nil {
		req.Header = headers
	}

	if regOpts != nil {
		if regOpts.Token != "" {
			req.Header.Set("Authorization", "Bearer "+regOpts.Token)
		} else if regOpts.Username != "" && regOpts.Password != "" {
			req.SetBasicAuth(regOpts.Username, regOpts.Password)
		}
	}

833
	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
Michael Yang's avatar
Michael Yang committed
834
835
836
837
838
839
840
841
842
843

	if s := req.Header.Get("Content-Length"); s != "" {
		contentLength, err := strconv.ParseInt(s, 10, 64)
		if err != nil {
			return nil, err
		}

		req.ContentLength = contentLength
	}

844
	c := &http.Client{
845
		CheckRedirect: regOpts.CheckRedirect,
Michael Yang's avatar
Michael Yang committed
846
	}
847
848
849
850
851
852
	if testMakeRequestDialContext != nil {
		tr := http.DefaultTransport.(*http.Transport).Clone()
		tr.DialContext = testMakeRequestDialContext
		c.Transport = tr
	}
	return c.Do(req)
Michael Yang's avatar
Michael Yang committed
853
854
}

Patrick Devine's avatar
Patrick Devine committed
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
func getValue(header, key string) string {
	startIdx := strings.Index(header, key+"=")
	if startIdx == -1 {
		return ""
	}

	// Move the index to the starting quote after the key.
	startIdx += len(key) + 2
	endIdx := startIdx

	for endIdx < len(header) {
		if header[endIdx] == '"' {
			if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
				endIdx++
				continue
			}
			break
		}
		endIdx++
	}
	return header[startIdx:endIdx]
}

Michael Yang's avatar
Michael Yang committed
878
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
879
880
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
881
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
882
883
884
885
886
887
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

888
var errDigestMismatch = errors.New("digest mismatch, file must be downloaded again")
889

Michael Yang's avatar
Michael Yang committed
890
891
892
893
894
895
896
897
898
899
900
901
902
903
func verifyBlob(digest string) error {
	fp, err := GetBlobsPath(digest)
	if err != nil {
		return err
	}

	f, err := os.Open(fp)
	if err != nil {
		return err
	}
	defer f.Close()

	fileDigest, _ := GetSHA256Digest(f)
	if digest != fileDigest {
904
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
905
906
907
908
	}

	return nil
}