Compare commits

...

4 Commits

Author SHA1 Message Date
Michael Yang
4b4e97ed10 tests 2024-07-08 13:42:51 -07:00
Michael Yang
8d62a65ca7 rebase main 2024-07-08 11:18:16 -07:00
Michael Yang
bf5ba6065b migrate registry domain 2024-07-08 11:18:16 -07:00
Michael Yang
09c0972a6a filepath.ToSlash 2024-07-08 11:18:16 -07:00
10 changed files with 220 additions and 53 deletions

View File

@ -468,7 +468,7 @@ curl http://localhost:11434/api/chat -d '{
```json
{
"model": "registry.ollama.ai/library/llama3:latest",
"model": "ollama.com/library/llama3:latest",
"created_at": "2023-12-12T14:13:43.416799Z",
"message": {
"role": "assistant",
@ -606,7 +606,7 @@ curl http://localhost:11434/api/chat -d '{
```json
{
"model": "registry.ollama.ai/library/llama3:latest",
"model": "ollama.com/library/llama3:latest",
"created_at": "2023-12-12T14:13:43.416799Z",
"message": {
"role": "assistant",

3
go.mod
View File

@ -18,6 +18,7 @@ require (
require (
github.com/agnivade/levenshtein v1.1.1
github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1
github.com/google/go-cmp v0.6.0
github.com/mattn/go-runewidth v0.0.14
github.com/nlpodyssey/gopickle v0.3.0
github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c
@ -71,7 +72,7 @@ require (
golang.org/x/net v0.25.0 // indirect
golang.org/x/sys v0.20.0
golang.org/x/term v0.20.0
golang.org/x/text v0.15.0 // indirect
golang.org/x/text v0.15.0
google.golang.org/protobuf v1.34.1
gopkg.in/yaml.v3 v3.0.1 // indirect
)

View File

@ -108,6 +108,7 @@ func startServer(t *testing.T, ctx context.Context, ollamaHost string) error {
if tmp := os.Getenv("OLLAMA_HOST"); tmp != ollamaHost {
slog.Info("setting env", "OLLAMA_HOST", ollamaHost)
t.Setenv("OLLAMA_HOST", ollamaHost)
envconfig.LoadConfig()
}
slog.Info("starting server", "url", ollamaHost)

View File

@ -3,6 +3,8 @@ package server
import (
"errors"
"fmt"
"io/fs"
"log/slog"
"net/url"
"os"
"path/filepath"
@ -21,7 +23,7 @@ type ModelPath struct {
}
const (
DefaultRegistry = "registry.ollama.ai"
DefaultRegistry = "ollama.com"
DefaultNamespace = "library"
DefaultTag = "latest"
DefaultProtocolScheme = "https"
@ -49,8 +51,7 @@ func ParseModelPath(name string) ModelPath {
name = after
}
name = strings.ReplaceAll(name, string(os.PathSeparator), "/")
parts := strings.Split(name, "/")
parts := strings.Split(filepath.ToSlash(name), "/")
switch len(parts) {
case 3:
mp.Registry = parts[0]
@ -152,3 +153,61 @@ func GetBlobsPath(digest string) (string, error) {
return path, nil
}
func migrateRegistryDomain() error {
manifests, err := GetManifestPath()
if err != nil {
return err
}
targetDomain := filepath.Join(manifests, DefaultRegistry)
if _, err := os.Stat(targetDomain); errors.Is(err, fs.ErrNotExist) {
// noop
} else if err != nil {
return err
} else {
// target directory already exists so skip migration
return nil
}
sourceDomain := filepath.Join(manifests, "registry.ollama.ai")
//nolint:errcheck
defer PruneDirectory(sourceDomain)
return filepath.Walk(sourceDomain, func(source string, info fs.FileInfo, err error) error {
if errors.Is(err, os.ErrNotExist) {
return nil
} else if err != nil {
return err
}
if !info.IsDir() {
slog.Info("migrating registry domain", "path", source)
rel, err := filepath.Rel(sourceDomain, source)
if err != nil {
return err
}
target := filepath.Join(targetDomain, rel)
if _, err := os.Stat(target); errors.Is(err, fs.ErrNotExist) {
// noop
} else if err != nil {
return err
} else {
return nil
}
if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
return err
}
if err := os.Rename(source, target); err != nil {
return err
}
}
return nil
})
}

View File

@ -157,3 +157,94 @@ func TestParseModelPath(t *testing.T) {
})
}
}
func TestMigrateRegistryDomain(t *testing.T) {
p := t.TempDir()
t.Setenv("OLLAMA_MODELS", p)
envconfig.LoadConfig()
manifests := []string{
filepath.Join("registry.ollama.ai", "library", "llama3", "7b"),
filepath.Join("registry.ollama.ai", "library", "mistral", "latest"),
filepath.Join("registry.other.com", "library", "llama3", "13b"),
}
for _, manifest := range manifests {
n := filepath.Join(p, "manifests", manifest)
if err := os.MkdirAll(filepath.Dir(n), 0o750); err != nil {
t.Fatal(err)
}
f, err := os.Create(n)
if err != nil {
t.Fatal(err)
}
if err := f.Close(); err != nil {
t.Fatal(err)
}
}
t.Run("migrate", func(t *testing.T) {
if err := migrateRegistryDomain(); err != nil {
t.Fatal(err)
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", DefaultRegistry, "library", "llama3", "7b"),
filepath.Join(p, "manifests", DefaultRegistry, "library", "mistral", "latest"),
filepath.Join(p, "manifests", "registry.other.com", "library", "llama3", "13b"),
})
})
t.Run("idempotent", func(t *testing.T) {
// subsequent run should be a noop
if err := migrateRegistryDomain(); err != nil {
t.Fatal(err)
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", DefaultRegistry, "library", "llama3", "7b"),
filepath.Join(p, "manifests", DefaultRegistry, "library", "mistral", "latest"),
filepath.Join(p, "manifests", "registry.other.com", "library", "llama3", "13b"),
})
})
t.Run("no migration needed", func(t *testing.T) {
n := filepath.Join(p, "manifests", "registry.ollama.ai", "library", "gemma", "7b")
if err := os.MkdirAll(filepath.Dir(n), 0o750); err != nil {
t.Fatal(err)
}
f, err := os.Create(n)
if err != nil {
t.Fatal(err)
}
if err := f.Close(); err != nil {
t.Fatal(err)
}
if err := migrateRegistryDomain(); err != nil {
t.Fatal(err)
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", DefaultRegistry, "library", "llama3", "7b"),
filepath.Join(p, "manifests", DefaultRegistry, "library", "mistral", "latest"),
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "gemma", "7b"),
filepath.Join(p, "manifests", "registry.other.com", "library", "llama3", "13b"),
})
})
t.Run("no migration source", func(t *testing.T) {
// cleanup premigration directories
if err := os.RemoveAll(filepath.Join(p, "manifests", "registry.ollama.ai")); err != nil {
t.Fatal(err)
}
if err := migrateRegistryDomain(); err != nil {
t.Fatal(err)
}
})
}

