images.go 22.9 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
108
109
110
111
112
	v, err := m.Template.Vars()
	if err != nil {
		slog.Warn("model template contains errors", "error", err)
	}
	if slices.Contains(v, "tools") || (builtinParser != nil && builtinParser.HasToolSupport()) {
113
114
115
116
		capabilities = append(capabilities, model.CapabilityTools)
	}

	// Check for insert capability
117
	if slices.Contains(v, "suffix") {
118
119
120
		capabilities = append(capabilities, model.CapabilityInsert)
	}

121
122
123
124
125
	// Check for vision capability in projector-based models
	if len(m.ProjectorPaths) > 0 {
		capabilities = append(capabilities, model.CapabilityVision)
	}

126
127
128
129
130
	// Skip the thinking check if it's already set
	if slices.Contains(capabilities, "thinking") {
		return capabilities
	}

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

139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
	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,
155
		model.CapabilityThinking:   errCapabilityThinking,
156
157
158
159
160
	}

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

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

170
	var err error
171
	if len(errs) > 0 {
172
173
174
175
176
177
178
179
		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
180
181
	}

182
	return err
183
184
}

Michael Yang's avatar
Michael Yang committed
185
func (m *Model) String() string {
186
	var modelfile parser.Modelfile
Michael Yang's avatar
Michael Yang committed
187

188
	modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
189
190
191
		Name: "model",
		Args: m.ModelPath,
	})
192

Michael Yang's avatar
Michael Yang committed
193
	for _, adapter := range m.AdapterPaths {
194
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
195
196
			Name: "adapter",
			Args: adapter,
Michael Yang's avatar
Michael Yang committed
197
		})
198
199
	}

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

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

Michael Yang's avatar
Michael Yang committed
214
	if m.System != "" {
215
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
216
217
			Name: "system",
			Args: m.System,
Michael Yang's avatar
Michael Yang committed
218
		})
219
220
	}

Devon Rifkin's avatar
Devon Rifkin committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
	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,
		})
	}

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

	for _, license := range m.License {
253
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
254
255
256
			Name: "license",
			Args: license,
		})
257
258
259
	}

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

Michael Yang's avatar
Michael Yang committed
266
	return modelfile.String()
267
268
}

269
type ConfigV2 struct {
270
271
272
	ModelFormat   string   `json:"model_format"`
	ModelFamily   string   `json:"model_family"`
	ModelFamilies []string `json:"model_families"`
273
274
	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
275
276
	Renderer      string   `json:"renderer,omitempty"`
	Parser        string   `json:"parser,omitempty"`
277

278
279
280
281
282
283
284
285
286
	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"`

287
	// required by spec
288
289
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
290
	RootFS       RootFS `json:"rootfs"`
291
292
293
294
295
296
297
}

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

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

Michael Yang's avatar
Michael Yang committed
304
	f, err := os.Open(fp)
305
	if err != nil {
Michael Yang's avatar
Michael Yang committed
306
		return nil, "", err
307
	}
Michael Yang's avatar
Michael Yang committed
308
	defer f.Close()
309

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

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

Michael Yang's avatar
Michael Yang committed
317
	return &manifest, hex.EncodeToString(sha256sum.Sum(nil)), nil
318
319
320
}

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

	model := &Model{
328
329
330
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
Michael Yang's avatar
Michael Yang committed
331
		Template:  template.DefaultTemplate,
332
333
	}

334
335
336
337
338
	if manifest.Config.Digest != "" {
		filename, err := GetBlobsPath(manifest.Config.Digest)
		if err != nil {
			return nil, err
		}
339

340
341
342
343
344
		configFile, err := os.Open(filename)
		if err != nil {
			return nil, err
		}
		defer configFile.Close()
345

346
347
348
		if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
			return nil, err
		}
349
350
	}

351
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
352
		filename, err := GetBlobsPath(layer.Digest)
353
354
355
356
		if err != nil {
			return nil, err
		}

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

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

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

394
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
395
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
396
397
				return nil, err
			}
398
399
400
401
402
403
404
405
406
407
		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
408
409
410
411
412
413
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
414
415
416
417
418
419
		}
	}

	return model, nil
}

Michael Yang's avatar
Michael Yang committed
420
func CopyModel(src, dst model.Name) error {
421
422
423
424
425
426
427
	if !dst.IsFullyQualified() {
		return model.Unqualified(dst)
	}
	if !src.IsFullyQualified() {
		return model.Unqualified(src)
	}

428
429
430
431
	if src.Filepath() == dst.Filepath() {
		return nil
	}

Michael Yang's avatar
Michael Yang committed
432
	manifests, err := GetManifestPath()
433
434
435
436
	if err != nil {
		return err
	}

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

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

Michael Yang's avatar
Michael Yang committed
449
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
450
451
452
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
453
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
454

Michael Yang's avatar
Michael Yang committed
455
456
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
457
458
}

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

