images.go 29.8 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/llama"
23
	"github.com/jmorganca/ollama/parser"
24
	"github.com/jmorganca/ollama/vector"
25
26
)

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

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

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

	tmpl, err := template.New("").Parse(t)
51
52
53
54
55
	if err != nil {
		return "", err
	}

	var vars struct {
Michael Yang's avatar
Michael Yang committed
56
		First  bool
57
58
		System string
		Prompt string
59
		Embed  string
60
61
62

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

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

71
72
73
74
	if request.System != "" {
		vars.System = request.System
	}

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

	return sb.String(), nil
}

83
84
85
86
87
88
89
90
91
92
93
94
95
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
96
type LayerReader struct {
97
	Layer
Michael Yang's avatar
Michael Yang committed
98
	io.Reader
99
100
101
102
103
104
105
106
107
108
109
110
111
}

type ConfigV2 struct {
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
	RootFS       RootFS `json:"rootfs"`
}

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

Patrick Devine's avatar
Patrick Devine committed
112
113
114
115
116
117
118
119
120
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
121
122
func GetManifest(mp ModelPath) (*ManifestV2, error) {
	fp, err := mp.GetManifestPath(false)
123
124
125
	if err != nil {
		return nil, err
	}
126

127
128
	if _, err = os.Stat(fp); err != nil {
		return nil, err
129
130
131
132
	}

	var manifest *ManifestV2

133
	bts, err := os.ReadFile(fp)
134
135
136
137
	if err != nil {
		return nil, fmt.Errorf("couldn't open file '%s'", fp)
	}

138
	if err := json.Unmarshal(bts, &manifest); err != nil {
139
140
141
142
143
144
145
		return nil, err
	}

	return manifest, nil
}

func GetModel(name string) (*Model, error) {
Patrick Devine's avatar
Patrick Devine committed
146
147
148
	mp := ParseModelPath(name)

	manifest, err := GetManifest(mp)
149
150
151
152
153
	if err != nil {
		return nil, err
	}

	model := &Model{
154
		Name:   mp.GetFullTagname(),
Jeffrey Morgan's avatar
Jeffrey Morgan committed
155
		Digest: manifest.Config.Digest,
156
157
158
	}

	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
159
		filename, err := GetBlobsPath(layer.Digest)
160
161
162
163
		if err != nil {
			return nil, err
		}

164
165
166
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
167
168
169
170
171
172
173
174
175
176
		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
			}
177
178
179
180
181
182
183
184
185
		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)
186
187
188
			if err != nil {
				return nil, err
			}
189
190

			model.System = string(bts)
191
192
193
194
195
196
197
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
198
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
199
200
201
202
203
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
204

205
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
206
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
207
208
				return nil, err
			}
209
210
211
212
213
214
		}
	}

	return model, nil
}

215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
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
}

