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

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

25
	"github.com/ollama/ollama/api"
26
	"github.com/ollama/ollama/auth"
Michael Yang's avatar
Michael Yang committed
27
	"github.com/ollama/ollama/envconfig"
Michael Yang's avatar
Michael Yang committed
28
	"github.com/ollama/ollama/format"
29
	"github.com/ollama/ollama/llm"
30
	"github.com/ollama/ollama/parser"
Michael Yang's avatar
Michael Yang committed
31
	"github.com/ollama/ollama/template"
32
	"github.com/ollama/ollama/types/errtypes"
Michael Yang's avatar
Michael Yang committed
33
	"github.com/ollama/ollama/types/model"
34
	"github.com/ollama/ollama/version"
35
36
)

Michael Yang's avatar
Michael Yang committed
37
38
39
40
type Capability string

const CapabilityCompletion = Capability("completion")

Michael Yang's avatar
Michael Yang committed
41
42
43
44
45
46
47
type registryOptions struct {
	Insecure bool
	Username string
	Password string
	Token    string
}

48
type Model struct {
Michael Yang's avatar
Michael Yang committed
49
	Name           string `json:"name"`
50
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
51
52
	ShortName      string
	ModelPath      string
53
	ParentModel    string
Michael Yang's avatar
Michael Yang committed
54
55
56
57
58
59
	AdapterPaths   []string
	ProjectorPaths []string
	System         string
	License        []string
	Digest         string
	Options        map[string]interface{}
60
	Messages       []Message
Michael Yang's avatar
Michael Yang committed
61
62

	Template *template.Template
63
64
}

Michael Yang's avatar
Michael Yang committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
func (m *Model) Has(caps ...Capability) bool {
	for _, cap := range caps {
		switch cap {
		case CapabilityCompletion:
			if slices.Contains(m.Config.ModelFamilies, "bert") || slices.Contains(m.Config.ModelFamilies, "nomic-bert") {
				return false
			}
		default:
			slog.Error("unknown capability", "capability", cap)
			return false
		}
	}

	return true
79
80
}

Michael Yang's avatar
Michael Yang committed
81
func (m *Model) String() string {
82
	var modelfile parser.File
Michael Yang's avatar
Michael Yang committed
83

84
	modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
85
86
87
		Name: "model",
		Args: m.ModelPath,
	})
88

Michael Yang's avatar
Michael Yang committed
89
	for _, adapter := range m.AdapterPaths {
90
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
91
92
			Name: "adapter",
			Args: adapter,
Michael Yang's avatar
Michael Yang committed
93
		})
94
95
	}

Michael Yang's avatar
Michael Yang committed
96
	for _, projector := range m.ProjectorPaths {
97
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
98
99
			Name: "model",
			Args: projector,
Michael Yang's avatar
Michael Yang committed
100
		})
101
102
	}

Michael Yang's avatar
Michael Yang committed
103
	if m.Template != nil {
104
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
105
			Name: "template",
Michael Yang's avatar
Michael Yang committed
106
			Args: m.Template.String(),
Michael Yang's avatar
Michael Yang committed
107
		})
108
109
	}

Michael Yang's avatar
Michael Yang committed
110
	if m.System != "" {
111
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
112
113
			Name: "system",
			Args: m.System,
Michael Yang's avatar
Michael Yang committed
114
		})
115
116
117
118
119
120
	}

	for k, v := range m.Options {
		switch v := v.(type) {
		case []any:
			for _, s := range v {
121
				modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
122
123
124
					Name: k,
					Args: fmt.Sprintf("%v", s),
				})
125
126
			}
		default:
127
			modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
128
129
130
				Name: k,
				Args: fmt.Sprintf("%v", v),
			})
131
132
133
134
		}
	}

	for _, license := range m.License {
135
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
136
137
138
			Name: "license",
			Args: license,
		})
139
140
141
	}

	for _, msg := range m.Messages {
142
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
143
144
145
			Name: "message",
			Args: fmt.Sprintf("%s %s", msg.Role, msg.Content),
		})
146
147
	}

Michael Yang's avatar
Michael Yang committed
148
	return modelfile.String()
149
150
}

151
152
153
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
154
155
156
}

