images.go 27.4 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
var errCapabilityCompletion = errors.New("completion")

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

Michael Yang's avatar
tools  
Michael Yang committed
41
42
43
44
const (
	CapabilityCompletion = Capability("completion")
	CapabilityTools      = Capability("tools")
)
Michael Yang's avatar
Michael Yang committed
45

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

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

	Template *template.Template
68
69
}

Michael Yang's avatar
Michael Yang committed
70
71
72
73
// CheckCapabilities checks if the model has the specified capabilities returning an error describing
// any missing or unknown capabilities
func (m *Model) CheckCapabilities(caps ...Capability) error {
	var errs []error
Michael Yang's avatar
Michael Yang committed
74
75
76
	for _, cap := range caps {
		switch cap {
		case CapabilityCompletion:
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
			f, err := os.Open(m.ModelPath)
			if err != nil {
				slog.Error("couldn't open model file", "error", err)
				continue
			}
			defer f.Close()

			// TODO(mxyng): decode the GGML into model to avoid doing this multiple times
			ggml, _, err := llm.DecodeGGML(f, 0)
			if err != nil {
				slog.Error("couldn't decode ggml", "error", err)
				continue
			}

			if _, ok := ggml.KV()[fmt.Sprintf("%s.pooling_type", ggml.KV().Architecture())]; ok {
Michael Yang's avatar
Michael Yang committed
92
				errs = append(errs, errCapabilityCompletion)
Michael Yang's avatar
Michael Yang committed
93
			}
Michael Yang's avatar
tools  
Michael Yang committed
94
95
96
97
		case CapabilityTools:
			if !slices.Contains(m.Template.Vars(), "tools") {
				errs = append(errs, errors.New("tools"))
			}
Michael Yang's avatar
Michael Yang committed
98
99
		default:
			slog.Error("unknown capability", "capability", cap)
Michael Yang's avatar
Michael Yang committed
100
			return fmt.Errorf("unknown capability: %s", cap)
Michael Yang's avatar
Michael Yang committed
101
102
103
		}
	}

Michael Yang's avatar
Michael Yang committed
104
	if err := errors.Join(errs...); err != nil {
Michael Yang's avatar
tools  
Michael Yang committed
105
		return fmt.Errorf("does not support %w", errors.Join(errs...))
Michael Yang's avatar
Michael Yang committed
106
107
108
	}

	return nil
109
110
}

Michael Yang's avatar
Michael Yang committed
111
func (m *Model) String() string {
112
	var modelfile parser.File
Michael Yang's avatar
Michael Yang committed
113

114
	modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
115
116
117
		Name: "model",
		Args: m.ModelPath,
	})
118

Michael Yang's avatar
Michael Yang committed
119
	for _, adapter := range m.AdapterPaths {
120
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
121
122
			Name: "adapter",
			Args: adapter,
Michael Yang's avatar
Michael Yang committed
123
		})
124
125
	}

Michael Yang's avatar
Michael Yang committed
126
	for _, projector := range m.ProjectorPaths {
127
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
128
129
			Name: "model",
			Args: projector,
Michael Yang's avatar
Michael Yang committed
130
		})
131
132
	}

Michael Yang's avatar
Michael Yang committed
133
	if m.Template != nil {
134
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
135
			Name: "template",
Michael Yang's avatar
Michael Yang committed
136
			Args: m.Template.String(),
Michael Yang's avatar
Michael Yang committed
137
		})
138
139
	}

Michael Yang's avatar
Michael Yang committed
140
	if m.System != "" {
141
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
142
143
			Name: "system",
			Args: m.System,
Michael Yang's avatar
Michael Yang committed
144
		})
145
146
147
148
149
150
	}

	for k, v := range m.Options {
		switch v := v.(type) {
		case []any:
			for _, s := range v {
151
				modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
152
153
154
					Name: k,
					Args: fmt.Sprintf("%v", s),
				})
155
156
			}
		default:
157
			modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
158
159
160
				Name: k,
				Args: fmt.Sprintf("%v", v),
			})
161
162
163
164
		}
	}

	for _, license := range m.License {
165
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
166
167
168
			Name: "license",
			Args: license,
		})
