images.go 33.2 KB
Newer Older
1
2
3
package server

import (
4
	"bufio"
5
	"bytes"
6
	"context"
7
8
9
10
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
11
	"html/template"
12
13
14
15
16
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
17
	"reflect"
18
19
20
21
	"strconv"
	"strings"

	"github.com/jmorganca/ollama/api"
22
	"github.com/jmorganca/ollama/llm"
23
	"github.com/jmorganca/ollama/parser"
24
	"github.com/jmorganca/ollama/vector"
25
26
)

27
28
const MaxRetries = 3

29
30
31
32
type RegistryOptions struct {
	Insecure bool
	Username string
	Password string
Patrick Devine's avatar
Patrick Devine committed
33
	Token    string
34
35
}

36
type Model struct {
37
38
39
40
41
42
43
44
	Name         string `json:"name"`
	ModelPath    string
	AdapterPaths []string
	Template     string
	System       string
	Digest       string
	Options      map[string]interface{}
	Embeddings   []vector.Embedding
45
46
}

Bruce MacDonald's avatar
Bruce MacDonald committed
47
func (m *Model) Prompt(request api.GenerateRequest, embedding string) (string, error) {
48
49
50
51
52
53
	t := m.Template
	if request.Template != "" {
		t = request.Template
	}

	tmpl, err := template.New("").Parse(t)
54
55
56
57
58
	if err != nil {
		return "", err
	}

	var vars struct {
Michael Yang's avatar
Michael Yang committed
59
		First  bool
60
61
		System string
		Prompt string
62
		Embed  string
63
64
65

		// deprecated: versions <= 0.0.7 used this to omit the system prompt
		Context []int
66
67
	}

Michael Yang's avatar
Michael Yang committed
68
	vars.First = len(request.Context) == 0
69
70
	vars.System = m.System
	vars.Prompt = request.Prompt
Michael Yang's avatar
Michael Yang committed
71
	vars.Context = request.Context
Bruce MacDonald's avatar
Bruce MacDonald committed
72
	vars.Embed = embedding
73

74
75
76
77
	if request.System != "" {
		vars.System = request.System
	}

78
79
80
81
82
83
84
85
	var sb strings.Builder
	if err := tmpl.Execute(&sb, vars); err != nil {
		return "", err
	}

	return sb.String(), nil
}

86
87
88
89
90
91
92
93
94
95
96
97
98
type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
	Config        Layer    `json:"config"`
	Layers        []*Layer `json:"layers"`
}

type Layer struct {
	MediaType string `json:"mediaType"`
	Digest    string `json:"digest"`
	Size      int    `json:"size"`
}

Michael Yang's avatar
Michael Yang committed
99
type LayerReader struct {
100
	Layer
Michael Yang's avatar
Michael Yang committed
101
	io.Reader
102
103
104
}

type ConfigV2 struct {
105
106
107
108
109
110
	ModelFamily llm.ModelFamily `json:"model_family"`
	ModelType   llm.ModelType   `json:"model_type"`
	FileType    llm.FileType    `json:"file_type"`
	RootFS      RootFS          `json:"rootfs"`

	// required by spec
111
112
113
114
115
116
117
118
119
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
}

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

Patrick Devine's avatar
Patrick Devine committed
120
121
122
123
124
125
126
127
128
func (m *ManifestV2) GetTotalSize() int {
	var total int
	for _, layer := range m.Layers {
		total += layer.Size
	}
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
129
130
func GetManifest(mp ModelPath) (*ManifestV2, error) {
	fp, err := mp.GetManifestPath(false)
131
132
133
	if err != nil {
		return nil, err
	}
134

135
136
	if _, err = os.Stat(fp); err != nil {
		return nil, err
137
138
139
140
	}

	var manifest *ManifestV2

141
	bts, err := os.ReadFile(fp)
142
143
144
145
	if err != nil {
		return nil, fmt.Errorf("couldn't open file '%s'", fp)
	}

146
	if err := json.Unmarshal(bts, &manifest); err != nil {
147
148
149
150
151
152
153
		return nil, err
	}

	return manifest, nil
}

func GetModel(name string) (*Model, error) {
Patrick Devine's avatar
Patrick Devine committed
154
155
156
	mp := ParseModelPath(name)

	manifest, err := GetManifest(mp)
157
158
159
160
161
	if err != nil {
		return nil, err
	}

	model := &Model{
162
		Name:   mp.GetFullTagname(),
Jeffrey Morgan's avatar
Jeffrey Morgan committed
163
		Digest: manifest.Config.Digest,
164
165
166
	}

	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
167
		filename, err := GetBlobsPath(layer.Digest)
168
169
170
171
		if err != nil {
			return nil, err
		}

172
173
174
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
175
176
177
178
179
180
181
182
183
184
		case "application/vnd.ollama.image.embed":
			file, err := os.Open(filename)
			if err != nil {
				return nil, fmt.Errorf("failed to open file: %s", filename)
			}
			defer file.Close()

			if err = json.NewDecoder(file).Decode(&model.Embeddings); err != nil {
				return nil, err
			}
185
186
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
187
188
189
190
191
192
193
194
195
		case "application/vnd.ollama.image.template":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
		case "application/vnd.ollama.image.system":
			bts, err := os.ReadFile(filename)
196
197
198
			if err != nil {
				return nil, err
			}
199
200

			model.System = string(bts)
201
202
203
204
205
206
207
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
208
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
209
210
211
212
213
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
214

215
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
216
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
217
218
				return nil, err
			}
219
220
221
222
223
224
		}
	}

	return model, nil
}