type ConfigV2 struct {
157
158
159
160
161
162
	ModelFormat   string   `json:"model_format"`
	ModelFamily   string   `json:"model_family"`
	ModelFamilies []string `json:"model_families"`
	ModelType     string   `json:"model_type"`
	FileType      string   `json:"file_type"`

163
	// required by spec
164
165
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
166
	RootFS       RootFS `json:"rootfs"`
167
168
169
170
171
172
173
}

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

Michael Yang's avatar
Michael Yang committed
174
func GetManifest(mp ModelPath) (*Manifest, string, error) {
175
	fp, err := mp.GetManifestPath()
176
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
177
		return nil, "", err
178
	}
179

180
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
181
		return nil, "", err
182
183
	}

Michael Yang's avatar
Michael Yang committed
184
	var manifest *Manifest
185

186
	bts, err := os.ReadFile(fp)
187
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
188
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
189
190
	}

Patrick Devine's avatar
Patrick Devine committed
191
192
193
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

194
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
195
		return nil, "", err
196
197
	}

Patrick Devine's avatar
Patrick Devine committed
198
	return manifest, shaStr, nil
199
200
201
}

func GetModel(name string) (*Model, error) {
202
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
203
	manifest, digest, err := GetManifest(mp)
204
205
206
207
208
	if err != nil {
		return nil, err
	}

	model := &Model{
209
210
211
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
Michael Yang's avatar
Michael Yang committed
212
		Template:  template.DefaultTemplate,
213
214
	}

215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
	filename, err := GetBlobsPath(manifest.Config.Digest)
	if err != nil {
		return nil, err
	}

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

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

230
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
231
		filename, err := GetBlobsPath(layer.Digest)
232
233
234
235
		if err != nil {
			return nil, err
		}

236
237
238
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
239
			model.ParentModel = layer.From
240
		case "application/vnd.ollama.image.embed":
241
242
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
243
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
244
245
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
246
247
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
Michael Yang's avatar
Michael Yang committed
248
249
		case "application/vnd.ollama.image.prompt",
			"application/vnd.ollama.image.template":
250
251
252
253
254
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
255
			model.Template, err = template.Parse(string(bts))
256
257
258
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
259
		case "application/vnd.ollama.image.system":
260
261
262
263
264
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
265
			model.System = string(bts)
266
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
267
268
269
270
271
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
272

273
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
274
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
275
276
				return nil, err
			}
277
278
279
280
281
282
283
284
285
286
		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
287
288
289
290
291
292
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
293
294
295
296
297
298
		}
	}

	return model, nil
}

Michael Yang's avatar
Michael Yang committed
299
func realpath(rel, from string) string {
300
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
301
	if err != nil {
302
		return from
303
304
	}

Michael Yang's avatar
Michael Yang committed
305
	home, err := os.UserHomeDir()
306
	if err != nil {
Michael Yang's avatar
Michael Yang committed
307
		return abspath
308
309
	}

310
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
311
		return home
312
313
314
315
	} else if strings.HasPrefix(from, "~/") {
		return filepath.Join(home, from[2:])
	}

Michael Yang's avatar
Michael Yang committed
316
	if _, err := os.Stat(filepath.Join(rel, from)); err == nil {
317
		// this is a file relative to the Modelfile
Michael Yang's avatar
Michael Yang committed
318
		return filepath.Join(rel, from)
319
320
	}

Michael Yang's avatar
Michael Yang committed
321
322
323
	return abspath
}

324
func CreateModel(ctx context.Context, name model.Name, modelFileDir, quantization string, modelfile *parser.File, fn func(resp api.ProgressResponse)) (err error) {
325
326
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
327
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
328
329
330
		RootFS: RootFS{
			Type: "layers",
		},
331
332
	}

Michael Yang's avatar
Michael Yang committed
333
334
	var messages []*api.Message
	parameters := make(map[string]any)
Michael Yang's avatar
Michael Yang committed
335

Michael Yang's avatar
Michael Yang committed
336
	var layers []*Layer