169
170
171
	}

	for _, msg := range m.Messages {
172
		modelfile.Commands = append(modelfile.Commands, parser.Command{
Michael Yang's avatar
Michael Yang committed
173
174
175
			Name: "message",
			Args: fmt.Sprintf("%s %s", msg.Role, msg.Content),
		})
176
177
	}

Michael Yang's avatar
Michael Yang committed
178
	return modelfile.String()
179
180
}

181
182
183
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
184
185
186
}

type ConfigV2 struct {
187
188
189
190
191
192
	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"`

193
	// required by spec
194
195
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
196
	RootFS       RootFS `json:"rootfs"`
197
198
199
200
201
202
203
}

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

Michael Yang's avatar
Michael Yang committed
204
func GetManifest(mp ModelPath) (*Manifest, string, error) {
205
	fp, err := mp.GetManifestPath()
206
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
207
		return nil, "", err
208
	}
209

210
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
211
		return nil, "", err
212
213
	}

Michael Yang's avatar
Michael Yang committed
214
	var manifest *Manifest
215

216
	bts, err := os.ReadFile(fp)
217
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
218
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
219
220
	}

Patrick Devine's avatar
Patrick Devine committed
221
222
223
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

224
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
225
		return nil, "", err
226
227
	}

Patrick Devine's avatar
Patrick Devine committed
228
	return manifest, shaStr, nil
229
230
231
}

func GetModel(name string) (*Model, error) {
232
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
233
	manifest, digest, err := GetManifest(mp)
234
235
236
237
238
	if err != nil {
		return nil, err
	}

	model := &Model{
239
240
241
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
Michael Yang's avatar
Michael Yang committed
242
		Template:  template.DefaultTemplate,
243
244
	}

245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
	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
	}

260
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
261
		filename, err := GetBlobsPath(layer.Digest)
262
263
264
265
		if err != nil {
			return nil, err
		}

266
267
268
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
269
			model.ParentModel = layer.From
270
		case "application/vnd.ollama.image.embed":
271
272
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
273
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
274
275
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
276
277
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
Michael Yang's avatar
Michael Yang committed
278
279
		case "application/vnd.ollama.image.prompt",
			"application/vnd.ollama.image.template":
280
281
282
283
284
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
285
			model.Template, err = template.Parse(string(bts))
286
287
288
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
289
		case "application/vnd.ollama.image.system":
290
291
292
293
294
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

Michael Yang's avatar
Michael Yang committed
295
			model.System = string(bts)
296
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
297
298
299
300
301
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
302

303
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
304
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
305
306
				return nil, err
			}
307
308
309
310
311
312
313
314
315
316
		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
317
318
319
320
321
322
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
323
324
325
326
327
328
		}
	}

	return model, nil
}

Michael Yang's avatar
Michael Yang committed
329
func realpath(rel, from string) string {
330
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
331
	if err != nil {
332
		return from
333
334
	}

Michael Yang's avatar
Michael Yang committed
335
	home, err := os.UserHomeDir()
336
	if err != nil {
Michael Yang's avatar
Michael Yang committed
337
		return abspath
338
339
	}

340
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
341
		return home
342
343
344
345
	} else if strings.HasPrefix(from, "~/") {
		return filepath.Join(home, from[2:])
	}

Michael Yang's avatar
Michael Yang committed
346
	if _, err := os.Stat(filepath.Join(rel, from)); err == nil {
347
		// this is a file relative to the Modelfile
Michael Yang's avatar
Michael Yang committed
348
		return filepath.Join(rel, from)
349
350
	}

Michael Yang's avatar
Michael Yang committed
351
352
353
	return abspath
}

354
func CreateModel(ctx context.Context, name model.Name, modelFileDir, quantization string, modelfile *parser.File, fn func(resp api.ProgressResponse)) (err error) {
355
356
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
357
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
358
359
360
		RootFS: RootFS{
			Type: "layers",
		},
361
362
	}

Michael Yang's avatar
Michael Yang committed
363
364
	var messages []*api.Message
	parameters := make(map[string]any)
Michael Yang's avatar
Michael Yang committed
365

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