225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
func filenameWithPath(path, f string) (string, error) {
	// if filePath starts with ~/, replace it with the user's home directory.
	if strings.HasPrefix(f, "~/") {
		parts := strings.Split(f, "/")
		home, err := os.UserHomeDir()
		if err != nil {
			return "", fmt.Errorf("failed to open file: %v", err)
		}

		f = filepath.Join(home, filepath.Join(parts[1:]...))
	}

	// if filePath is not an absolute path, make it relative to the modelfile path
	if !filepath.IsAbs(f) {
		f = filepath.Join(filepath.Dir(path), f)
	}

	return f, nil
}

245
func CreateModel(ctx context.Context, name string, path string, fn func(resp api.ProgressResponse)) error {
246
247
	mf, err := os.Open(path)
	if err != nil {
248
		fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't open modelfile '%s'", path)})
249
		return fmt.Errorf("failed to open file: %w", err)
250
	}
251
	defer mf.Close()
252

253
	fn(api.ProgressResponse{Status: "parsing modelfile"})
254
255
256
257
258
	commands, err := parser.Parse(mf)
	if err != nil {
		return err
	}

259
260
261
262
263
	config := ConfigV2{
		Architecture: "amd64",
		OS:           "linux",
	}

Michael Yang's avatar
Michael Yang committed
264
	var layers []*LayerReader
265
	params := make(map[string][]string)
266
	embed := EmbeddingParams{fn: fn}
267
	for _, c := range commands {
268
		log.Printf("[%s] - %s\n", c.Name, c.Args)
269
270
		switch c.Name {
		case "model":
271
			fn(api.ProgressResponse{Status: "looking for model"})
272
			embed.model = c.Args
273
			mf, err := GetManifest(ParseModelPath(c.Args))
274
			if err != nil {
275
276
277
				modelFile, err := filenameWithPath(path, c.Args)
				if err != nil {
					return err
278
				}
279
				if _, err := os.Stat(modelFile); err != nil {
280
281
282
					// the model file does not exist, try pulling it
					if errors.Is(err, os.ErrNotExist) {
						fn(api.ProgressResponse{Status: "pulling model file"})
283
						if err := PullModel(ctx, c.Args, &RegistryOptions{}, fn); err != nil {
284
285
							return err
						}
286
						mf, err = GetManifest(ParseModelPath(c.Args))
287
288
289
290
291
292
293
						if err != nil {
							return fmt.Errorf("failed to open file after pull: %v", err)
						}
					} else {
						return err
					}
				} else {
294
					embed.model = modelFile
295
296
					// create a model from this specified file
					fn(api.ProgressResponse{Status: "creating model layer"})
297
					file, err := os.Open(modelFile)
298
299
300
301
302
					if err != nil {
						return fmt.Errorf("failed to open file: %v", err)
					}
					defer file.Close()

303
304
305
306
307
308
309
310
311
312
313
314
					ggml, err := llm.DecodeGGML(file, llm.ModelFamilyLlama)
					if err != nil {
						return err
					}

					config.ModelFamily = ggml.ModelFamily
					config.ModelType = ggml.ModelType
					config.FileType = ggml.FileType

					// reset the file
					file.Seek(0, io.SeekStart)

315
316
317
318
319
320
					l, err := CreateLayer(file)
					if err != nil {
						return fmt.Errorf("failed to create layer: %v", err)
					}
					l.MediaType = "application/vnd.ollama.image.model"
					layers = append(layers, l)
321
				}
322
			}
323

324
			if mf != nil {
325
326
327
328
329
330
331
332
333
				log.Printf("manifest = %#v", mf)
				for _, l := range mf.Layers {
					newLayer, err := GetLayerWithBufferFromLayer(l)
					if err != nil {
						return err
					}
					layers = append(layers, newLayer)
				}
			}
334
335
		case "embed":
			embedFilePath, err := filenameWithPath(path, c.Args)
Michael Yang's avatar
Michael Yang committed
336
337
338
			if err != nil {
				return err
			}
339
			embed.files = append(embed.files, embedFilePath)
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
		case "adapter":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})

			fp := c.Args
			if strings.HasPrefix(fp, "~/") {
				parts := strings.Split(fp, "/")
				home, err := os.UserHomeDir()
				if err != nil {
					return fmt.Errorf("failed to open file: %v", err)
				}

				fp = filepath.Join(home, filepath.Join(parts[1:]...))
			}

			// If filePath is not an absolute path, make it relative to the modelfile path
			if !filepath.IsAbs(fp) {
				fp = filepath.Join(filepath.Dir(path), fp)
			}

			// create a model from this specified file
			fn(api.ProgressResponse{Status: "creating model layer"})

			file, err := os.Open(fp)
			if err != nil {
				return fmt.Errorf("failed to open file: %v", err)
			}
			defer file.Close()

			l, err := CreateLayer(file)
			if err != nil {
				return fmt.Errorf("failed to create layer: %v", err)
			}
			l.MediaType = "application/vnd.ollama.image.adapter"
			layers = append(layers, l)
Bruce MacDonald's avatar
Bruce MacDonald committed
374
375
376
377
378
379
380
381
382
383
384
385
		case "license":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
			mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

			layer, err := CreateLayer(strings.NewReader(c.Args))
			if err != nil {
				return err
			}

			layer.MediaType = mediaType
			layers = append(layers, layer)
		case "template", "system", "prompt":
386
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
387
			// remove the layer if one exists
388
389
			mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
			layers = removeLayerFromLayers(layers, mediaType)
390

391
			layer, err := CreateLayer(strings.NewReader(c.Args))
392
			if err != nil {
393
				return err
394
			}
395
396
397

			layer.MediaType = mediaType
			layers = append(layers, layer)
398
		default:
399
400
			// runtime parameters, build a list of args for each parameter to allow multiple values to be specified (ex: multiple stop tokens)
			params[c.Name] = append(params[c.Name], c.Args)
401
402
403
404
		}
	}

	// Create a single layer for the parameters