Michael Yang's avatar
Michael Yang committed
466
	for _, manifest := range manifests {
Michael Yang's avatar
Michael Yang committed
467
468
469
470
471
		for _, layer := range manifest.Layers {
			delete(deleteMap, layer.Digest)
		}

		delete(deleteMap, manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
472
	}
473
474

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

487
488
489
490
	return nil
}

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

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

	for _, blob := range blobs {
		name := blob.Name()
505
		name = strings.ReplaceAll(name, "-", ":")
506
507
508
509
510
511
512
513
514
515
516

		_, 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
517
		}
518
519

		deleteMap[name] = struct{}{}
520
521
	}

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

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

529
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
530
531
532
533

	return nil
}

Michael Yang's avatar
Michael Yang committed
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
563
564
565
566
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
567
func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
568
	mp := ParseModelPath(name)
569
570
	fn(api.ProgressResponse{Status: "retrieving manifest"})

571
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
CYJiang's avatar
CYJiang committed
572
		return errInsecureProtocol
573
574
	}

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

581
	var layers []Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
582
	layers = append(layers, manifest.Layers...)
583
	if manifest.Config.Digest != "" {
584
		layers = append(layers, manifest.Config)
585
	}
586
587

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

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

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

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

611
	fn(api.ProgressResponse{Status: "success"})
612
613
614
615

	return nil
}

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

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

635
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
CYJiang's avatar
CYJiang committed
636
		return errInsecureProtocol
637
	}
638

639
	fn(api.ProgressResponse{Status: "pulling manifest"})
640

641
	manifest, err = pullModelManifest(ctx, mp, regOpts)
642
	if err != nil {
643
		return fmt.Errorf("pull model manifest: %s", err)
644
645
	}

646
	var layers []Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
647
	layers = append(layers, manifest.Layers...)
648
	if manifest.Config.Digest != "" {
649
		layers = append(layers, manifest.Config)
650
	}
651

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

Michael Yang's avatar
Michael Yang committed
668
669
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
670
671
672
		if skipVerify[layer.Digest] {
			continue
		}
Michael Yang's avatar
Michael Yang committed
673
		if err := verifyBlob(layer.Digest); err != nil {
674
675
676
677
678
679
680
681
			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
682
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
683
684
				}
			}
Michael Yang's avatar
Michael Yang committed
685
686
687
688
			return err
		}
	}

689
	fn(api.ProgressResponse{Status: "writing manifest"})
690

691
	manifestJSON, err := json.Marshal(manifest)
692
693
694
695
	if err != nil {
		return err
	}

696
	fp, err := mp.GetManifestPath()
697
698
699
	if err != nil {
		return err
	}
700
701
702
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
703

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

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

717
	fn(api.ProgressResponse{Status: "success"})
718
719
720
721

	return nil
}

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

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

Michael Yang's avatar
Michael Yang committed
733
	var m Manifest
734
735
736
737
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
738
	return &m, err
739
740
741
}

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

Michael Yang's avatar
Michael Yang committed
749
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
750
751
}

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

Michael Yang's avatar
Michael Yang committed
754
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
755
	for range 2 {
Michael Yang's avatar
Michael Yang committed
756
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
757
		if err != nil {
Michael Yang's avatar
Michael Yang committed
758
			if !errors.Is(err, context.Canceled) {
759
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
760
761
			}

Michael Yang's avatar
Michael Yang committed
762
763
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
764
765
766

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
767
768
			resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
769
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
770
771
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
772
773
774
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
775
776
777
778
779
780
781
782
			regOpts.Token = token
			if body != nil {
				_, err = body.Seek(0, io.SeekStart)
				if err != nil {
					return nil, err
				}
			}
		case resp.StatusCode == http.StatusNotFound:
783
			resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
784
785
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
786
			defer resp.Body.Close()
Michael Yang's avatar
Michael Yang committed
787
788
789
790
791
792
793
			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
794
795
796
		}
	}

Michael Yang's avatar
Michael Yang committed
797
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
798
799
}

800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
// 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
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
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)
		}
	}

837
	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
838
839
840
841
842
843
844
845
846
847

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

		req.ContentLength = contentLength
	}

848
	c := &http.Client{
849
		CheckRedirect: regOpts.CheckRedirect,
Michael Yang's avatar
Michael Yang committed
850
	}
851
852
853
854
855
856
	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
857
858
}

Patrick Devine's avatar
Patrick Devine committed
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
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
882
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
883
884
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
885
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
886
887
888
889
890
891
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
894
895
896
897
898
899
900
901
902
903
904
905
906
907
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 {
908
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
909
910
911
912
	}

	return nil
}