Michael Yang's avatar
Michael Yang committed
337
	for _, c := range modelfile.Commands {
Michael Yang's avatar
Michael Yang committed
338
339
		mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

340
		switch c.Name {
Michael Yang's avatar
Michael Yang committed
341
		case "model", "adapter":
342
			var baseLayers []*layerGGML
Michael Yang's avatar
rebase  
Michael Yang committed
343
			if name := model.ParseName(c.Args); name.IsValid() {
Michael Yang's avatar
Michael Yang committed
344
				baseLayers, err = parseFromModel(ctx, name, fn)
Michael Yang's avatar
Michael Yang committed
345
346
347
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
348
			} else if strings.HasPrefix(c.Args, "@") {
349
				digest := strings.TrimPrefix(c.Args, "@")
Michael Yang's avatar
Michael Yang committed
350
351
				if ib, ok := intermediateBlobs[digest]; ok {
					p, err := GetBlobsPath(ib)
352
353
354
355
356
357
358
359
360
					if err != nil {
						return err
					}

					if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
						// pass
					} else if err != nil {
						return err
					} else {
Michael Yang's avatar
Michael Yang committed
361
362
						fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib)})
						digest = ib
363
364
365
366
					}
				}

				blobpath, err := GetBlobsPath(digest)
Michael Yang's avatar
Michael Yang committed
367
				if err != nil {
Michael Yang's avatar
Michael Yang committed
368
					return err
369
				}
370

Michael Yang's avatar
Michael Yang committed
371
				blob, err := os.Open(blobpath)
Michael Yang's avatar
Michael Yang committed
372
373
374
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
375
				defer blob.Close()
Michael Yang's avatar
Michael Yang committed
376

377
				baseLayers, err = parseFromFile(ctx, blob, digest, fn)
Michael Yang's avatar
Michael Yang committed
378
379
380
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
381
382
			} else if file, err := os.Open(realpath(modelFileDir, c.Args)); err == nil {
				defer file.Close()
Michael Yang's avatar
Michael Yang committed
383

384
				baseLayers, err = parseFromFile(ctx, file, "", fn)
Michael Yang's avatar
Michael Yang committed
385
				if err != nil {
Michael Yang's avatar
Michael Yang committed
386
387
					return err
				}
Michael Yang's avatar
Michael Yang committed
388
389
390
			} else {
				return fmt.Errorf("invalid model reference: %s", c.Args)
			}
Michael Yang's avatar
Michael Yang committed
391

Michael Yang's avatar
Michael Yang committed
392
			for _, baseLayer := range baseLayers {
Michael Yang's avatar
Michael Yang committed
393
394
395
396
				if quantization != "" &&
					baseLayer.MediaType == "application/vnd.ollama.image.model" &&
					baseLayer.GGML != nil &&
					baseLayer.GGML.Name() == "gguf" {
Michael Yang's avatar
Michael Yang committed
397
					want, err := llm.ParseFileType(quantization)
398
					if err != nil {
Michael Yang's avatar
Michael Yang committed
399
						return err
400
					}
Michael Yang's avatar
Michael Yang committed
401

Michael Yang's avatar
Michael Yang committed
402
403
					ft := baseLayer.GGML.KV().FileType()
					if !slices.Contains([]string{"F16", "F32"}, ft.String()) {
Michael Yang's avatar
Michael Yang committed
404
						return errors.New("quantization is only supported for F16 and F32 models")
Michael Yang's avatar
Michael Yang committed
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
					} else if want != ft {
						fn(api.ProgressResponse{Status: fmt.Sprintf("quantizing %s model to %s", ft, quantization)})

						blob, err := GetBlobsPath(baseLayer.Digest)
						if err != nil {
							return err
						}

						temp, err := os.CreateTemp(filepath.Dir(blob), quantization)
						if err != nil {
							return err
						}
						defer temp.Close()
						defer os.Remove(temp.Name())

						if err := llm.Quantize(blob, temp.Name(), want); err != nil {
							return err
						}

Michael Yang's avatar
Michael Yang committed
424
						layer, err := NewLayer(temp, baseLayer.MediaType)
Michael Yang's avatar
Michael Yang committed
425
426
427
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
428

Michael Yang's avatar
Michael Yang committed
429
430
431
432
						if _, err := temp.Seek(0, io.SeekStart); err != nil {
							return err
						}

433
						ggml, _, err := llm.DecodeGGML(temp, 0)
Michael Yang's avatar
Michael Yang committed
434
435
						if err != nil {
							return err
Michael Yang's avatar
Michael Yang committed
436
437
						}

Michael Yang's avatar
Michael Yang committed
438
439
						baseLayer.Layer = layer
						baseLayer.GGML = ggml
Michael Yang's avatar
Michael Yang committed
440
					}
441
				}
442

Michael Yang's avatar
Michael Yang committed
443
444
445
446
				if baseLayer.GGML != nil {
					config.ModelFormat = cmp.Or(config.ModelFormat, baseLayer.GGML.Name())
					config.ModelFamily = cmp.Or(config.ModelFamily, baseLayer.GGML.KV().Architecture())
					config.ModelType = cmp.Or(config.ModelType, format.HumanNumber(baseLayer.GGML.KV().ParameterCount()))
Michael Yang's avatar
Michael Yang committed
447
					config.FileType = cmp.Or(config.FileType, baseLayer.GGML.KV().FileType().String())
Michael Yang's avatar
Michael Yang committed
448
					config.ModelFamilies = append(config.ModelFamilies, baseLayer.GGML.KV().Architecture())
449
450
				}

Michael Yang's avatar
Michael Yang committed
451
				layers = append(layers, baseLayer.Layer)
452
			}