Michael Yang's avatar
Michael Yang committed
405
	if len(params) > 0 {
406
		fn(api.ProgressResponse{Status: "creating parameter layer"})
407
		layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params")
408
		formattedParams, err := formatParams(params)
409
410
411
		if err != nil {
			return fmt.Errorf("couldn't create params json: %v", err)
		}
412
413
414
415
416
417
418

		bts, err := json.Marshal(formattedParams)
		if err != nil {
			return err
		}

		l, err := CreateLayer(bytes.NewReader(bts))
419
420
421
422
423
		if err != nil {
			return fmt.Errorf("failed to create layer: %v", err)
		}
		l.MediaType = "application/vnd.ollama.image.params"
		layers = append(layers, l)
424
425

		// apply these parameters to the embedding options, in case embeddings need to be generated using this model
426
		embed.opts = formattedParams
427
428
	}

429
430
431
432
433
434
435
	// generate the embedding layers
	embeddingLayers, err := embeddingLayers(embed)
	if err != nil {
		return err
	}
	layers = append(layers, embeddingLayers...)

436
437
438
439
440
441
442
443
444
445
446
	digests, err := getLayerDigests(layers)
	if err != nil {
		return err
	}

	var manifestLayers []*Layer
	for _, l := range layers {
		manifestLayers = append(manifestLayers, &l.Layer)
	}

	// Create a layer for the config object
447
	fn(api.ProgressResponse{Status: "creating config layer"})
448
	cfg, err := createConfigLayer(config, digests)
449
450
451
452
453
454
455
456
457
458
459
	if err != nil {
		return err
	}
	layers = append(layers, cfg)

	err = SaveLayers(layers, fn, false)
	if err != nil {
		return err
	}

	// Create the manifest
460
	fn(api.ProgressResponse{Status: "writing manifest"})
461
462
463
464
465
	err = CreateManifest(name, cfg, manifestLayers)
	if err != nil {
		return err
	}

466
	fn(api.ProgressResponse{Status: "success"})
467
468
469
	return nil
}

470
471
type EmbeddingParams struct {
	model string
472
	opts  map[string]interface{}
473
474
475
476
477
478
479
480
	files []string // paths to files to embed
	fn    func(resp api.ProgressResponse)
}