370
		switch c.Name {
Michael Yang's avatar
Michael Yang committed
371
		case "model", "adapter":
372
			var baseLayers []*layerGGML
Michael Yang's avatar
rebase  
Michael Yang committed
373
			if name := model.ParseName(c.Args); name.IsValid() {
Michael Yang's avatar
Michael Yang committed
374
				baseLayers, err = parseFromModel(ctx, name, fn)
Michael Yang's avatar
Michael Yang committed
375
376
377
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
378
			} else if strings.HasPrefix(c.Args, "@") {
379
				digest := strings.TrimPrefix(c.Args, "@")
Michael Yang's avatar
Michael Yang committed
380
381
				if ib, ok := intermediateBlobs[digest]; ok {
					p, err := GetBlobsPath(ib)
382
383
384
385
386
387
388
389
390
					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
391
392
						fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib)})
						digest = ib
393
394
395
396
					}
				}

				blobpath, err := GetBlobsPath(digest)
Michael Yang's avatar
Michael Yang committed
397
				if err != nil {
Michael Yang's avatar
Michael Yang committed
398
					return err
399
				}
400

Michael Yang's avatar
Michael Yang committed
401
				blob, err := os.Open(blobpath)
Michael Yang's avatar
Michael Yang committed
402
403
404
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
405
				defer blob.Close()
Michael Yang's avatar
Michael Yang committed
406

407
				baseLayers, err = parseFromFile(ctx, blob, digest, fn)
Michael Yang's avatar
Michael Yang committed
408
409
410
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
411
412
			} else if file, err := os.Open(realpath(modelFileDir, c.Args)); err == nil {
				defer file.Close()
Michael Yang's avatar
Michael Yang committed
413

414
				baseLayers, err = parseFromFile(ctx, file, "", fn)
Michael Yang's avatar
Michael Yang committed
415
				if err != nil {
Michael Yang's avatar
Michael Yang committed
416
417
					return err
				}
Michael Yang's avatar
Michael Yang committed
418
419
420
			} else {
				return fmt.Errorf("invalid model reference: %s", c.Args)
			}
Michael Yang's avatar
Michael Yang committed
421

Michael Yang's avatar
Michael Yang committed
422
			for _, baseLayer := range baseLayers {
Michael Yang's avatar
Michael Yang committed
423
424
425
426
				if quantization != "" &&
					baseLayer.MediaType == "application/vnd.ollama.image.model" &&
					baseLayer.GGML != nil &&
					baseLayer.GGML.Name() == "gguf" {
Michael Yang's avatar
Michael Yang committed
427
					want, err := llm.ParseFileType(quantization)
428
					if err != nil {
Michael Yang's avatar
Michael Yang committed
429
						return err
430
					}
Michael Yang's avatar
Michael Yang committed
431

Michael Yang's avatar
Michael Yang committed
432
433
					ft := baseLayer.GGML.KV().FileType()
					if !slices.Contains([]string{"F16", "F32"}, ft.String()) {
Michael Yang's avatar
Michael Yang committed
434
						return errors.New("quantization is only supported for F16 and F32 models")
Michael Yang's avatar
Michael Yang committed
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
					} 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
454
						layer, err := NewLayer(temp, baseLayer.MediaType)
Michael Yang's avatar
Michael Yang committed
455
456
457
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
458

Michael Yang's avatar
Michael Yang committed
459
460
461
462
						if _, err := temp.Seek(0, io.SeekStart); err != nil {
							return err
						}

463
						ggml, _, err := llm.DecodeGGML(temp, 0)
Michael Yang's avatar
Michael Yang committed
464
465
						if err != nil {
							return err
Michael Yang's avatar
Michael Yang committed
466
467
						}

Michael Yang's avatar
Michael Yang committed
468
469
						baseLayer.Layer = layer
						baseLayer.GGML = ggml
Michael Yang's avatar
Michael Yang committed
470
					}
471
				}
472

Michael Yang's avatar
Michael Yang committed
473
474
475
476
				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
477
					config.FileType = cmp.Or(config.FileType, baseLayer.GGML.KV().FileType().String())
Michael Yang's avatar
Michael Yang committed
478
					config.ModelFamilies = append(config.ModelFamilies, baseLayer.GGML.KV().Architecture())
479
480
				}

Michael Yang's avatar
Michael Yang committed
481
				layers = append(layers, baseLayer.Layer)
482
			}