235
func CreateModel(ctx context.Context, name string, path string, fn func(resp api.ProgressResponse)) error {
236
237
	mf, err := os.Open(path)
	if err != nil {
238
		fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't open modelfile '%s'", path)})
239
		return fmt.Errorf("failed to open file: %w", err)
240
	}
241
	defer mf.Close()
242

243
	fn(api.ProgressResponse{Status: "parsing modelfile"})
244
245
246
247
248
	commands, err := parser.Parse(mf)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
249
	var layers []*LayerReader
250
	params := make(map[string][]string)
251
	embed := EmbeddingParams{fn: fn, opts: api.DefaultOptions()}
252
	for _, c := range commands {
253
		log.Printf("[%s] - %s\n", c.Name, c.Args)
254
255
		switch c.Name {
		case "model":
256
			fn(api.ProgressResponse{Status: "looking for model"})
257
			embed.model = c.Args
258
			mf, err := GetManifest(ParseModelPath(c.Args))
259
			if err != nil {
260
261
262
				modelFile, err := filenameWithPath(path, c.Args)
				if err != nil {
					return err
263
				}
264
				if _, err := os.Stat(modelFile); err != nil {
265
266
267
					// the model file does not exist, try pulling it
					if errors.Is(err, os.ErrNotExist) {
						fn(api.ProgressResponse{Status: "pulling model file"})
268
						if err := PullModel(ctx, c.Args, &RegistryOptions{}, fn); err != nil {
269
270
							return err
						}
271
						mf, err = GetManifest(ParseModelPath(c.Args))
272
273
274
275
276
277
278
279
280
						if err != nil {
							return fmt.Errorf("failed to open file after pull: %v", err)
						}
					} else {
						return err
					}
				} else {
					// create a model from this specified file
					fn(api.ProgressResponse{Status: "creating model layer"})
281
					file, err := os.Open(modelFile)
282
283
284
285
286
287
288
289
290
291
292
					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.model"
					layers = append(layers, l)
293
				}
294
295
			}
			if mf != nil {
296
297
298
299
300
301
302
303
304
				log.Printf("manifest = %#v", mf)
				for _, l := range mf.Layers {
					newLayer, err := GetLayerWithBufferFromLayer(l)
					if err != nil {
						return err
					}
					layers = append(layers, newLayer)
				}
			}
305
306
		case "embed":
			embedFilePath, err := filenameWithPath(path, c.Args)
Michael Yang's avatar
Michael Yang committed
307
308
309
			if err != nil {
				return err
			}
310
			embed.files = append(embed.files, embedFilePath)
Bruce MacDonald's avatar
Bruce MacDonald committed
311
312
313
314
315
316
317
318
319
320
321
322
		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":
323
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
324
			// remove the prompt layer if one exists
325
326
			mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
			layers = removeLayerFromLayers(layers, mediaType)
327

328
			layer, err := CreateLayer(strings.NewReader(c.Args))
329
			if err != nil {
330
				return err
331
			}
332
333
334

			layer.MediaType = mediaType
			layers = append(layers, layer)
335
		default:
336
337
			// 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)
338
339
340
341
		}
	}

	// Create a single layer for the parameters
Michael Yang's avatar
Michael Yang committed
342
	if len(params) > 0 {
343
		fn(api.ProgressResponse{Status: "creating parameter layer"})
344
		layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params")
345
		formattedParams, err := formatParams(params)
346
347
348
		if err != nil {
			return fmt.Errorf("couldn't create params json: %v", err)
		}
349
350
351
352
353
354
355

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

		l, err := CreateLayer(bytes.NewReader(bts))
356
357
358
359
360
		if err != nil {
			return fmt.Errorf("failed to create layer: %v", err)
		}
		l.MediaType = "application/vnd.ollama.image.params"
		layers = append(layers, l)
361
362
363
364

		// apply these parameters to the embedding options, in case embeddings need to be generated using this model
		embed.opts = api.DefaultOptions()
		embed.opts.FromMap(formattedParams)
365
366
	}

367
368
369
370
371
372
373
	// generate the embedding layers
	embeddingLayers, err := embeddingLayers(embed)
	if err != nil {
		return err
	}
	layers = append(layers, embeddingLayers...)

374
375
376
377
378
379
380
381
382
383
384
	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
385
	fn(api.ProgressResponse{Status: "creating config layer"})
386
387
388
389
390
391
392
393
394
395
396
397
	cfg, err := createConfigLayer(digests)
	if err != nil {
		return err
	}
	layers = append(layers, cfg)

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

	// Create the manifest
398
	fn(api.ProgressResponse{Status: "writing manifest"})
399
400
401
402
403
	err = CreateManifest(name, cfg, manifestLayers)
	if err != nil {
		return err
	}

404
	fn(api.ProgressResponse{Status: "success"})
405
406
407
	return nil
}