// embeddingLayers loads the associated LLM and generates the embeddings to be stored from an input file
func embeddingLayers(e EmbeddingParams) ([]*LayerReader, error) {
	layers := []*LayerReader{}
	if len(e.files) > 0 {
481
482
483
484
485
		// check if the model is a file path or a model name
		model, err := GetModel(e.model)
		if err != nil {
			if !strings.Contains(err.Error(), "couldn't open file") {
				return nil, fmt.Errorf("unexpected error opening model to generate embeddings: %v", err)
486
			}
487
488
			// the model may be a file path, create a model from this file
			model = &Model{ModelPath: e.model}
489
490
		}

491
		if err := load(model, e.opts, defaultSessionDuration); err != nil {
492
493
494
			return nil, fmt.Errorf("load model to generate embeddings: %v", err)
		}

495
		// this will be used to check if we already have embeddings for a file
496
		modelInfo, err := os.Stat(model.ModelPath)
497
		if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
498
			return nil, fmt.Errorf("failed to get model file info: %v", err)
499
500
		}

Bruce MacDonald's avatar
Bruce MacDonald committed
501
502
503
		addedFiles := make(map[string]bool) // keep track of files that have already been added
		for _, filePattern := range e.files {
			matchingFiles, err := filepath.Glob(filePattern)
504
			if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
505
				return nil, fmt.Errorf("could not find files with pattern %s: %w", filePattern, err)
506
507
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
508
509
			for _, filePath := range matchingFiles {
				if addedFiles[filePath] {
510
511
					continue
				}
Bruce MacDonald's avatar
Bruce MacDonald committed
512
				addedFiles[filePath] = true
Bruce MacDonald's avatar
Bruce MacDonald committed
513
				// check if we already have embeddings for this file path
514
				layerIdentifier := fmt.Sprintf("%s:%s:%s:%d", filePath, e.model, modelInfo.ModTime().Format("2006-01-02 15:04:05"), modelInfo.Size())
Bruce MacDonald's avatar
Bruce MacDonald committed
515
516
517
518
519
520
				digest, _ := GetSHA256Digest(strings.NewReader(layerIdentifier))
				existing, err := existingFileEmbeddings(digest)
				if err != nil {
					return nil, fmt.Errorf("failed to check existing embeddings for file %s: %v", filePath, err)
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
521
522
				// TODO: check file type
				f, err := os.Open(filePath)
523
				if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
524
					return nil, fmt.Errorf("could not open embed file: %w", err)
525
				}
Bruce MacDonald's avatar
Bruce MacDonald committed
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
				scanner := bufio.NewScanner(f)
				scanner.Split(bufio.ScanLines)

				data := []string{}
				for scanner.Scan() {
					data = append(data, scanner.Text())
				}
				f.Close()

				// the digest of the file is set here so that the client knows a new operation is in progress
				fileDigest, _ := GetSHA256Digest(bytes.NewReader([]byte(filePath)))

				embeddings := []vector.Embedding{}
				for i, d := range data {
					if strings.TrimSpace(d) == "" {
						continue
					}
					e.fn(api.ProgressResponse{
						Status:    fmt.Sprintf("creating embeddings for file %s", filePath),
						Digest:    fileDigest,
						Total:     len(data) - 1,
						Completed: i,
					})
Bruce MacDonald's avatar
Bruce MacDonald committed
549
550
551
552
553
					if len(existing[d]) > 0 {
						// already have an embedding for this line
						embeddings = append(embeddings, vector.Embedding{Data: d, Vector: existing[d]})
						continue
					}
554
					embed, err := loaded.llm.Embedding(d)
Bruce MacDonald's avatar
Bruce MacDonald committed
555
					if err != nil {
556
557
						log.Printf("failed to generate embedding for '%s' line %d: %v", filePath, i+1, err)
						continue
Bruce MacDonald's avatar
Bruce MacDonald committed
558
559
					}
					embeddings = append(embeddings, vector.Embedding{Data: d, Vector: embed})
560
561
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
562
563
564
565
566
				b, err := json.Marshal(embeddings)
				if err != nil {
					return nil, fmt.Errorf("failed to encode embeddings: %w", err)
				}
				r := bytes.NewReader(b)
567

Bruce MacDonald's avatar
Bruce MacDonald committed
568
569
570
571
				layer := &LayerReader{
					Layer: Layer{
						MediaType: "application/vnd.ollama.image.embed",
						Digest:    digest,
Bruce MacDonald's avatar
Bruce MacDonald committed
572
						Size:      r.Len(),
Bruce MacDonald's avatar
Bruce MacDonald committed
573
574
575
					},
					Reader: r,
				}
576

Bruce MacDonald's avatar
Bruce MacDonald committed
577
578
				layers = append(layers, layer)
			}
579
580
581
582
583
		}
	}
	return layers, nil
}

Bruce MacDonald's avatar
Bruce MacDonald committed
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
// existingFileEmbeddings checks if we already have embeddings for a file and loads them into a look-up map
func existingFileEmbeddings(digest string) (map[string][]float64, error) {
	path, err := GetBlobsPath(digest)
	if err != nil {
		return nil, fmt.Errorf("embeddings blobs path: %w", err)
	}
	existingFileEmbeddings := make(map[string][]float64)
	if _, err := os.Stat(path); err == nil {
		// already have some embeddings for this file, load embeddings previously generated
		file, err := os.Open(path)
		if err != nil {
			return nil, fmt.Errorf("failed to open existing embedding file: %s", err)
		}
		defer file.Close()

		existing := []vector.Embedding{}
		if err = json.NewDecoder(file).Decode(&existing); err != nil {
			return nil, err
		}
		for _, e := range existing {
			existingFileEmbeddings[e.Data] = e.Vector
		}
	}
	return existingFileEmbeddings, nil
}

Michael Yang's avatar
Michael Yang committed
610
func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
611
612
613
614
615
616
617
618
619
620
	j := 0
	for _, l := range layers {
		if l.MediaType != mediaType {
			layers[j] = l
			j++
		}
	}
	return layers[:j]
}