Michael Yang's avatar
Michael Yang committed
483
484
485
486
		case "license", "template", "system":
			if c.Name != "license" {
				// replace
				layers = slices.DeleteFunc(layers, func(layer *Layer) bool {
487
488
489
490
491
492
493
494
495
					if layer.MediaType != mediatype {
						return false
					}

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

					return true
Michael Yang's avatar
Michael Yang committed
496
				})
Michael Yang's avatar
Michael Yang committed
497
498
			}

499
500
501
502
503
504
			blob := strings.NewReader(c.Args)
			layer, err := NewLayer(blob, mediatype)
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
505
506
507
508
509
510
511
512
513
514
			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}})
515
			if err != nil {
Michael Yang's avatar
Michael Yang committed
516
				return err
517
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
518

Michael Yang's avatar
Michael Yang committed
519
520
521
522
523
524
525
526
527
528
529
			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
530

Michael Yang's avatar
Michael Yang committed
531
532
533
534
535
536
537
538
539
540
541
542
543
	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
544
			if err != nil {
Michael Yang's avatar
Michael Yang committed
545
546
				err2 = err
				return false
Bruce MacDonald's avatar
Bruce MacDonald committed
547
			}
Michael Yang's avatar
Michael Yang committed
548
			defer r.Close()
Bruce MacDonald's avatar
Bruce MacDonald committed
549

Michael Yang's avatar
Michael Yang committed
550
551
552
553
554
			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
555

Michael Yang's avatar
Michael Yang committed
556
557
558
559
			for k, v := range ps {
				if _, ok := parameters[k]; !ok {
					parameters[k] = v
				}
560
			}
561

Michael Yang's avatar
Michael Yang committed
562
			return true
563
		default:
Michael Yang's avatar
Michael Yang committed
564
			return false
565
		}
Michael Yang's avatar
Michael Yang committed
566
567
568
569
	})

	if err2 != nil {
		return err2
570
571
	}

572
573
	if len(messages) > 0 {
		var b bytes.Buffer
Michael Yang's avatar
Michael Yang committed
574
		if err := json.NewEncoder(&b).Encode(messages); err != nil {
575
576
577
578
579
580
581
582
			return err
		}

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

Michael Yang's avatar
Michael Yang committed
583
		layers = append(layers, layer)
584
585
	}

Michael Yang's avatar
Michael Yang committed
586
	if len(parameters) > 0 {
Michael Yang's avatar
Michael Yang committed
587
		var b bytes.Buffer
Michael Yang's avatar
Michael Yang committed
588
		if err := json.NewEncoder(&b).Encode(parameters); err != nil {
589
590
591
			return err
		}

Michael Yang's avatar
Michael Yang committed
592
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
593
		if err != nil {
Michael Yang's avatar
Michael Yang committed
594
			return err
595
		}
Michael Yang's avatar
Michael Yang committed
596

Michael Yang's avatar
Michael Yang committed
597
		layers = append(layers, layer)
598
599
	}

Michael Yang's avatar
Michael Yang committed
600
601
	digests := make([]string, len(layers))
	for i, layer := range layers {
Michael Yang's avatar
Michael Yang committed
602
		digests[i] = layer.Digest
603
604
	}

Michael Yang's avatar
Michael Yang committed
605
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
606

Michael Yang's avatar
Michael Yang committed
607
608
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
609
610
611
		return err
	}

Michael Yang's avatar
Michael Yang committed
612
	layer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
Michael Yang's avatar
Michael Yang committed
613
	if err != nil {
614
615
616
		return err
	}

Michael Yang's avatar
Michael Yang committed
617
618
619
	for _, layer := range append(layers, layer) {
		if layer.status != "" {
			fn(api.ProgressResponse{Status: layer.status})
620
		}
Michael Yang's avatar
Michael Yang committed
621
	}
622

623
	old, _ := ParseNamedManifest(name)
624

Michael Yang's avatar
Michael Yang committed
625
	fn(api.ProgressResponse{Status: "writing manifest"})
Michael Yang's avatar
Michael Yang committed
626
	if err := WriteManifest(name, layer, layers); err != nil {
627
628
		return err
	}
629

630
631
	if !envconfig.NoPrune && old != nil {
		if err := old.RemoveLayers(); err != nil {
Michael Yang's avatar
Michael Yang committed
632
			return err
633
634
635
		}
	}

