From 4b4e97ed103f07002503076ebc250cd08f88a367 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Mon, 20 May 2024 14:04:27 -0700 Subject: [PATCH] tests --- go.mod | 3 +- integration/utils_test.go | 1 + server/modelpath.go | 44 +++++++++++++---- server/modelpath_test.go | 91 ++++++++++++++++++++++++++++++++++++ server/routes_create_test.go | 44 +++++++++-------- server/routes_delete_test.go | 13 +++--- 6 files changed, 161 insertions(+), 35 deletions(-) diff --git a/go.mod b/go.mod index 6807b9b48..2e0c6614c 100644 --- a/go.mod +++ b/go.mod @@ -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 ) diff --git a/integration/utils_test.go b/integration/utils_test.go index 7e1fcc10e..1d44aa852 100644 --- a/integration/utils_test.go +++ b/integration/utils_test.go @@ -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) diff --git a/server/modelpath.go b/server/modelpath.go index 39f42305f..69b152801 100644 --- a/server/modelpath.go +++ b/server/modelpath.go @@ -160,22 +160,50 @@ func migrateRegistryDomain() error { return err } - olddomainpath := filepath.Join(manifests, "registry.ollama.ai") - newdomainpath := filepath.Join(manifests, DefaultRegistry) + 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 + } - return filepath.Walk(olddomainpath, func(path string, info fs.FileInfo, err error) error { - if err != 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", path) - newpath := filepath.Join(newdomainpath, strings.TrimPrefix(path, olddomainpath)) - if err := os.MkdirAll(filepath.Dir(newpath), 0o755); err != nil { + slog.Info("migrating registry domain", "path", source) + + rel, err := filepath.Rel(sourceDomain, source) + if err != nil { return err } - if err := os.Rename(path, newpath); err != nil { + 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 } } diff --git a/server/modelpath_test.go b/server/modelpath_test.go index 6c4dfbee4..6cb5d7e07 100644 --- a/server/modelpath_test.go +++ b/server/modelpath_test.go @@ -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) + } + }) +} diff --git a/server/routes_create_test.go b/server/routes_create_test.go index 340612822..0b57ab03f 100644 --- a/server/routes_create_test.go +++ b/server/routes_create_test.go @@ -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{ diff --git a/server/routes_delete_test.go b/server/routes_delete_test.go index 00303bd17..c2f0be7bc 100644 --- a/server/routes_delete_test.go +++ b/server/routes_delete_test.go @@ -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) }