Michael Yang's avatar
Michael Yang committed
453
454
455
456
		case "license", "template", "system":
			if c.Name != "license" {
				// replace
				layers = slices.DeleteFunc(layers, func(layer *Layer) bool {
457
458
459
460
461
462
463
464
465
					if layer.MediaType != mediatype {
						return false
					}

					if err := layer.Remove(); err != nil {
						return false
					}

					return true
Michael Yang's avatar
Michael Yang committed
466
				})
Michael Yang's avatar
Michael Yang committed
467
468
			}

469
470
471
472
473
474
			blob := strings.NewReader(c.Args)
			layer, err := NewLayer(blob, mediatype)
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
475
476
477
478
479
480
481
482
483
484
			layers = append(layers, layer)
		case "message":
			role, content, ok := strings.Cut(c.Args, ": ")
			if !ok {
				return fmt.Errorf("invalid message: %s", c.Args)
			}

			messages = append(messages, &api.Message{Role: role, Content: content})
		default:
			ps, err := api.FormatParams(map[string][]string{c.Name: {c.Args}})
485
			if err != nil {
Michael Yang's avatar
Michael Yang committed
486
				return err
487
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
488

Michael Yang's avatar
Michael Yang committed
489
490
491
492
493
494
495
496
497
498
499
			for k, v := range ps {
				if ks, ok := parameters[k].([]string); ok {
					parameters[k] = append(ks, v.([]string)...)
				} else if vs, ok := v.([]string); ok {
					parameters[k] = vs
				} else {
					parameters[k] = v
				}
			}
		}
	}
Michael Yang's avatar
Michael Yang committed
500

Michael Yang's avatar
Michael Yang committed
501
502
503
504
505
506
507
508
509
510
511
512
513
	var err2 error
	layers = slices.DeleteFunc(layers, func(layer *Layer) bool {
		switch layer.MediaType {
		case "application/vnd.ollama.image.message":
			// if there are new messages, remove the inherited ones
			if len(messages) > 0 {
				return true
			}

			return false
		case "application/vnd.ollama.image.params":
			// merge inherited parameters with new ones
			r, err := layer.Open()
Bruce MacDonald's avatar
Bruce MacDonald committed
514
			if err != nil {
Michael Yang's avatar
Michael Yang committed
515
516
				err2 = err
				return false
Bruce MacDonald's avatar
Bruce MacDonald committed
517
			}
Michael Yang's avatar
Michael Yang committed
518
			defer r.Close()
Bruce MacDonald's avatar
Bruce MacDonald committed
519

Michael Yang's avatar
Michael Yang committed
520
521
522
523
524
			var ps map[string]any
			if err := json.NewDecoder(r).Decode(&ps); err != nil {
				err2 = err
				return false
			}
Michael Yang's avatar
Michael Yang committed
525

Michael Yang's avatar
Michael Yang committed
526
527
528
529
			for k, v := range ps {
				if _, ok := parameters[k]; !ok {
					parameters[k] = v
				}
530
			}
531

Michael Yang's avatar
Michael Yang committed
532
			return true
533
		default:
Michael Yang's avatar
Michael Yang committed
534
			return false
535
		}
Michael Yang's avatar
Michael Yang committed
536
537
538
539
	})

	if err2 != nil {
		return err2
540
541
	}