408
409
410
411
412
413
414
415
416
417
418
type EmbeddingParams struct {
	model string
	opts  api.Options
	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 {
419
420
421
422
423
424
425
426
427
428
429
		if _, err := os.Stat(e.model); err != nil {
			if os.IsNotExist(err) {
				// this is a model name rather than the file
				model, err := GetModel(e.model)
				if err != nil {
					return nil, fmt.Errorf("failed to get model to generate embeddings: %v", err)
				}
				e.model = model.ModelPath
			} else {
				return nil, fmt.Errorf("failed to get model file to generate embeddings: %v", err)
			}
430
431
432
		}

		e.opts.EmbeddingOnly = true
433
		llm, err := llama.New(e.model, e.opts)
434
435
436
		if err != nil {
			return nil, fmt.Errorf("load model to generate embeddings: %v", err)
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
437
438
439
440
441
		defer func() {
			if llm != nil {
				llm.Close()
			}
		}()
442

Bruce MacDonald's avatar
Bruce MacDonald committed
443
444
445
		addedFiles := make(map[string]bool) // keep track of files that have already been added
		for _, filePattern := range e.files {
			matchingFiles, err := filepath.Glob(filePattern)
446
			if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
447
				return nil, fmt.Errorf("could not find files with pattern %s: %w", filePattern, err)
448
449
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
450
451
			for _, filePath := range matchingFiles {
				if addedFiles[filePath] {
452
453
					continue
				}
Bruce MacDonald's avatar
Bruce MacDonald committed
454
455
456
				addedFiles[filePath] = true
				// TODO: check file type
				f, err := os.Open(filePath)
457
				if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
458
					return nil, fmt.Errorf("could not open embed file: %w", err)
459
				}
Bruce MacDonald's avatar
Bruce MacDonald committed
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
				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,
					})
					embed, err := llm.Embedding(d)
					if err != nil {
485
486
						log.Printf("failed to generate embedding for '%s' line %d: %v", filePath, i+1, err)
						continue
Bruce MacDonald's avatar
Bruce MacDonald committed
487
488
					}
					embeddings = append(embeddings, vector.Embedding{Data: d, Vector: embed})
489
490
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
491
492
493
494
495
				b, err := json.Marshal(embeddings)
				if err != nil {
					return nil, fmt.Errorf("failed to encode embeddings: %w", err)
				}
				r := bytes.NewReader(b)
496

Bruce MacDonald's avatar
Bruce MacDonald committed
497
498
499
500
501
				digest, size := GetSHA256Digest(r)
				// Reset the position of the reader after calculating the digest
				if _, err := r.Seek(0, io.SeekStart); err != nil {
					return nil, fmt.Errorf("could not reset embed reader: %w", err)
				}
502

Bruce MacDonald's avatar
Bruce MacDonald committed
503
504
505
506
507
508
509
510
				layer := &LayerReader{
					Layer: Layer{
						MediaType: "application/vnd.ollama.image.embed",
						Digest:    digest,
						Size:      size,
					},
					Reader: r,
				}
511

Bruce MacDonald's avatar
Bruce MacDonald committed
512
513
				layers = append(layers, layer)
			}
514
515
516
517
518
		}
	}
	return layers, nil
}

Michael Yang's avatar
Michael Yang committed
519
func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
520
521
522
523
524
525
526
527
528
529
	j := 0
	for _, l := range layers {
		if l.MediaType != mediaType {
			layers[j] = l
			j++
		}
	}
	return layers[:j]
}

530
func SaveLayers(layers []*LayerReader, fn func(resp api.ProgressResponse), force bool) error {
531
532
	// Write each of the layers to disk
	for _, layer := range layers {
Patrick Devine's avatar
Patrick Devine committed
533
		fp, err := GetBlobsPath(layer.Digest)
534
535
536
		if err != nil {
			return err
		}
537
538
539

		_, err = os.Stat(fp)
		if os.IsNotExist(err) || force {
540
541
			fn(api.ProgressResponse{Status: fmt.Sprintf("writing layer %s", layer.Digest)})

542
543
544
545
546
547
548
			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
549
			if _, err = io.Copy(out, layer.Reader); err != nil {
550
551
				return err
			}
Michael Yang's avatar
Michael Yang committed
552

553
		} else {
554
			fn(api.ProgressResponse{Status: fmt.Sprintf("using already created layer %s", layer.Digest)})
555
556
557
558
559
560
		}
	}

	return nil
}

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