621
func SaveLayers(layers []*LayerReader, fn func(resp api.ProgressResponse), force bool) error {
622
623
	// Write each of the layers to disk
	for _, layer := range layers {
Patrick Devine's avatar
Patrick Devine committed
624
		fp, err := GetBlobsPath(layer.Digest)
625
626
627
		if err != nil {
			return err
		}
628
629

		_, err = os.Stat(fp)
Bruce MacDonald's avatar
Bruce MacDonald committed
630
631
		// note: embed layers are always written since their digest doesnt indicate anything about the contents
		if os.IsNotExist(err) || force || layer.MediaType == "application/vnd.ollama.image.embed" {
632
633
			fn(api.ProgressResponse{Status: fmt.Sprintf("writing layer %s", layer.Digest)})

634
635
636
637
638
639
640
			out, err := os.Create(fp)
			if err != nil {
				log.Printf("couldn't create %s", fp)
				return err
			}
			defer out.Close()

Michael Yang's avatar
Michael Yang committed
641
			if _, err = io.Copy(out, layer.Reader); err != nil {
642
643
				return err
			}
Michael Yang's avatar
Michael Yang committed
644

645
		} else {
646
			fn(api.ProgressResponse{Status: fmt.Sprintf("using already created layer %s", layer.Digest)})
647
648
649
650
651
652
		}
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
653
func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error {
Patrick Devine's avatar
Patrick Devine committed
654
655
	mp := ParseModelPath(name)

656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
	manifest := ManifestV2{
		SchemaVersion: 2,
		MediaType:     "application/vnd.docker.distribution.manifest.v2+json",
		Config: Layer{
			MediaType: cfg.MediaType,
			Size:      cfg.Size,
			Digest:    cfg.Digest,
		},
		Layers: layers,
	}

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

Patrick Devine's avatar
Patrick Devine committed
672
	fp, err := mp.GetManifestPath(true)
673
674
675
	if err != nil {
		return err
	}
676
	return os.WriteFile(fp, manifestJSON, 0o644)
677
678
}

Michael Yang's avatar
Michael Yang committed
679
func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
Patrick Devine's avatar
Patrick Devine committed
680
	fp, err := GetBlobsPath(layer.Digest)
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
	if err != nil {
		return nil, err
	}

	file, err := os.Open(fp)
	if err != nil {
		return nil, fmt.Errorf("could not open blob: %w", err)
	}
	defer file.Close()

	newLayer, err := CreateLayer(file)
	if err != nil {
		return nil, err
	}
	newLayer.MediaType = layer.MediaType
	return newLayer, nil
}

699
700
// formatParams converts specified parameter options to their correct types
func formatParams(params map[string][]string) (map[string]interface{}, error) {
701
702
703
	opts := api.Options{}
	valueOpts := reflect.ValueOf(&opts).Elem() // names of the fields in the options struct
	typeOpts := reflect.TypeOf(opts)           // types of the fields in the options struct
Michael Yang's avatar
Michael Yang committed
704

705
	// build map of json struct tags to their types
Michael Yang's avatar
Michael Yang committed
706
707
708
709
710
711
712
713
	jsonOpts := make(map[string]reflect.StructField)
	for _, field := range reflect.VisibleFields(typeOpts) {
		jsonTag := strings.Split(field.Tag.Get("json"), ",")[0]
		if jsonTag != "" {
			jsonOpts[jsonTag] = field
		}
	}

714
	out := make(map[string]interface{})
Michael Yang's avatar
Michael Yang committed
715
	// iterate params and set values based on json struct tags
716
	for key, vals := range params {
Michael Yang's avatar
Michael Yang committed
717
718
719
720
721
		if opt, ok := jsonOpts[key]; ok {
			field := valueOpts.FieldByName(opt.Name)
			if field.IsValid() && field.CanSet() {
				switch field.Kind() {
				case reflect.Float32:
722
					floatVal, err := strconv.ParseFloat(vals[0], 32)
Michael Yang's avatar
Michael Yang committed
723
					if err != nil {
724
						return nil, fmt.Errorf("invalid float value %s", vals)
Michael Yang's avatar
Michael Yang committed
725
726
					}

727
					out[key] = floatVal
Michael Yang's avatar
Michael Yang committed
728
				case reflect.Int:
729
					intVal, err := strconv.ParseInt(vals[0], 10, 0)
Michael Yang's avatar
Michael Yang committed
730
					if err != nil {
731
						return nil, fmt.Errorf("invalid int value %s", vals)
Michael Yang's avatar
Michael Yang committed
732
733
					}

734
					out[key] = intVal
Michael Yang's avatar
Michael Yang committed
735
				case reflect.Bool:
736
					boolVal, err := strconv.ParseBool(vals[0])
Michael Yang's avatar
Michael Yang committed
737
					if err != nil {
738
						return nil, fmt.Errorf("invalid bool value %s", vals)
Michael Yang's avatar
Michael Yang committed
739
740
					}

741
					out[key] = boolVal
Michael Yang's avatar
Michael Yang committed
742
				case reflect.String:
743
					out[key] = vals[0]
744
				case reflect.Slice:
745
746
					// TODO: only string slices are supported right now
					out[key] = vals
Michael Yang's avatar
Michael Yang committed
747
748
749
750
751
752
753
				default:
					return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key)
				}
			}
		}
	}

754
	return out, nil
755
756
}

Michael Yang's avatar
Michael Yang committed
757
func getLayerDigests(layers []*LayerReader) ([]string, error) {
758
759
760
761
762
763
764
765
766
767
768
	var digests []string
	for _, l := range layers {
		if l.Digest == "" {
			return nil, fmt.Errorf("layer is missing a digest")
		}
		digests = append(digests, l.Digest)
	}
	return digests, nil
}

// CreateLayer creates a Layer object from a given file
Michael Yang's avatar
Michael Yang committed
769
770
func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
	digest, size := GetSHA256Digest(f)
