Clean up the manifest and modelpath (#13807)

manifest/layer.go

@@ -0,0 +1,135 @@
+package manifest
+
+import (
+	"crypto/sha256"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+)
+
+type Layer struct {
+	MediaType string `json:"mediaType"`
+	Digest    string `json:"digest"`
+	Size      int64  `json:"size"`
+	From      string `json:"from,omitempty"`
+	Name      string `json:"name,omitempty"` // tensor name, e.g., "text_encoder/model.embed_tokens.weight"
+	Status    string `json:"-"`
+}
+
+const (
+	MediaTypeImageTensor = "application/vnd.ollama.image.tensor"
+)
+
+func NewLayer(r io.Reader, mediatype string) (Layer, error) {
+	blobs, err := BlobsPath("")
+	if err != nil {
+		return Layer{}, err
+	}
+
+	temp, err := os.CreateTemp(blobs, "sha256-")
+	if err != nil {
+		return Layer{}, err
+	}
+	defer temp.Close()
+	defer os.Remove(temp.Name())
+
+	sha256sum := sha256.New()
+	n, err := io.Copy(io.MultiWriter(temp, sha256sum), r)
+	if err != nil {
+		return Layer{}, err
+	}
+
+	if err := temp.Close(); err != nil {
+		return Layer{}, err
+	}
+
+	digest := fmt.Sprintf("sha256:%x", sha256sum.Sum(nil))
+	blob, err := BlobsPath(digest)
+	if err != nil {
+		return Layer{}, err
+	}
+
+	status := "using existing layer"
+	if _, err := os.Stat(blob); err != nil {
+		status = "creating new layer"
+		if err := os.Rename(temp.Name(), blob); err != nil {
+			return Layer{}, err
+		}
+		if err := os.Chmod(blob, 0o644); err != nil {
+			return Layer{}, err
+		}
+	}
+
+	return Layer{
+		MediaType: mediatype,
+		Digest:    digest,
+		Size:      n,
+		Status:    fmt.Sprintf("%s %s", status, digest),
+	}, nil
+}
+
+func NewLayerFromLayer(digest, mediatype, from string) (Layer, error) {
+	if digest == "" {
+		return Layer{}, errors.New("creating new layer from layer with empty digest")
+	}
+
+	blob, err := BlobsPath(digest)
+	if err != nil {
+		return Layer{}, err
+	}
+
+	fi, err := os.Stat(blob)
+	if err != nil {
+		return Layer{}, err
+	}
+
+	return Layer{
+		MediaType: mediatype,
+		Digest:    digest,
+		Size:      fi.Size(),
+		From:      from,
+		Status:    fmt.Sprintf("using existing layer %s", digest),
+	}, nil
+}
+
+func (l *Layer) Open() (io.ReadSeekCloser, error) {
+	if l.Digest == "" {
+		return nil, errors.New("opening layer with empty digest")
+	}
+
+	blob, err := BlobsPath(l.Digest)
+	if err != nil {
+		return nil, err
+	}
+
+	return os.Open(blob)
+}
+
+func (l *Layer) Remove() error {
+	if l.Digest == "" {
+		return nil
+	}
+
+	// Ignore corrupt manifests to avoid blocking deletion of layers that are freshly orphaned
+	ms, err := Manifests(true)
+	if err != nil {
+		return err
+	}
+
+	for _, m := range ms {
+		for _, layer := range append(m.Layers, m.Config) {
+			if layer.Digest == l.Digest {
+				// something is using this layer
+				return nil
+			}
+		}
+	}
+
+	blob, err := BlobsPath(l.Digest)
+	if err != nil {
+		return err
+	}
+
+	return os.Remove(blob)
+}

manifest/manifest.go

@@ -0,0 +1,227 @@
+package manifest
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"path/filepath"
+
+	"github.com/ollama/ollama/types/model"
+)
+
+type Manifest struct {
+	SchemaVersion int     `json:"schemaVersion"`
+	MediaType     string  `json:"mediaType"`
+	Config        Layer   `json:"config"`
+	Layers        []Layer `json:"layers"`
+
+	filepath string
+	fi       os.FileInfo
+	digest   string
+}
+
+func (m *Manifest) Size() (size int64) {
+	for _, layer := range append(m.Layers, m.Config) {
+		size += layer.Size
+	}
+
+	return
+}
+
+func (m *Manifest) Digest() string {
+	return m.digest
+}
+
+func (m *Manifest) FileInfo() os.FileInfo {
+	return m.fi
+}
+
+// ReadConfigJSON reads and unmarshals a config layer as JSON.
+func (m *Manifest) ReadConfigJSON(configPath string, v any) error {
+	for _, layer := range m.Layers {
+		if layer.MediaType == "application/vnd.ollama.image.json" && layer.Name == configPath {
+			blobPath, err := BlobsPath(layer.Digest)
+			if err != nil {
+				return err
+			}
+			data, err := os.ReadFile(blobPath)
+			if err != nil {
+				return err
+			}
+			return json.Unmarshal(data, v)
+		}
+	}
+	return fmt.Errorf("config %q not found in manifest", configPath)
+}
+
+func (m *Manifest) Remove() error {
+	if err := os.Remove(m.filepath); err != nil {
+		return err
+	}
+
+	manifests, err := Path()
+	if err != nil {
+		return err
+	}
+
+	return PruneDirectory(manifests)
+}
+
+func (m *Manifest) RemoveLayers() error {
+	ms, err := Manifests(true)
+	if err != nil {
+		return err
+	}
+
+	// Build set of digests still in use by other manifests
+	inUse := make(map[string]struct{})
+	for _, other := range ms {
+		for _, layer := range append(other.Layers, other.Config) {
+			if layer.Digest != "" {
+				inUse[layer.Digest] = struct{}{}
+			}
+		}
+	}
+
+	// Remove layers not used by any other manifest
+	for _, layer := range append(m.Layers, m.Config) {
+		if layer.Digest == "" {
+			continue
+		}
+		if _, used := inUse[layer.Digest]; used {
+			continue
+		}
+		blob, err := BlobsPath(layer.Digest)
+		if err != nil {
+			return err
+		}
+		if err := os.Remove(blob); os.IsNotExist(err) {
+			slog.Debug("layer does not exist", "digest", layer.Digest)
+		} else if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func ParseNamedManifest(n model.Name) (*Manifest, error) {
+	if !n.IsFullyQualified() {
+		return nil, model.Unqualified(n)
+	}
+
+	manifests, err := Path()
+	if err != nil {
+		return nil, err
+	}
+
+	p := filepath.Join(manifests, n.Filepath())
+
+	var m Manifest
+	f, err := os.Open(p)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	fi, err := f.Stat()
+	if err != nil {
+		return nil, err
+	}
+
+	sha256sum := sha256.New()
+	if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&m); err != nil {
+		return nil, err
+	}
+
+	m.filepath = p
+	m.fi = fi
+	m.digest = hex.EncodeToString(sha256sum.Sum(nil))
+
+	return &m, nil
+}
+
+func WriteManifest(name model.Name, config Layer, layers []Layer) error {
+	manifests, err := Path()
+	if err != nil {
+		return err
+	}
+
+	p := filepath.Join(manifests, name.Filepath())
+	if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
+		return err
+	}
+
+	f, err := os.Create(p)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	m := Manifest{
+		SchemaVersion: 2,
+		MediaType:     "application/vnd.docker.distribution.manifest.v2+json",
+		Config:        config,
+		Layers:        layers,
+	}
+
+	return json.NewEncoder(f).Encode(m)
+}
+
+func Manifests(continueOnError bool) (map[model.Name]*Manifest, error) {
+	manifests, err := Path()
+	if err != nil {
+		return nil, err
+	}
+
+	// TODO(mxyng): use something less brittle
+	matches, err := filepath.Glob(filepath.Join(manifests, "*", "*", "*", "*"))
+	if err != nil {
+		return nil, err
+	}
+
+	ms := make(map[model.Name]*Manifest)
+	for _, match := range matches {
+		fi, err := os.Stat(match)
+		if err != nil {
+			return nil, err
+		}
+
+		if !fi.IsDir() {
+			rel, err := filepath.Rel(manifests, match)
+			if err != nil {
+				if !continueOnError {
+					return nil, fmt.Errorf("%s %w", match, err)
+				}
+				slog.Warn("bad filepath", "path", match, "error", err)
+				continue
+			}
+
+			n := model.ParseNameFromFilepath(rel)
+			if !n.IsValid() {
+				if !continueOnError {
+					return nil, fmt.Errorf("%s %w", rel, err)
+				}
+				slog.Warn("bad manifest name", "path", rel)
+				continue
+			}
+
+			m, err := ParseNamedManifest(n)
+			if err != nil {
+				if !continueOnError {
+					return nil, fmt.Errorf("%s %w", n, err)
+				}
+				slog.Warn("bad manifest", "name", n, "error", err)
+				continue
+			}
+
+			ms[n] = m
+		}
+	}
+
+	return ms, nil
+}