564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
	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
580
	fp, err := mp.GetManifestPath(true)
581
582
583
	if err != nil {
		return err
	}
584
	return os.WriteFile(fp, manifestJSON, 0o644)
585
586
}

Michael Yang's avatar
Michael Yang committed
587
func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
Patrick Devine's avatar
Patrick Devine committed
588
	fp, err := GetBlobsPath(layer.Digest)
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
	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
}

607
608
// formatParams converts specified parameter options to their correct types
func formatParams(params map[string][]string) (map[string]interface{}, error) {
609
610
611
	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
612

613
	// build map of json struct tags to their types
Michael Yang's avatar
Michael Yang committed
614
615
616
617
618
619
620
621
	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
		}
	}

622
	out := make(map[string]interface{})
Michael Yang's avatar
Michael Yang committed
623
	// iterate params and set values based on json struct tags
624
	for key, vals := range params {
Michael Yang's avatar
Michael Yang committed
625
626
627
628
629
		if opt, ok := jsonOpts[key]; ok {
			field := valueOpts.FieldByName(opt.Name)
			if field.IsValid() && field.CanSet() {
				switch field.Kind() {
				case reflect.Float32:
630
					floatVal, err := strconv.ParseFloat(vals[0], 32)
Michael Yang's avatar
Michael Yang committed
631
					if err != nil {
632
						return nil, fmt.Errorf("invalid float value %s", vals)
Michael Yang's avatar
Michael Yang committed
633
634
					}

635
					out[key] = floatVal
Michael Yang's avatar
Michael Yang committed
636
				case reflect.Int:
637
					intVal, err := strconv.ParseInt(vals[0], 10, 0)
Michael Yang's avatar
Michael Yang committed
638
					if err != nil {
639
						return nil, fmt.Errorf("invalid int value %s", vals)
Michael Yang's avatar
Michael Yang committed
640
641
					}

642
					out[key] = intVal
Michael Yang's avatar
Michael Yang committed
643
				case reflect.Bool:
644
					boolVal, err := strconv.ParseBool(vals[0])
Michael Yang's avatar
Michael Yang committed
645
					if err != nil {
646
						return nil, fmt.Errorf("invalid bool value %s", vals)
Michael Yang's avatar
Michael Yang committed
647
648
					}

649
					out[key] = boolVal
Michael Yang's avatar
Michael Yang committed
650
				case reflect.String:
651
					out[key] = vals[0]
652
				case reflect.Slice:
653
654
					// TODO: only string slices are supported right now
					out[key] = vals
Michael Yang's avatar
Michael Yang committed
655
656
657
658
659
660
661
				default:
					return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key)
				}
			}
		}
	}

662
	return out, nil
663
664
}

Michael Yang's avatar
Michael Yang committed
665
func getLayerDigests(layers []*LayerReader) ([]string, error) {
666
667
668
669
670
671
672
673
674
675
676
	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
677
678
679
func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
	digest, size := GetSHA256Digest(f)
	f.Seek(0, 0)
680

Michael Yang's avatar
Michael Yang committed
681
	layer := &LayerReader{
682
683
684
685
686
		Layer: Layer{
			MediaType: "application/vnd.docker.image.rootfs.diff.tar",
			Digest:    digest,
			Size:      size,
		},
Michael Yang's avatar
Michael Yang committed
687
		Reader: f,
688
689
690
691
692
	}

	return layer, nil
}

Patrick Devine's avatar
Patrick Devine committed
693
694
695
696
697
698
699
700
701
702
703
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
704
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
705
706
707
708
709
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
710
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
711
712
713
714
715
716
717
718
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