771
	f.Seek(0, io.SeekStart)
772

Michael Yang's avatar
Michael Yang committed
773
	layer := &LayerReader{
774
775
776
777
778
		Layer: Layer{
			MediaType: "application/vnd.docker.image.rootfs.diff.tar",
			Digest:    digest,
			Size:      size,
		},
Michael Yang's avatar
Michael Yang committed
779
		Reader: f,
780
781
782
783
784
	}

	return layer, nil
}

Patrick Devine's avatar
Patrick Devine committed
785
786
787
788
789
790
791
792
793
794
795
func CopyModel(src, dest string) error {
	srcPath, err := ParseModelPath(src).GetManifestPath(false)
	if err != nil {
		return err
	}
	destPath, err := ParseModelPath(dest).GetManifestPath(true)
	if err != nil {
		return err
	}

	// copy the file
Michael Yang's avatar
Michael Yang committed
796
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
797
798
799
800
801
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
802
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
803
804
805
806
807
808
809
810
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

811
func DeleteModel(name string) error {
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
	mp := ParseModelPath(name)

	manifest, err := GetManifest(mp)
	if err != nil {
		return err
	}
	deleteMap := make(map[string]bool)
	for _, layer := range manifest.Layers {
		deleteMap[layer.Digest] = true
	}
	deleteMap[manifest.Config.Digest] = true

	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
	err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() {
			path := path[len(fp)+1:]
			slashIndex := strings.LastIndex(path, "/")
			if slashIndex == -1 {
				return nil
			}
			tag := path[:slashIndex] + ":" + path[slashIndex+1:]
			fmp := ParseModelPath(tag)

			// skip the manifest we're trying to delete
			if mp.GetFullTagname() == fmp.GetFullTagname() {
				return nil
			}

			// save (i.e. delete from the deleteMap) any files used in other manifests
			manifest, err := GetManifest(fmp)
			if err != nil {
				log.Printf("skipping file: %s", fp)
				return nil
			}
			for _, layer := range manifest.Layers {
				delete(deleteMap, layer.Digest)
			}
			delete(deleteMap, manifest.Config.Digest)
		}
		return nil
	})
Michael Yang's avatar
Michael Yang committed
859
860
861
	if err != nil {
		return err
	}
862
863
864
865

	// only delete the files which are still in the deleteMap
	for k, v := range deleteMap {
		if v {
866
			fp, err := GetBlobsPath(k)
867
			if err != nil {
868
869
870
871
872
				log.Printf("couldn't get file path for '%s': %v", k, err)
				continue
			}
			if err := os.Remove(fp); err != nil {
				log.Printf("couldn't remove file '%s': %v", fp, err)
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
				continue
			}
		}
	}

	fp, err = mp.GetManifestPath(false)
	if err != nil {
		return err
	}
	err = os.Remove(fp)
	if err != nil {
		log.Printf("couldn't remove manifest file '%s': %v", fp, err)
		return err
	}

	return nil
}

891
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
Patrick Devine's avatar
Patrick Devine committed
892
893
	mp := ParseModelPath(name)

894
895
	fn(api.ProgressResponse{Status: "retrieving manifest"})

Patrick Devine's avatar
Patrick Devine committed
896
	manifest, err := GetManifest(mp)
897
	if err != nil {
898
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
899
900
901
902
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
903
	layers = append(layers, manifest.Layers...)
904
905
906
	layers = append(layers, &manifest.Config)

	for _, layer := range layers {
907
		exists, err := checkBlobExistence(ctx, mp, layer.Digest, regOpts)
908
909
910
911
912
		if err != nil {
			return err
		}

		if exists {
913
914
915
			fn(api.ProgressResponse{
				Status:    "using existing layer",
				Digest:    layer.Digest,
916
917
				Total:     layer.Size,
				Completed: layer.Size,
918
			})
919
			log.Printf("Layer %s already exists", layer.Digest)
920
921
922
			continue
		}

923
		fn(api.ProgressResponse{
924
925
926
			Status: "starting upload",
			Digest: layer.Digest,
			Total:  layer.Size,
927
		})
928

929
		location, err := startUpload(ctx, mp, regOpts)
930
931
932
933
934
		if err != nil {
			log.Printf("couldn't start upload: %v", err)
			return err
		}

935
		err = uploadBlobChunked(ctx, mp, location, layer, regOpts, fn)
936
937
938
939
		if err != nil {
			log.Printf("error uploading blob: %v", err)
			return err
		}
940
941
	}

942
	fn(api.ProgressResponse{Status: "pushing manifest"})
943
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
944
945
946
947
948
949
950
951
952
	headers := map[string]string{
		"Content-Type": "application/vnd.docker.distribution.manifest.v2+json",
	}

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

953
	resp, err := makeRequest(ctx, "PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
954
955
956
957
958
959
960
961
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
	if resp.StatusCode != http.StatusCreated {
		body, _ := io.ReadAll(resp.Body)
962
		return fmt.Errorf("on push registry responded with code %d: %v", resp.StatusCode, string(body))
963
964
	}

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

	return nil
}

970
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
Patrick Devine's avatar
Patrick Devine committed
971
	mp := ParseModelPath(name)
972

973
	fn(api.ProgressResponse{Status: "pulling manifest"})
974

975
	manifest, err := pullModelManifest(ctx, mp, regOpts)
976
	if err != nil {
977
		return fmt.Errorf("pull model manifest: %s", err)
978
979
980
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
981
	layers = append(layers, manifest.Layers...)
982
983
984
	layers = append(layers, &manifest.Config)

	for _, layer := range layers {
985
		if err := downloadBlob(ctx, mp, layer.Digest, regOpts, fn); err != nil {
986
987
988
989
			return err
		}
	}

Michael Yang's avatar
Michael Yang committed
990
991
992
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
993
994
995
996
997
998
999
1000
1001
1002
1003
			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
					log.Printf("couldn't remove file with digest mismatch '%s': %v", fp, err)
				}
			}