Michael Yang's avatar
Michael Yang committed
636
637
	fn(api.ProgressResponse{Status: "success"})
	return nil
638
639
}

Michael Yang's avatar
Michael Yang committed
640
func CopyModel(src, dst model.Name) error {
641
642
643
644
645
646
647
	if !dst.IsFullyQualified() {
		return model.Unqualified(dst)
	}
	if !src.IsFullyQualified() {
		return model.Unqualified(src)
	}

648
649
650
651
	if src.Filepath() == dst.Filepath() {
		return nil
	}

Michael Yang's avatar
Michael Yang committed
652
	manifests, err := GetManifestPath()
653
654
655
656
	if err != nil {
		return err
	}

657
	dstpath := filepath.Join(manifests, dst.Filepath())
Michael Yang's avatar
Michael Yang committed
658
	if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
659
660
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
661

662
	srcpath := filepath.Join(manifests, src.Filepath())
Michael Yang's avatar
Michael Yang committed
663
	srcfile, err := os.Open(srcpath)
Patrick Devine's avatar
Patrick Devine committed
664
665
666
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
667
	defer srcfile.Close()
Patrick Devine's avatar
Patrick Devine committed
668

Michael Yang's avatar
Michael Yang committed
669
	dstfile, err := os.Create(dstpath)
Patrick Devine's avatar
Patrick Devine committed
670
671
672
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
673
	defer dstfile.Close()
Patrick Devine's avatar
Patrick Devine committed
674

Michael Yang's avatar
Michael Yang committed
675
676
	_, err = io.Copy(dstfile, srcfile)
	return err
Patrick Devine's avatar
Patrick Devine committed
677
678
}

679
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}) error {
680
681
682
683
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
684
685
686
687

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

Michael Yang's avatar
Michael Yang committed
690
691
692
693
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
694

Michael Yang's avatar
Michael Yang committed
695
		// skip the manifest we're trying to delete
696
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
697
			return nil
698
		}
Michael Yang's avatar
Michael Yang committed
699
700
701
702

		// 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
703
			//nolint:nilerr
Michael Yang's avatar
Michael Yang committed
704
705
706
707
708
709
710
711
			return nil
		}

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

		delete(deleteMap, manifest.Config.Digest)
712
		return nil
Michael Yang's avatar
Michael Yang committed
713
714
715
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
716
717
		return err
	}
718
719

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
720
721
722
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
723
			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
Michael Yang's avatar
Michael Yang committed
724
725
			continue
		}
726
727
728
		if err := os.Remove(fp); err != nil {
			slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
			continue
729
730
731
		}
	}

732
733
734
735
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
736
	deleteMap := make(map[string]struct{})
737
738
739
740
741
742
743
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
744
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
745
746
747
748
749
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
750
		name = strings.ReplaceAll(name, "-", ":")
751
752
753
754
755
756
757
758
759
760
761

		_, 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
762
		}
763
764

		deleteMap[name] = struct{}{}
765
766
	}

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

769
	err = deleteUnusedLayers(nil, deleteMap)
770
771
772
773
	if err != nil {
		return err
	}

774
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
775
776
777
778

	return nil
}

Michael Yang's avatar
Michael Yang committed
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
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
812
func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
813
	mp := ParseModelPath(name)
814
815
	fn(api.ProgressResponse{Status: "retrieving manifest"})

816
817
818
819
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
820
	manifest, _, err := GetManifest(mp)
821
	if err != nil {
822
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
823
824
825
826
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
827
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
828
	layers = append(layers, manifest.Config)
829
830

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
831
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
832
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
833
834
			return err
		}
835
836
	}

837
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
838
839
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
840
841
842
843
844
845

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

Michael Yang's avatar
Michael Yang committed
846
847
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
848
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
849
850
851
852
853
	if err != nil {
		return err
	}
	defer resp.Body.Close()

854
	fn(api.ProgressResponse{Status: "success"})
855
856
857
858

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
862
	var manifest *Manifest
863
864
865
866
	var err error
	var noprune string

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

869
	if !envconfig.NoPrune {
870
871
872
873
874
875
876
		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
877
				deleteMap[l.Digest] = struct{}{}
878
			}