719
func DeleteModel(name string) error {
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
	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
767
768
769
	if err != nil {
		return err
	}
770

Michael Yang's avatar
Michael Yang committed
771
772
773
774
	if err != nil {
		return err
	}

775
776
777
	// only delete the files which are still in the deleteMap
	for k, v := range deleteMap {
		if v {
778
			fp, err := GetBlobsPath(k)
779
			if err != nil {
780
781
782
783
784
				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)
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
				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
}

803
func PushModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
Patrick Devine's avatar
Patrick Devine committed
804
805
	mp := ParseModelPath(name)

806
807
	fn(api.ProgressResponse{Status: "retrieving manifest"})

Patrick Devine's avatar
Patrick Devine committed
808
	manifest, err := GetManifest(mp)
809
	if err != nil {
810
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
811
812
813
814
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
815
	layers = append(layers, manifest.Layers...)
816
817
818
	layers = append(layers, &manifest.Config)

	for _, layer := range layers {
819
		exists, err := checkBlobExistence(mp, layer.Digest, regOpts)
820
821
822
823
824
		if err != nil {
			return err
		}

		if exists {
825
826
827
			fn(api.ProgressResponse{
				Status:    "using existing layer",
				Digest:    layer.Digest,
828
829
				Total:     layer.Size,
				Completed: layer.Size,
830
			})
831
			log.Printf("Layer %s already exists", layer.Digest)
832
833
834
			continue
		}

835
		fn(api.ProgressResponse{
836
837
838
			Status: "starting upload",
			Digest: layer.Digest,
			Total:  layer.Size,
839
		})
840

841
		location, err := startUpload(mp, regOpts)
842
843
844
845
846
		if err != nil {
			log.Printf("couldn't start upload: %v", err)
			return err
		}

847
		err = uploadBlobChunked(mp, location, layer, regOpts, fn)
848
849
850
851
		if err != nil {
			log.Printf("error uploading blob: %v", err)
			return err
		}
852
853
	}

854
	fn(api.ProgressResponse{Status: "pushing manifest"})
855
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
856
857
858
859
860
861
862
863
864
	headers := map[string]string{
		"Content-Type": "application/vnd.docker.distribution.manifest.v2+json",
	}

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

865
	resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
866
867
868
869
870
871
872
873
	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)
874
		return fmt.Errorf("on push registry responded with code %d: %v", resp.StatusCode, string(body))
875
876
	}

877
	fn(api.ProgressResponse{Status: "success"})
878
879
880
881

	return nil
}

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

885
	fn(api.ProgressResponse{Status: "pulling manifest"})
886

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

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
893
	layers = append(layers, manifest.Layers...)
894
895
896
	layers = append(layers, &manifest.Config)

	for _, layer := range layers {
897
		if err := downloadBlob(ctx, mp, layer.Digest, regOpts, fn); err != nil {
898
899
900
901
			return err
		}
	}

Michael Yang's avatar
Michael Yang committed
902
903
904
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
905
906
907
908
909
910
911
912
913
914
915
			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
916
917
918
919
			return err
		}
	}

920
	fn(api.ProgressResponse{Status: "writing manifest"})
921

922
	manifestJSON, err := json.Marshal(manifest)
923
924
925
926
	if err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
927
	fp, err := mp.GetManifestPath(true)
928
929
930
931
	if err != nil {
		return err
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
932
	err = os.WriteFile(fp, manifestJSON, 0o644)
933
934
935
936
937
	if err != nil {
		log.Printf("couldn't write to %s", fp)
		return err
	}

938
	fn(api.ProgressResponse{Status: "success"})
939
940
941
942

	return nil
}

943
944
func pullModelManifest(mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
	url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
945
946
947
948
	headers := map[string]string{
		"Accept": "application/vnd.docker.distribution.manifest.v2+json",
	}

949
	resp, err := makeRequest("GET", url, headers, nil, regOpts)
950
951
952
953
954
955
956
957
	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 {
958
		if resp.StatusCode == http.StatusNotFound {
Bruce MacDonald's avatar
Bruce MacDonald committed
959
			return nil, fmt.Errorf("model not found")
960
		}
961
		body, _ := io.ReadAll(resp.Body)
962
		return nil, fmt.Errorf("on pull registry responded with code %d: %s", resp.StatusCode, body)
963
964
965
966
967
968
969
970
971
972
	}

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

	return m, err
}