542
543
	if len(messages) > 0 {
		var b bytes.Buffer
Michael Yang's avatar
Michael Yang committed
544
		if err := json.NewEncoder(&b).Encode(messages); err != nil {
545
546
547
548
549
550
551
552
			return err
		}

		layer, err := NewLayer(&b, "application/vnd.ollama.image.messages")
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
553
		layers = append(layers, layer)
554
555
	}

Michael Yang's avatar
Michael Yang committed
556
	if len(parameters) > 0 {
Michael Yang's avatar
Michael Yang committed
557
		var b bytes.Buffer
Michael Yang's avatar
Michael Yang committed
558
		if err := json.NewEncoder(&b).Encode(parameters); err != nil {
559
560
561
			return err
		}

Michael Yang's avatar
Michael Yang committed
562
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
563
		if err != nil {
Michael Yang's avatar
Michael Yang committed
564
			return err
565
		}
Michael Yang's avatar
Michael Yang committed
566

Michael Yang's avatar
Michael Yang committed
567
		layers = append(layers, layer)
568
569
	}

Michael Yang's avatar
Michael Yang committed
570
571
	digests := make([]string, len(layers))
	for i, layer := range layers {
Michael Yang's avatar
Michael Yang committed
572
		digests[i] = layer.Digest
573
574
	}

Michael Yang's avatar
Michael Yang committed
575
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
576

Michael Yang's avatar
Michael Yang committed
577
578
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
579
580
581
		return err
	}

Michael Yang's avatar
Michael Yang committed
582
	layer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
Michael Yang's avatar
Michael Yang committed
583
	if err != nil {
584
585
586
		return err
	}

Michael Yang's avatar
Michael Yang committed
587
588
589
	for _, layer := range append(layers, layer) {
		if layer.status != "" {
			fn(api.ProgressResponse{Status: layer.status})
590
		}
Michael Yang's avatar
Michael Yang committed
591
	}
592

593
	old, _ := ParseNamedManifest(name)
594

Michael Yang's avatar
Michael Yang committed
595
	fn(api.ProgressResponse{Status: "writing manifest"})
Michael Yang's avatar
Michael Yang committed
596
	if err := WriteManifest(name, layer, layers); err != nil {
597
598
		return err
	}
599

600
601
	if !envconfig.NoPrune && old != nil {
		if err := old.RemoveLayers(); err != nil {
Michael Yang's avatar
Michael Yang committed
602
			return err
603
604
605
		}
	}

Michael Yang's avatar
Michael Yang committed
606
607
	fn(api.ProgressResponse{Status: "success"})
	return nil
608
609
}

Michael Yang's avatar
Michael Yang committed
610
func CopyModel(src, dst model.Name) error {
611
612
613
614
615
616
617
	if !dst.IsFullyQualified() {
		return model.Unqualified(dst)
	}
	if !src.IsFullyQualified() {
		return model.Unqualified(src)
	}

618
619
620
621
	if src.Filepath() == dst.Filepath() {
		return nil
	}

Michael Yang's avatar
Michael Yang committed
622
	manifests, err := GetManifestPath()
623
624
625
626
	if err != nil {
		return err
	}

627
	dstpath := filepath.Join(manifests, dst.Filepath())
Michael Yang's avatar
Michael Yang committed
628
	if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
629
630
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
631

632
	srcpath := filepath.Join(manifests, src.Filepath())
Michael Yang's avatar
Michael Yang committed
633
	srcfile, err := os.Open(srcpath)
Patrick Devine's avatar
Patrick Devine committed
634
635
636
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
637
	defer srcfile.Close()
Patrick Devine's avatar
Patrick Devine committed
638

Michael Yang's avatar
Michael Yang committed
639
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
640
641
642
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
643
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
644

Michael Yang's avatar
Michael Yang committed
645
646
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
647
648
}