Michael Yang's avatar
Michael Yang committed
879
			deleteMap[manifest.Config.Digest] = struct{}{}
880
881
882
		}
	}

883
884
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
885
	}
886

887
	fn(api.ProgressResponse{Status: "pulling manifest"})
888

889
	manifest, err = pullModelManifest(ctx, mp, regOpts)
890
	if err != nil {
891
		return fmt.Errorf("pull model manifest: %s", err)
892
893
894
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
895
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
896
	layers = append(layers, manifest.Config)
897

898
	skipVerify := make(map[string]bool)
899
	for _, layer := range layers {
900
901
902
903
904
905
906
		cacheHit, err := downloadBlob(ctx, downloadOpts{
			mp:      mp,
			digest:  layer.Digest,
			regOpts: regOpts,
			fn:      fn,
		})
		if err != nil {
907
908
			return err
		}
909
		skipVerify[layer.Digest] = cacheHit
910
		delete(deleteMap, layer.Digest)
911
	}
912
	delete(deleteMap, manifest.Config.Digest)
913

Michael Yang's avatar
Michael Yang committed
914
915
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
916
917
918
		if skipVerify[layer.Digest] {
			continue
		}
Michael Yang's avatar
Michael Yang committed
919
		if err := verifyBlob(layer.Digest); err != nil {
920
921
922
923
924
925
926
927
			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
928
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
929
930
				}
			}
Michael Yang's avatar
Michael Yang committed
931
932
933
934
			return err
		}
	}

935
	fn(api.ProgressResponse{Status: "writing manifest"})
936

937
	manifestJSON, err := json.Marshal(manifest)
938
939
940
941
	if err != nil {
		return err
	}

942
	fp, err := mp.GetManifestPath()
943
944
945
	if err != nil {
		return err
	}
946
947
948
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
949

Bruce MacDonald's avatar
Bruce MacDonald committed
950
	err = os.WriteFile(fp, manifestJSON, 0o644)
951
	if err != nil {
952
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
953
954
955
		return err
	}

956
957
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
958
		err = deleteUnusedLayers(nil, deleteMap)
959
960
961
962
963
		if err != nil {
			return err
		}
	}

964
	fn(api.ProgressResponse{Status: "success"})
965
966
967
968

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
972
973
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
974
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
975
976
977
978
979
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
980
	var m *Manifest
981
982
983
984
985
986
987
988
	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
989
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
990
991
992
993
994
995
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
996
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
997
998
}

999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
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)
}
1029

Michael Yang's avatar
Michael Yang committed
1030
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
1031
	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
1032
	for range 2 {
Michael Yang's avatar
Michael Yang committed
1033
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1034
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1035
			if !errors.Is(err, context.Canceled) {
1036
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1037
1038
			}

Michael Yang's avatar
Michael Yang committed
1039
1040
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1041
1042
1043
1044

		switch {
		case resp.StatusCode == http.StatusUnauthorized:
			// Handle authentication error with one retry
Michael Yang's avatar
Michael Yang committed
1045
1046
			challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
			token, err := getAuthorizationToken(ctx, challenge)
Michael Yang's avatar
Michael Yang committed
1047
1048
1049
			if err != nil {
				return nil, err
			}
1050
			anonymous = getTokenSubject(token) == "anonymous"
Michael Yang's avatar
Michael Yang committed
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
			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
1068
1069
1070
		}
	}

1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
	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
1081
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1082
1083
}

Michael Yang's avatar
Michael Yang committed
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
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
1117
	resp, err := http.DefaultClient.Do(req)
Michael Yang's avatar
Michael Yang committed
1118
1119
1120
1121
1122
1123
1124
	if err != nil {
		return nil, err
	}

	return resp, nil
}

Patrick Devine's avatar
Patrick Devine committed
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
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
1148
func parseRegistryChallenge(authStr string) registryChallenge {
Patrick Devine's avatar
Patrick Devine committed
1149
1150
	authStr = strings.TrimPrefix(authStr, "Bearer ")

Michael Yang's avatar
Michael Yang committed
1151
	return registryChallenge{
Patrick Devine's avatar
Patrick Devine committed
1152
1153
1154
1155
1156
1157
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
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 {
1174
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1175
1176
1177
1178
	}

	return nil
}