Michael Yang's avatar
Michael Yang committed
973
func createConfigLayer(layers []string) (*LayerReader, error) {
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
	// TODO change architecture and OS
	config := ConfigV2{
		Architecture: "arm64",
		OS:           "linux",
		RootFS: RootFS{
			Type:    "layers",
			DiffIDs: layers,
		},
	}

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

989
	digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
990

Michael Yang's avatar
Michael Yang committed
991
	layer := &LayerReader{
992
993
994
995
996
		Layer: Layer{
			MediaType: "application/vnd.docker.container.image.v1+json",
			Digest:    digest,
			Size:      size,
		},
997
		Reader: bytes.NewBuffer(configJSON),
998
999
1000
1001
1002
	}
	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
1003
1004
1005
1006
1007
1008
1009
1010
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)
1011
1012
}

1013
1014
func startUpload(mp ModelPath, regOpts *RegistryOptions) (string, error) {
	url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
1015

1016
	resp, err := makeRequest("POST", url, nil, nil, regOpts)
1017
1018
1019
1020
1021
1022
1023
1024
1025
	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)
1026
		return "", fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
	}

	// 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
1039
1040
func checkBlobExistence(mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
	url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
1041

1042
	resp, err := makeRequest("HEAD", url, nil, nil, regOpts)
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
	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
}

Michael Yang's avatar
Michael Yang committed
1053
func uploadBlobChunked(mp ModelPath, url string, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
1054
1055
1056
1057
	// TODO allow resumability
	// TODO allow canceling uploads via DELETE
	// TODO allow cross repo blob mount

Patrick Devine's avatar
Patrick Devine committed
1058
	fp, err := GetBlobsPath(layer.Digest)
1059
1060
1061
1062
	if err != nil {
		return err
	}

1063
1064
1065
1066
1067
	f, err := os.Open(fp)
	if err != nil {
		return err
	}

1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
	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
			}
1086

1087
			totalUploaded += int(n)
1088
1089

			fn(api.ProgressResponse{
1090
				Status:    fmt.Sprintf("uploading %s", layer.Digest),
1091
				Digest:    layer.Digest,
1092
1093
				Total:     layer.Size,
				Completed: totalUploaded,
1094
			})
1095
1096
1097
1098

			if totalUploaded >= layer.Size {
				return
			}
1099
		}
1100
	}()
1101

1102
	url = fmt.Sprintf("%s&digest=%s", url, layer.Digest)
1103

1104
1105
1106
1107
	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))
1108

1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
	// finish the upload
	resp, err := makeRequest("PUT", url, headers, r, regOpts)
	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))
1120
	}
1121
1122
1123
	return nil
}

1124
1125
1126
1127
1128
1129
1130
1131
1132
func makeRequest(method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
	if !strings.HasPrefix(url, "http") {
		if regOpts.Insecure {
			url = "http://" + url
		} else {
			url = "https://" + url
		}
	}

Patrick Devine's avatar
Patrick Devine committed
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
	// 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)
1145
1146
1147
1148
	if err != nil {
		return nil, err
	}

Patrick Devine's avatar
Patrick Devine committed
1149
1150
1151
1152
	if regOpts.Token != "" {
		req.Header.Set("Authorization", "Bearer "+regOpts.Token)
	} else if regOpts.Username != "" && regOpts.Password != "" {
		req.SetBasicAuth(regOpts.Username, regOpts.Password)
1153
1154
	}

Patrick Devine's avatar
Patrick Devine committed
1155
1156
	for k, v := range headers {
		req.Header.Set(k, v)
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
	}

	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
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
	// 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))
		token, err := getAuthToken(authRedir, regOpts)
		if err != nil {
			return nil, err
		}
		regOpts.Token = token
		bodyCopy = bytes.NewReader(buf.Bytes())
		return makeRequest(method, url, headers, bodyCopy, regOpts)
	}

1186
1187
	return resp, nil
}
Michael Yang's avatar
Michael Yang committed
1188

Patrick Devine's avatar
Patrick Devine committed
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
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"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
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 {
1238
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1239
1240
1241
1242
	}

	return nil
}