Michael Yang's avatar
Michael Yang committed
1004
1005
1006
1007
			return err
		}
	}

1008
	fn(api.ProgressResponse{Status: "writing manifest"})
1009

1010
	manifestJSON, err := json.Marshal(manifest)
1011
1012
1013
1014
	if err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
1015
	fp, err := mp.GetManifestPath(true)
1016
1017
1018
1019
	if err != nil {
		return err
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
1020
	err = os.WriteFile(fp, manifestJSON, 0o644)
1021
1022
1023
1024
1025
	if err != nil {
		log.Printf("couldn't write to %s", fp)
		return err
	}

1026
	fn(api.ProgressResponse{Status: "success"})
1027
1028
1029
1030

	return nil
}

1031
func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
1032
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
1033
1034
1035
1036
	headers := map[string]string{
		"Accept": "application/vnd.docker.distribution.manifest.v2+json",
	}

1037
	resp, err := makeRequest(ctx, "GET", url, headers, nil, regOpts)
1038
1039
1040
1041
1042
1043
1044
1045
	if err != nil {
		log.Printf("couldn't get manifest: %v", err)
		return nil, err
	}
	defer resp.Body.Close()

	// Check for success: For a successful upload, the Docker registry will respond with a 201 Created
	if resp.StatusCode != http.StatusOK {
1046
		if resp.StatusCode == http.StatusNotFound {
Bruce MacDonald's avatar
Bruce MacDonald committed
1047
			return nil, fmt.Errorf("model not found")
1048
		}
1049
		body, _ := io.ReadAll(resp.Body)
1050
		return nil, fmt.Errorf("on pull registry responded with code %d: %s", resp.StatusCode, body)
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
	}

	var m *ManifestV2
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

	return m, err
}