View File

@ -1046,14 +1046,6 @@ func Serve(ln net.Listener) error {
slog.SetDefault(slog.New(handler))
blobsDir, err := GetBlobsPath("")
if err != nil {
return err
}
if err := fixBlobs(blobsDir); err != nil {
return err
}
if !envconfig.NoPrune {
// clean up unused layers and manifests
if err := PruneLayers(); err != nil {
@ -1070,6 +1062,24 @@ func Serve(ln net.Listener) error {
}
}
// MIGRATION
blobsDir, err := GetBlobsPath("")
if err != nil {
return err
}
if err := fixBlobs(blobsDir); err != nil {
return err
}
// migrate registry.ollama.ai to ollama.com
if err := migrateRegistryDomain(); err != nil {
if !errors.Is(err, os.ErrNotExist) {
return err
}
}
// END MIGRATION
ctx, done := context.WithCancel(context.Background())
schedCtx, schedDone := context.WithCancel(ctx)
sched := InitScheduler(schedCtx)

View File

@ -14,6 +14,7 @@ import (
"testing"
"github.com/gin-gonic/gin"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/envconfig"
"github.com/ollama/ollama/llm"
@ -79,8 +80,11 @@ func checkFileExists(t *testing.T, p string, expect []string) {
t.Fatal(err)
}
if !slices.Equal(actual, expect) {
t.Fatalf("expected slices to be equal %v", actual)
slices.Sort(actual)
slices.Sort(expect)
if diff := cmp.Diff(actual, expect); diff != "" {
t.Errorf("mismatch (-got, +want):\n%s", diff)
}
}
@ -101,7 +105,7 @@ func TestCreateFromBin(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -127,7 +131,7 @@ func TestCreateFromModel(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
w = createRequest(t, s.CreateModelHandler, api.CreateRequest{
@ -141,8 +145,8 @@ func TestCreateFromModel(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test2", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -168,7 +172,7 @@ func TestCreateRemovesLayers(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -188,7 +192,7 @@ func TestCreateRemovesLayers(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -215,7 +219,7 @@ func TestCreateUnsetsSystem(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -235,7 +239,7 @@ func TestCreateUnsetsSystem(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -271,7 +275,7 @@ func TestCreateMergeParameters(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -292,8 +296,8 @@ func TestCreateMergeParameters(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test2", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -330,8 +334,8 @@ func TestCreateMergeParameters(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test2", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -374,7 +378,7 @@ func TestCreateReplacesMessages(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -394,8 +398,8 @@ func TestCreateReplacesMessages(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test2", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -450,7 +454,7 @@ func TestCreateTemplateSystem(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -496,7 +500,7 @@ func TestCreateLicenses(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{

View File

@ -39,8 +39,8 @@ func TestDelete(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test2", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -57,7 +57,7 @@ func TestDelete(t *testing.T) {
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
filepath.Join(p, "manifests", "ollama.com", "library", "test2", "latest"),
})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -72,13 +72,14 @@ func TestDelete(t *testing.T) {
t.Fatalf("expected status code 200, actual %d", w.Code)
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{})
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{})
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), nil)
checkFileExists(t, filepath.Join(p, "blobs", "*"), nil)
}
func TestDeleteDuplicateLayers(t *testing.T) {
p := t.TempDir()
t.Setenv("OLLAMA_MODELS", p)
envconfig.LoadConfig()
var s Server
n := model.ParseName("test")
@ -103,5 +104,5 @@ func TestDeleteDuplicateLayers(t *testing.T) {
t.Errorf("expected status code 200, actual %d", w.Code)
}
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{})
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), nil)
}

View File

@ -35,7 +35,7 @@ func Unqualified(n Name) error {
const MissingPart = "!MISSING!"
const (
defaultHost = "registry.ollama.ai"
defaultHost = "ollama.com"
defaultNamespace = "library"
defaultTag = "latest"
)
@ -43,7 +43,7 @@ const (
// DefaultName returns a name with the default values for the host, namespace,
// and tag parts. The model and digest parts are empty.
//
// - The default host is ("registry.ollama.ai")
// - The default host is ("ollama.com")
// - The default namespace is ("library")
// - The default tag is ("latest")
func DefaultName() Name {

View File

@ -20,14 +20,14 @@ func TestParseNameParts(t *testing.T) {
wantValidDigest bool
}{
{
in: "registry.ollama.ai/library/dolphin-mistral:7b-v2.6-dpo-laser-q6_K",
in: "ollama.com/library/dolphin-mistral:7b-v2.6-dpo-laser-q6_K",
want: Name{
Host: "registry.ollama.ai",
Host: "ollama.com",
Namespace: "library",
Model: "dolphin-mistral",
Tag: "7b-v2.6-dpo-laser-q6_K",
},
wantFilepath: filepath.Join("registry.ollama.ai", "library", "dolphin-mistral", "7b-v2.6-dpo-laser-q6_K"),
wantFilepath: filepath.Join("ollama.com", "library", "dolphin-mistral", "7b-v2.6-dpo-laser-q6_K"),
},
{
in: "scheme://host:port/namespace/model:tag",
@ -83,14 +83,14 @@ func TestParseNameParts(t *testing.T) {
Namespace: "namespace",
Model: "model",
},
wantFilepath: filepath.Join("registry.ollama.ai", "namespace", "model", "latest"),
wantFilepath: filepath.Join("ollama.com", "namespace", "model", "latest"),
},
{
in: "model",
want: Name{
Model: "model",
},
wantFilepath: filepath.Join("registry.ollama.ai", "library", "model", "latest"),
wantFilepath: filepath.Join("ollama.com", "library", "model", "latest"),
},
{
in: "h/nn/mm:t",
@ -193,7 +193,7 @@ func TestNameparseNameDefault(t *testing.T) {
const name = "xx"
n := ParseName(name)
got := n.String()
want := "registry.ollama.ai/library/xx:latest"
want := "ollama.com/library/xx:latest"
if got != want {
t.Errorf("parseName(%q).String() = %q; want %q", name, got, want)
}
@ -290,11 +290,11 @@ func TestParseNameFromFilepath(t *testing.T) {
func TestDisplayShortest(t *testing.T) {
cases := map[string]string{
"registry.ollama.ai/library/model:latest": "model:latest",
"registry.ollama.ai/library/model:tag": "model:tag",
"registry.ollama.ai/namespace/model:tag": "namespace/model:tag",
"host/namespace/model:tag": "host/namespace/model:tag",
"host/library/model:tag": "host/library/model:tag",
"ollama.com/library/model:latest": "model:latest",
"ollama.com/library/model:tag": "model:tag",
"ollama.com/namespace/model:tag": "namespace/model:tag",
"host/namespace/model:tag": "host/namespace/model:tag",
"host/library/model:tag": "host/library/model:tag",
}
for in, want := range cases {