manifest.go 3.18 KB
Newer Older
1
2
3
4
package server

import (
	"crypto/sha256"
Michael Yang's avatar
lint  
Michael Yang committed
5
	"encoding/hex"
6
	"encoding/json"
7
	"errors"
8
	"fmt"
9
	"io"
10
	"log/slog"
11
12
13
14
15
16
17
	"os"
	"path/filepath"

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

type Manifest struct {
18
19
20
21
	SchemaVersion int     `json:"schemaVersion"`
	MediaType     string  `json:"mediaType"`
	Config        Layer   `json:"config"`
	Layers        []Layer `json:"layers"`
22
23

	filepath string
24
	fi       os.FileInfo
25
	digest   string
26
27
28
}

func (m *Manifest) Size() (size int64) {
29
	for _, layer := range append(m.Layers, m.Config) {
30
31
32
33
34
35
		size += layer.Size
	}

	return
}

36
37
38
39
40
41
42
43
44
45
46
47
48
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)
}

49
func (m *Manifest) RemoveLayers() error {
50
	for _, layer := range append(m.Layers, m.Config) {
51
52
53
54
55
56
		if layer.Digest != "" {
			if err := layer.Remove(); errors.Is(err, os.ErrNotExist) {
				slog.Debug("layer does not exist", "digest", layer.Digest)
			} else if err != nil {
				return err
			}
57
58
59
60
61
62
		}
	}

	return nil
}

63
64
65
func ParseNamedManifest(n model.Name) (*Manifest, error) {
	if !n.IsFullyQualified() {
		return nil, model.Unqualified(n)
66
67
68
69
70
71
72
	}

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

73
74
	p := filepath.Join(manifests, n.Filepath())

Michael Yang's avatar
Michael Yang committed
75
	var m Manifest
76
	f, err := os.Open(p)
77
78
79
	if err != nil {
		return nil, err
	}
80
	defer f.Close()
81

82
83
84
85
86
	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

87
	sha256sum := sha256.New()
88
	if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&m); err != nil {
89
90
91
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
92
93
	m.filepath = p
	m.fi = fi
Michael Yang's avatar
lint  
Michael Yang committed
94
	m.digest = hex.EncodeToString(sha256sum.Sum(nil))
Michael Yang's avatar
Michael Yang committed
95
96

	return &m, nil
97
98
}

99
func WriteManifest(name model.Name, config Layer, layers []Layer) error {
100
101
102
	manifests, err := GetManifestPath()
	if err != nil {
		return err
103
104
	}

105
106
	p := filepath.Join(manifests, name.Filepath())
	if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
107
108
109
		return err
	}

110
	f, err := os.Create(p)
111
112
113
	if err != nil {
		return err
	}
114
	defer f.Close()
115

Michael Yang's avatar
Michael Yang committed
116
	m := Manifest{
117
118
		SchemaVersion: 2,
		MediaType:     "application/vnd.docker.distribution.manifest.v2+json",
119
		Config:        config,
120
		Layers:        layers,
121
122
	}

123
	return json.NewEncoder(f).Encode(m)
124
}
125
126
127
128
129
130
131
132

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
133
	matches, err := filepath.Glob(filepath.Join(manifests, "*", "*", "*", "*"))
134
135
136
137
138
139
	if err != nil {
		return nil, err
	}

	ms := make(map[model.Name]*Manifest)
	for _, match := range matches {
Michael Yang's avatar
Michael Yang committed
140
		fi, err := os.Stat(match)
141
142
143
144
		if err != nil {
			return nil, err
		}

Michael Yang's avatar
Michael Yang committed
145
146
147
148
149
150
151
152
153
		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() {
Michael Yang's avatar
Michael Yang committed
154
				slog.Warn("bad manifest name", "path", rel)
Michael Yang's avatar
Michael Yang committed
155
156
157
				continue
			}

158
			m, err := ParseNamedManifest(n)
159
			if syntax := &(json.SyntaxError{}); errors.As(err, &syntax) {
160
161
				slog.Warn("bad manifest", "name", n, "error", err)
				continue
162
163
			} else if err != nil {
				return nil, fmt.Errorf("%s: %w", n, err)
164
165
166
167
168
169
170
171
			}

			ms[n] = m
		}
	}

	return ms, nil
}