manifest/manifest_test.go

@@ -0,0 +1,150 @@
+package manifest
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"slices"
+	"testing"
+
+	"github.com/ollama/ollama/types/model"
+)
+
+func createManifest(t *testing.T, path, name string) {
+	t.Helper()
+
+	p := filepath.Join(path, "manifests", name)
+	if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
+		t.Fatal(err)
+	}
+
+	f, err := os.Create(p)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	if err := json.NewEncoder(f).Encode(Manifest{}); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestManifests(t *testing.T) {
+	cases := map[string]struct {
+		ps               []string
+		wantValidCount   int
+		wantInvalidCount int
+	}{
+		"empty": {},
+		"single": {
+			ps: []string{
+				filepath.Join("host", "namespace", "model", "tag"),
+			},
+			wantValidCount: 1,
+		},
+		"multiple": {
+			ps: []string{
+				filepath.Join("registry.ollama.ai", "library", "llama3", "latest"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q4_0"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q4_1"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q8_0"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q5_0"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q5_1"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q2_K"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q3_K_S"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q3_K_M"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q3_K_L"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q4_K_S"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q4_K_M"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q5_K_S"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q5_K_M"),
+				filepath.Join("registry.ollama.ai", "library", "llama3", "q6_K"),
+			},
+			wantValidCount: 15,
+		},
+		"hidden": {
+			ps: []string{
+				filepath.Join("host", "namespace", "model", "tag"),
+				filepath.Join("host", "namespace", "model", ".hidden"),
+			},
+			wantValidCount:   1,
+			wantInvalidCount: 1,
+		},
+		"subdir": {
+			ps: []string{
+				filepath.Join("host", "namespace", "model", "tag", "one"),
+				filepath.Join("host", "namespace", "model", "tag", "another", "one"),
+			},
+			wantInvalidCount: 2,
+		},
+		"upper tag": {
+			ps: []string{
+				filepath.Join("host", "namespace", "model", "TAG"),
+			},
+			wantValidCount: 1,
+		},
+		"upper model": {
+			ps: []string{
+				filepath.Join("host", "namespace", "MODEL", "tag"),
+			},
+			wantValidCount: 1,
+		},
+		"upper namespace": {
+			ps: []string{
+				filepath.Join("host", "NAMESPACE", "model", "tag"),
+			},
+			wantValidCount: 1,
+		},
+		"upper host": {
+			ps: []string{
+				filepath.Join("HOST", "namespace", "model", "tag"),
+			},
+			wantValidCount: 1,
+		},
+	}
+
+	for n, wants := range cases {
+		t.Run(n, func(t *testing.T) {
+			d := t.TempDir()
+			t.Setenv("OLLAMA_MODELS", d)
+
+			for _, p := range wants.ps {
+				createManifest(t, d, p)
+			}
+
+			ms, err := Manifests(true)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			var ns []model.Name
+			for k := range ms {
+				ns = append(ns, k)
+			}
+
+			var gotValidCount, gotInvalidCount int
+			for _, p := range wants.ps {
+				n := model.ParseNameFromFilepath(p)
+				if n.IsValid() {
+					gotValidCount++
+				} else {
+					gotInvalidCount++
+				}
+
+				if !n.IsValid() && slices.Contains(ns, n) {
+					t.Errorf("unexpected invalid name: %s", p)
+				} else if n.IsValid() && !slices.Contains(ns, n) {
+					t.Errorf("missing valid name: %s", p)
+				}
+			}
+
+			if gotValidCount != wants.wantValidCount {
+				t.Errorf("got valid count %d, want %d", gotValidCount, wants.wantValidCount)
+			}
+
+			if gotInvalidCount != wants.wantInvalidCount {
+				t.Errorf("got invalid count %d, want %d", gotInvalidCount, wants.wantInvalidCount)
+			}
+		})
+	}
+}

manifest/paths.go

@@ -0,0 +1,95 @@
+package manifest
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
+	"strings"
+
+	"github.com/ollama/ollama/envconfig"
+	"github.com/ollama/ollama/types/model"
+)
+
+var ErrInvalidDigestFormat = errors.New("invalid digest format")
+
+func Path() (string, error) {
+	path := filepath.Join(envconfig.Models(), "manifests")
+	if err := os.MkdirAll(path, 0o755); err != nil {
+		return "", fmt.Errorf("%w: ensure path elements are traversable", err)
+	}
+
+	return path, nil
+}
+
+// PathForName returns the path to the manifest file for a specific model name.
+func PathForName(n model.Name) (string, error) {
+	if !n.IsValid() {
+		return "", os.ErrNotExist
+	}
+
+	manifests, err := Path()
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(manifests, n.Filepath()), nil
+}
+
+func BlobsPath(digest string) (string, error) {
+	// only accept actual sha256 digests
+	pattern := "^sha256[:-][0-9a-fA-F]{64}$"
+	re := regexp.MustCompile(pattern)
+
+	if digest != "" && !re.MatchString(digest) {
+		return "", ErrInvalidDigestFormat
+	}
+
+	digest = strings.ReplaceAll(digest, ":", "-")
+	path := filepath.Join(envconfig.Models(), "blobs", digest)
+	dirPath := filepath.Dir(path)
+	if digest == "" {
+		dirPath = path
+	}
+
+	if err := os.MkdirAll(dirPath, 0o755); err != nil {
+		return "", fmt.Errorf("%w: ensure path elements are traversable", err)
+	}
+
+	return path, nil
+}
+
+// PruneDirectory removes empty directories recursively.
+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
+}