649
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}) error {
650
651
652
653
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
654
655
656
657

	walkFunc := func(path string, info os.FileInfo, _ error) error {
		if info.IsDir() {
			return nil
658
659
		}

Michael Yang's avatar
Michael Yang committed
660
661
662
663
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
664

Michael Yang's avatar
Michael Yang committed
665
		// skip the manifest we're trying to delete
666
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
667
			return nil
668
		}
Michael Yang's avatar
Michael Yang committed
669
670
671
672

		// save (i.e. delete from the deleteMap) any files used in other manifests
		manifest, _, err := GetManifest(fmp)
		if err != nil {
Michael Yang's avatar
Michael Yang committed
673
			//nolint:nilerr
Michael Yang's avatar
Michael Yang committed
674
675
676
677
678
679
680
681
			return nil
		}

		for _, layer := range manifest.Layers {
			delete(deleteMap, layer.Digest)
		}

		delete(deleteMap, manifest.Config.Digest)
682
		return nil
Michael Yang's avatar
Michael Yang committed
683
684
685
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
686
687
		return err
	}
688
689

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
690
691
692
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
693
			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
Michael Yang's avatar
Michael Yang committed
694
695
			continue
		}
696
697
698
		if err := os.Remove(fp); err != nil {
			slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
			continue
699
700
701
		}
	}

702
703
704
705
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
706
	deleteMap := make(map[string]struct{})
707
708
709
710
711
712
713
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
714
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
715
716
717
718
719
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
720
		name = strings.ReplaceAll(name, "-", ":")
721
722
723
724
725
726
727
728
729
730
731

		_, 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
732
		}
733
734

		deleteMap[name] = struct{}{}
735
736
	}

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

739
	err = deleteUnusedLayers(nil, deleteMap)
740
741
742
743
	if err != nil {
		return err
	}

744
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
745
746
747
748

	return nil
}

Michael Yang's avatar
Michael Yang committed
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
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
782
func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
783
	mp := ParseModelPath(name)
784
785
	fn(api.ProgressResponse{Status: "retrieving manifest"})

786
787
788
789
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
790
	manifest, _, err := GetManifest(mp)
791
	if err != nil {
792
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
793
794
795
796
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
797
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
798
	layers = append(layers, manifest.Config)
799
800

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
801
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
802
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
803
804
			return err
		}
805
806
	}

807
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
808
809
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
810
811
812
813
814
815

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

Michael Yang's avatar
Michael Yang committed
816
817
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
818
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
819
820
821
822
823
	if err != nil {
		return err
	}
	defer resp.Body.Close()

824
	fn(api.ProgressResponse{Status: "success"})
825
826
827
828

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
832
	var manifest *Manifest
833
834
835
836
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
837
	deleteMap := make(map[string]struct{})
838

839
	if !envconfig.NoPrune {
840
841
842
843
844
845
846
		manifest, _, err = GetManifest(mp)
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}

		if manifest != nil {
			for _, l := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
847
				deleteMap[l.Digest] = struct{}{}
848
			}
Michael Yang's avatar
Michael Yang committed
849
			deleteMap[manifest.Config.Digest] = struct{}{}
850
851
852
		}
	}

853
854
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
855
	}
856

857
	fn(api.ProgressResponse{Status: "pulling manifest"})
858

859
	manifest, err = pullModelManifest(ctx, mp, regOpts)
860
	if err != nil {
861
		return fmt.Errorf("pull model manifest: %s", err)
862
863
864
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
865
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
866
	layers = append(layers, manifest.Config)
867

868
	skipVerify := make(map[string]bool)
869
	for _, layer := range layers {
870
871
872
873
874
875
876
		cacheHit, err := downloadBlob(ctx, downloadOpts{
			mp:      mp,
			digest:  layer.Digest,
			regOpts: regOpts,
			fn:      fn,
		})
		if err != nil {
877
878
			return err
		}
879
		skipVerify[layer.Digest] = cacheHit
880
		delete(deleteMap, layer.Digest)
881
	}
882
	delete(deleteMap, manifest.Config.Digest)
883

Michael Yang's avatar
Michael Yang committed
884
885
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
886
887
888
		if skipVerify[layer.Digest] {
			continue
		}
Michael Yang's avatar
Michael Yang committed
889
		if err := verifyBlob(layer.Digest); err != nil {
890
891
892
893
894
895
896
897
			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
898
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
899
900
				}
			}
