diff --git a/build/blob/ref.go b/build/blob/ref.go index 9a033fcb..dc7bdcef 100644 --- a/build/blob/ref.go +++ b/build/blob/ref.go @@ -2,6 +2,8 @@ package blob import ( "cmp" + "path" + "path/filepath" "strings" ) @@ -11,26 +13,31 @@ import ( // // Users or Ref must check Valid before using it. type Ref struct { - name string - tag string - build string + domain string + name string + tag string + build string } // WithBuild returns a copy of r with the provided build. If the provided // build is empty, it returns the short, unqualified copy of r. func (r Ref) WithBuild(build string) Ref { if build == "" { - return Ref{r.name, r.tag, ""} + return Ref{r.domain, r.name, r.tag, ""} } if !isValidPart(build) { return Ref{} } - return makeRef(r.name, r.tag, build) + return makeRef(r.domain, r.name, r.tag, build) } // String returns the fully qualified ref string. func (r Ref) String() string { var b strings.Builder + if r.domain != "" { + b.WriteString(r.domain) + b.WriteString("/") + } b.WriteString(r.name) if r.tag != "" { b.WriteString(":") @@ -49,7 +56,7 @@ func (r Ref) Full() string { if !r.Valid() { return "" } - return makeRef(r.name, r.tag, cmp.Or(r.build, "!(MISSING BUILD)")).String() + return makeRef(r.domain, r.name, r.tag, cmp.Or(r.build, "!(MISSING BUILD)")).String() } // Short returns the short ref string which does not include the build. @@ -65,9 +72,18 @@ func (r Ref) FullyQualified() bool { return r.name != "" && r.tag != "" && r.build != "" } -func (r Ref) Name() string { return r.name } -func (r Ref) Tag() string { return r.tag } -func (r Ref) Build() string { return r.build } +func (r Ref) Path() string { + return path.Join(r.domain, r.name, r.tag, r.build) +} + +func (r Ref) Filepath() string { + return filepath.Join(r.domain, r.name, r.tag, r.build) +} + +func (r Ref) Domain() string { return r.domain } +func (r Ref) Name() string { return r.name } +func (r Ref) Tag() string { return r.tag } +func (r Ref) Build() string { return r.build } // ParseRef parses a ref string into a Ref. A ref string is a name, an // optional tag, and an optional build, separated by colons and pluses. @@ -107,12 +123,14 @@ func ParseRef(s string) Ref { if expectBuild && !isValidPart(build) { return Ref{} } - return makeRef(name, tag, build) + + const TODO = "registry.ollama.ai" + return makeRef(TODO, name, tag, build) } // makeRef makes a ref, skipping validation. -func makeRef(name, tag, build string) Ref { - return Ref{name, cmp.Or(tag, "latest"), strings.ToUpper(build)} +func makeRef(domain, name, tag, build string) Ref { + return Ref{domain, name, cmp.Or(tag, "latest"), strings.ToUpper(build)} } // isValidPart returns true if given part is valid ascii [a-zA-Z0-9_\.-] diff --git a/build/blob/ref_test.go b/build/blob/ref_test.go index b49d39df..1322022f 100644 --- a/build/blob/ref_test.go +++ b/build/blob/ref_test.go @@ -12,24 +12,24 @@ func TestParseRef(t *testing.T) { in string want Ref }{ - {"mistral:latest", Ref{"mistral", "latest", ""}}, - {"mistral", Ref{"mistral", "latest", ""}}, - {"mistral:30B", Ref{"mistral", "30B", ""}}, - {"mistral:7b", Ref{"mistral", "7b", ""}}, - {"mistral:7b+Q4_0", Ref{"mistral", "7b", "Q4_0"}}, - {"mistral+KQED", Ref{"mistral", "latest", "KQED"}}, - {"mistral.x-3:7b+Q4_0", Ref{"mistral.x-3", "7b", "Q4_0"}}, + {"mistral:latest", Ref{"registry.ollama.ai", "mistral", "latest", ""}}, + {"mistral", Ref{"registry.ollama.ai", "mistral", "latest", ""}}, + {"mistral:30B", Ref{"registry.ollama.ai", "mistral", "30B", ""}}, + {"mistral:7b", Ref{"registry.ollama.ai", "mistral", "7b", ""}}, + {"mistral:7b+Q4_0", Ref{"registry.ollama.ai", "mistral", "7b", "Q4_0"}}, + {"mistral+KQED", Ref{"registry.ollama.ai", "mistral", "latest", "KQED"}}, + {"mistral.x-3:7b+Q4_0", Ref{"registry.ollama.ai", "mistral.x-3", "7b", "Q4_0"}}, // lowecase build - {"mistral:7b+q4_0", Ref{"mistral", "7b", "Q4_0"}}, + {"mistral:7b+q4_0", Ref{"registry.ollama.ai", "mistral", "7b", "Q4_0"}}, // Invalid - {"mistral:7b+Q4_0:latest", Ref{"", "", ""}}, - {"mi tral", Ref{"", "", ""}}, - {"llama2:+", Ref{"", "", ""}}, + {"mistral:7b+Q4_0:latest", Ref{"", "", "", ""}}, + {"mi tral", Ref{"", "", "", ""}}, + {"llama2:+", Ref{"", "", "", ""}}, // too long - {refTooLong, Ref{"", "", ""}}, + {refTooLong, Ref{"", "", "", ""}}, } for _, tt := range cases { t.Run(tt.in, func(t *testing.T) { @@ -48,11 +48,11 @@ func TestRefFull(t *testing.T) { wantFull string }{ {"", "", ""}, - {"mistral:7b+x", "mistral:7b", "mistral:7b+X"}, - {"mistral:7b+Q4_0", "mistral:7b", "mistral:7b+Q4_0"}, - {"mistral:latest", "mistral:latest", "mistral:latest+!(MISSING BUILD)"}, - {"mistral", "mistral:latest", "mistral:latest+!(MISSING BUILD)"}, - {"mistral:30b", "mistral:30b", "mistral:30b+!(MISSING BUILD)"}, + {"mistral:7b+x", "registry.ollama.ai/mistral:7b", "registry.ollama.ai/mistral:7b+X"}, + {"mistral:7b+Q4_0", "registry.ollama.ai/mistral:7b", "registry.ollama.ai/mistral:7b+Q4_0"}, + {"mistral:latest", "registry.ollama.ai/mistral:latest", "registry.ollama.ai/mistral:latest+!(MISSING BUILD)"}, + {"mistral", "registry.ollama.ai/mistral:latest", "registry.ollama.ai/mistral:latest+!(MISSING BUILD)"}, + {"mistral:30b", "registry.ollama.ai/mistral:30b", "registry.ollama.ai/mistral:30b+!(MISSING BUILD)"}, } for _, tt := range cases { diff --git a/build/build.go b/build/build.go index 21491102..f66138db 100644 --- a/build/build.go +++ b/build/build.go @@ -6,7 +6,6 @@ import ( "fmt" "io/fs" "os" - "path" "bllamo.com/build/blob" "bllamo.com/build/internal/blobstore" @@ -22,10 +21,6 @@ var ( ErrNotFound = errors.New("not found") ) -func ManifestKey(domain string, ref blob.Ref) string { - return path.Join("manifests", domain, ref.Name(), ref.Tag(), ref.Build()) -} - type mediaType string // Known media types diff --git a/build/internal/blobstore/blob.go b/build/internal/blobstore/blob.go index e5981416..3c29538a 100644 --- a/build/internal/blobstore/blob.go +++ b/build/internal/blobstore/blob.go @@ -1,7 +1,3 @@ -// Copyright 2017 The Go Authors. All rights reserved. -// Use of this source code is governed by a BSD-style -// license that can be found in the LICENSE file. - // Package blobstore implements a blob store. package blobstore @@ -228,8 +224,7 @@ func (s *Store) refFileName(ref blob.Ref) (string, error) { if !ref.FullyQualified() { return "", fmt.Errorf("ref not fully qualified: %q", ref) } - const cheatTODO = "registry.ollama.ai/library" - return filepath.Join(s.dir, "manifests", cheatTODO, ref.Name(), ref.Tag(), ref.Build()), nil + return filepath.Join(s.dir, "manifests", ref.Domain(), ref.Name(), ref.Tag(), ref.Build()), nil } // Get looks up the blob ID in the store, diff --git a/registry/server.go b/registry/server.go index b22042bb..f22d0b1f 100644 --- a/registry/server.go +++ b/registry/server.go @@ -1,4 +1,4 @@ -// Package implements an Ollama registry client and server +// Package implements an Ollama registry client and server package registry package registry import ( @@ -8,9 +8,10 @@ import ( "errors" "log" "net/http" + "os" + "path" "time" - "bllamo.com/build" "bllamo.com/build/blob" "bllamo.com/client/ollama" "bllamo.com/oweb" @@ -19,6 +20,13 @@ import ( "github.com/minio/minio-go/v7/pkg/credentials" ) +// TODO(bmizerany): move all env things to package envkobs? +var defaultLibrary = cmp.Or(os.Getenv("OLLAMA_REGISTRY"), "registry.ollama.ai/library") + +func DefaultLibrary() string { + return defaultLibrary +} + type Server struct{} func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -80,7 +88,8 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error { } if !pushed { const expires = 1 * time.Hour - signedURL, err := mc.PresignedPutObject(r.Context(), "test", l.Digest, expires) + key := path.Join("blobs", l.Digest) + signedURL, err := mc.PresignedPutObject(r.Context(), "test", key, expires) if err != nil { return err } @@ -95,9 +104,10 @@ func (s *Server) handlePush(w http.ResponseWriter, r *http.Request) error { } if len(requirements) == 0 { - const cheatTODO = "registry.ollama.ai/library" - key := build.ManifestKey(cheatTODO, ref) - _, err := mc.PutObject(r.Context(), "test", key, bytes.NewReader(pr.Manifest), int64(len(pr.Manifest)), minio.PutObjectOptions{}) + // Commit the manifest + body := bytes.NewReader(pr.Manifest) + path := path.Join("manifests", ref.Path()) + _, err := mc.PutObject(r.Context(), "test", path, body, int64(len(pr.Manifest)), minio.PutObjectOptions{}) if err != nil { return err } @@ -122,7 +132,8 @@ func (s *Server) statObject(ctx context.Context, digest string) (pushed bool, er } // HEAD the object - _, err = mc.StatObject(ctx, "test", digest, minio.StatObjectOptions{}) + path := path.Join("blobs", digest) + _, err = mc.StatObject(ctx, "test", path, minio.StatObjectOptions{}) if err != nil { if isNoSuchKey(err) { err = nil diff --git a/registry/server_test.go b/registry/server_test.go index 79c1875b..0e268c76 100644 --- a/registry/server_test.go +++ b/registry/server_test.go @@ -2,6 +2,7 @@ package registry import ( "context" + "encoding/json" "net/http/httptest" "os/exec" "strings" @@ -57,6 +58,50 @@ func TestPush(t *testing.T) { if len(got) != 0 { t.Fatalf("unexpected requirements: % #v", pretty.Formatter(got)) } + + mc, err := minio.New("localhost:9000", &minio.Options{ + Creds: credentials.NewStaticV4("minioadmin", "minioadmin", ""), + Secure: false, + }) + if err != nil { + t.Fatal(err) + } + + var paths []string + keys := mc.ListObjects(context.Background(), "test", minio.ListObjectsOptions{ + Recursive: true, + }) + for k := range keys { + paths = append(paths, k.Key) + } + + t.Logf("paths: %v", paths) + + diff.Test(t, t.Errorf, paths, []string{ + "blobs/sha256-1", + "blobs/sha256-2", + "blobs/sha256-3", + "manifests/registry.ollama.ai/x/latest/Y", + }) + + obj, err := mc.GetObject(context.Background(), "test", "manifests/registry.ollama.ai/x/latest/Y", minio.GetObjectOptions{}) + if err != nil { + t.Fatal(err) + } + defer obj.Close() + + var gotM apitype.Manifest + if err := json.NewDecoder(obj).Decode(&gotM); err != nil { + t.Fatal(err) + } + + diff.Test(t, t.Errorf, gotM, apitype.Manifest{ + Layers: []apitype.Layer{ + {Digest: "sha256-1", Size: 1}, + {Digest: "sha256-2", Size: 2}, + {Digest: "sha256-3", Size: 3}, + }, + }) } func startMinio(t *testing.T) {