manifest.go 2.81 KB
Newer Older
1
2
3
4
5
6
7
package server

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
8
	"log/slog"
9
10
11
12
13
14
15
16
	"os"
	"path/filepath"

	"github.com/ollama/ollama/types/model"
)

type Manifest struct {
	ManifestV2
17
18

	filepath string
19
	fi       os.FileInfo
20
	digest   string
21
22
23
24
25
26
27
28
29
30
}

func (m *Manifest) Size() (size int64) {
	for _, layer := range append(m.Layers, m.Config) {
		size += layer.Size
	}

	return
}

31
32
33
34
35
36
37
38
39
40
41
42
43
func (m *Manifest) Remove() error {
	if err := os.Remove(m.filepath); err != nil {
		return err
	}

	manifests, err := GetManifestPath()
	if err != nil {
		return err
	}

	return PruneDirectory(manifests)
}

44
45
46
47
48
49
50
51
52
53
func (m *Manifest) RemoveLayers() error {
	for _, layer := range append(m.Layers, m.Config) {
		if err := layer.Remove(); err != nil {
			return err
		}
	}

	return nil
}

54
55
56
func ParseNamedManifest(n model.Name) (*Manifest, error) {
	if !n.IsFullyQualified() {
		return nil, model.Unqualified(n)
57
58
59
60
61
62
63
	}

	manifests, err := GetManifestPath()
	if err != nil {
		return nil, err
	}

64
65
66
67
	p := filepath.Join(manifests, n.Filepath())

	var m ManifestV2
	f, err := os.Open(p)
68
69
70
	if err != nil {
		return nil, err
	}
71
	defer f.Close()
72

73
74
75
76
77
	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

78
	sha256sum := sha256.New()
79
	if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&m); err != nil {
80
81
82
83
		return nil, err
	}

	return &Manifest{
84
85
		ManifestV2: m,
		filepath:   p,
86
		fi:         fi,
87
		digest:     fmt.Sprintf("%x", sha256sum.Sum(nil)),
88
89
90
	}, nil
}

91
92
93
94
func WriteManifest(name model.Name, config *Layer, layers []*Layer) error {
	manifests, err := GetManifestPath()
	if err != nil {
		return err
95
96
	}

97
98
	p := filepath.Join(manifests, name.Filepath())
	if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
99
100
101
		return err
	}

102
	f, err := os.Create(p)
103
104
105
	if err != nil {
		return err
	}
106
	defer f.Close()
107

108
109
110
111
112
	m := ManifestV2{
		SchemaVersion: 2,
		MediaType:     "application/vnd.docker.distribution.manifest.v2+json",
		Config:        config,
		Layers:        layers,
113
114
	}

115
	return json.NewEncoder(f).Encode(m)
116
}
117
118
119
120
121
122
123
124

func Manifests() (map[model.Name]*Manifest, error) {
	manifests, err := GetManifestPath()
	if err != nil {
		return nil, err
	}

	// TODO(mxyng): use something less brittle
Michael Yang's avatar
Michael Yang committed
125
	matches, err := filepath.Glob(filepath.Join(manifests, "*", "*", "*", "*"))
126
127
128
129
130
131
	if err != nil {
		return nil, err
	}

	ms := make(map[model.Name]*Manifest)
	for _, match := range matches {
Michael Yang's avatar
Michael Yang committed
132
		fi, err := os.Stat(match)
133
134
135
136
		if err != nil {
			return nil, err
		}

Michael Yang's avatar
Michael Yang committed
137
138
139
140
141
142
143
144
145
146
147
148
149
		if !fi.IsDir() {
			rel, err := filepath.Rel(manifests, match)
			if err != nil {
				slog.Warn("bad filepath", "path", match, "error", err)
				continue
			}

			n := model.ParseNameFromFilepath(rel)
			if !n.IsValid() {
				slog.Warn("bad manifest name", "path", rel, "error", err)
				continue
			}

150
151
			m, err := ParseNamedManifest(n)
			if err != nil {
152
153
				slog.Warn("bad manifest", "name", n, "error", err)
				continue
154
155
156
157
158
159
160
161
			}

			ms[n] = m
		}
	}

	return ms, nil
}