Michael Yang's avatar
Michael Yang committed
901
902
903
904
			return err
		}
	}

905
	fn(api.ProgressResponse{Status: "writing manifest"})
906

907
	manifestJSON, err := json.Marshal(manifest)
908
909
910
911
	if err != nil {
		return err
	}

912
	fp, err := mp.GetManifestPath()
913
914
915
	if err != nil {
		return err
	}
916
917
918
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
919

Bruce MacDonald's avatar
Bruce MacDonald committed
920
	err = os.WriteFile(fp, manifestJSON, 0o644)
921
	if err != nil {
922
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
923
924
925
		return err
	}

926
927
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
928
		err = deleteUnusedLayers(nil, deleteMap)
929
930
931
932
933
		if err != nil {
			return err
		}
	}

934
	fn(api.ProgressResponse{Status: "success"})
935
936
937
938

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
942
943
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
944
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
945
946
947
948
949
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
950
	var m *Manifest
951
952
953
954
955
956
957
958
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

	return m, err
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
959
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
960
961
962
963
964
965
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
966
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
967
968
}

969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
var errUnauthorized = fmt.Errorf("unauthorized: access denied")

// getTokenSubject returns the subject of a JWT token, it does not validate the token
func getTokenSubject(token string) string {
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return ""
	}

	payload := parts[1]
	payloadBytes, err := base64.RawURLEncoding.DecodeString(payload)
	if err != nil {
		slog.Error(fmt.Sprintf("failed to decode jwt payload: %v", err))
		return ""
	}

	var payloadMap map[string]interface{}
	if err := json.Unmarshal(payloadBytes, &payloadMap); err != nil {
		slog.Error(fmt.Sprintf("failed to unmarshal payload JSON: %v", err))
		return ""
	}

	sub, ok := payloadMap["sub"]
	if !ok {
		slog.Error("jwt does not contain 'sub' field")
		return ""
	}

	return fmt.Sprintf("%s", sub)
}
999

Michael Yang's avatar
Michael Yang committed
1000
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
1001
	anonymous := true // access will default to anonymous if no user is found associated with the public key
Michael Yang's avatar
lint  
Michael Yang committed
1002
	for range 2 {
Michael Yang's avatar
Michael Yang committed
1003
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1004
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1005
			if !errors.Is(err, context.Canceled) {
1006
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1007
1008
			}

Michael Yang's avatar
Michael Yang committed
1009
1010
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1011
1012
1013
1014

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
1015
1016
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
1017
1018
1019
			if err != nil {
				return nil, err
			}
1020
			anonymous = getTokenSubject(token) == "anonymous"
Michael Yang's avatar
Michael Yang committed
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
			regOpts.Token = token
			if body != nil {
				_, err = body.Seek(0, io.SeekStart)
				if err != nil {
					return nil, err
				}
			}
		case resp.StatusCode == http.StatusNotFound:
			return nil, os.ErrNotExist
		case resp.StatusCode >= http.StatusBadRequest:
			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
1038
1039
1040
		}
	}

1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
	if anonymous {
		// no user is associated with the public key, and the request requires non-anonymous access
		pubKey, nestedErr := auth.GetPublicKey()
		if nestedErr != nil {
			slog.Error(fmt.Sprintf("couldn't get public key: %v", nestedErr))
			return nil, errUnauthorized
		}
		return nil, &errtypes.UnknownOllamaKey{Key: pubKey}
	}
	// user is associated with the public key, but is not authorized to make the request
Michael Yang's avatar
Michael Yang committed
1051
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1052
1053
}

Michael Yang's avatar
Michael Yang committed
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
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)
		}
	}

	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))

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

		req.ContentLength = contentLength
	}

Michael Yang's avatar
Michael Yang committed
1087
	resp, err := http.DefaultClient.Do(req)
Michael Yang's avatar
Michael Yang committed
1088
1089
1090
1091
1092
1093
1094
	if err != nil {
		return nil, err
	}

	return resp, nil
}

Patrick Devine's avatar
Patrick Devine committed
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
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
1118
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
1119
1120
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
1121
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
1122
1123
1124
1125
1126
1127
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
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 {
1144
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1145
1146
1147
1148
	}

	return nil
}