1061
1062
1063
1064
func createConfigLayer(config ConfigV2, layers []string) (*LayerReader, error) {
	config.RootFS = RootFS{
		Type:    "layers",
		DiffIDs: layers,
1065
1066
1067
1068
1069
1070
1071
	}

	configJSON, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

1072
	digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
1073

Michael Yang's avatar
Michael Yang committed
1074
	layer := &LayerReader{
1075
1076
1077
1078
1079
		Layer: Layer{
			MediaType: "application/vnd.docker.container.image.v1+json",
			Digest:    digest,
			Size:      size,
		},
1080
		Reader: bytes.NewBuffer(configJSON),
1081
1082
1083
1084
1085
	}
	return layer, nil
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
1086
1087
1088
1089
1090
1091
1092
1093
func GetSHA256Digest(r io.Reader) (string, int) {
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

	return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
1094
1095
}

1096
func startUpload(ctx context.Context, mp ModelPath, regOpts *RegistryOptions) (string, error) {
1097
	url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
1098

1099
	resp, err := makeRequest(ctx, "POST", url, nil, nil, regOpts)
1100
1101
1102
1103
1104
1105
1106
1107
1108
	if err != nil {
		log.Printf("couldn't start upload: %v", err)
		return "", err
	}
	defer resp.Body.Close()

	// Check for success
	if resp.StatusCode != http.StatusAccepted {
		body, _ := io.ReadAll(resp.Body)
1109
		return "", fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
	}

	// Extract UUID location from header
	location := resp.Header.Get("Location")
	if location == "" {
		return "", fmt.Errorf("location header is missing in response")
	}

	return location, nil
}

// Function to check if a blob already exists in the Docker registry
1122
func checkBlobExistence(ctx context.Context, mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
1123
	url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
1124

1125
	resp, err := makeRequest(ctx, "HEAD", url, nil, nil, regOpts)
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
	if err != nil {
		log.Printf("couldn't check for blob: %v", err)
		return false, err
	}
	defer resp.Body.Close()

	// Check for success: If the blob exists, the Docker registry will respond with a 200 OK
	return resp.StatusCode == http.StatusOK, nil
}

1136
func uploadBlobChunked(ctx context.Context, mp ModelPath, url string, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
1137
1138
1139
1140
	// TODO allow resumability
	// TODO allow canceling uploads via DELETE
	// TODO allow cross repo blob mount

Patrick Devine's avatar
Patrick Devine committed
1141
	fp, err := GetBlobsPath(layer.Digest)
1142
1143
1144
1145
	if err != nil {
		return err
	}

1146
1147
1148
1149
	f, err := os.Open(fp)
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
1150
	defer f.Close()
1151

1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
	totalUploaded := 0

	r, w := io.Pipe()
	defer r.Close()

	go func() {
		defer w.Close()
		for {
			n, err := io.CopyN(w, f, 1024*1024)
			if err != nil && !errors.Is(err, io.EOF) {
				fn(api.ProgressResponse{
					Status:    fmt.Sprintf("error copying pipe: %v", err),
					Digest:    layer.Digest,
					Total:     layer.Size,
					Completed: totalUploaded,
				})
				return
			}
1170

1171
			totalUploaded += int(n)
1172
1173

			fn(api.ProgressResponse{
1174
				Status:    fmt.Sprintf("uploading %s", layer.Digest),
1175
				Digest:    layer.Digest,
1176
1177
				Total:     layer.Size,
				Completed: totalUploaded,
1178
			})
1179
1180
1181
1182

			if totalUploaded >= layer.Size {
				return
			}
1183
		}
1184
	}()
1185

1186
	url = fmt.Sprintf("%s&digest=%s", url, layer.Digest)
1187

1188
1189
1190
1191
	headers := make(map[string]string)
	headers["Content-Type"] = "application/octet-stream"
	headers["Content-Range"] = fmt.Sprintf("0-%d", layer.Size-1)
	headers["Content-Length"] = strconv.Itoa(int(layer.Size))
1192

1193
	// finish the upload
1194
	resp, err := makeRequest(ctx, "PUT", url, headers, r, regOpts)
1195
1196
1197
1198
1199
1200
1201
1202
1203
	if err != nil {
		log.Printf("couldn't finish upload: %v", err)
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusCreated {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("on finish upload registry responded with code %d: %v", resp.StatusCode, string(body))
1204
	}
1205
1206
1207
	return nil
}

1208
1209
1210
1211
1212
1213
func makeRequest(ctx context.Context, method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
	retryCtx := ctx.Value("retries")
	var retries int
	var ok bool
	if retries, ok = retryCtx.(int); ok {
		if retries > MaxRetries {
Bruce MacDonald's avatar
Bruce MacDonald committed
1214
			return nil, fmt.Errorf("maximum retries hit; are you sure you have access to this resource?")
1215
1216
1217
		}
	}

1218
1219
1220
1221
1222
1223
1224
1225
	if !strings.HasPrefix(url, "http") {
		if regOpts.Insecure {
			url = "http://" + url
		} else {
			url = "https://" + url
		}
	}

Patrick Devine's avatar
Patrick Devine committed
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
	// make a copy of the body in case we need to try the call to makeRequest again
	var buf bytes.Buffer
	if body != nil {
		_, err := io.Copy(&buf, body)
		if err != nil {
			return nil, err
		}
	}

	bodyCopy := bytes.NewReader(buf.Bytes())

	req, err := http.NewRequest(method, url, bodyCopy)
1238
1239
1240
1241
	if err != nil {
		return nil, err
	}

Patrick Devine's avatar
Patrick Devine committed
1242
1243
1244
1245
	if regOpts.Token != "" {
		req.Header.Set("Authorization", "Bearer "+regOpts.Token)
	} else if regOpts.Username != "" && regOpts.Password != "" {
		req.SetBasicAuth(regOpts.Username, regOpts.Password)
1246
1247
	}

Patrick Devine's avatar
Patrick Devine committed
1248
1249
	for k, v := range headers {
		req.Header.Set(k, v)
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
	}

	client := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= 10 {
				return fmt.Errorf("too many redirects")
			}
			log.Printf("redirected to: %s\n", req.URL)
			return nil
		},
	}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

Patrick Devine's avatar
Patrick Devine committed
1266
1267
1268
1269
	// if the request is unauthenticated, try to authenticate and make the request again
	if resp.StatusCode == http.StatusUnauthorized {
		auth := resp.Header.Get("Www-Authenticate")
		authRedir := ParseAuthRedirectString(string(auth))
1270
		token, err := getAuthToken(ctx, authRedir, regOpts)
Patrick Devine's avatar
Patrick Devine committed
1271
1272
1273
1274
1275
		if err != nil {
			return nil, err
		}
		regOpts.Token = token
		bodyCopy = bytes.NewReader(buf.Bytes())
1276
1277
		ctx = context.WithValue(ctx, "retries", retries+1)
		return makeRequest(ctx, method, url, headers, bodyCopy, regOpts)
Patrick Devine's avatar
Patrick Devine committed
1278
1279
	}

1280
1281
	return resp, nil
}
Michael Yang's avatar
Michael Yang committed
1282

Patrick Devine's avatar
Patrick Devine committed
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
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]
}

func ParseAuthRedirectString(authStr string) AuthRedirect {
	authStr = strings.TrimPrefix(authStr, "Bearer ")

	return AuthRedirect{
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

1316
1317
var errDigestMismatch = fmt.Errorf("digest mismatch, file must be downloaded again")

Michael Yang's avatar
Michael Yang committed
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
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 {
1332
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1333
1334
1335
1336
	}

	return nil
}