Michael Yang 11 maanden geleden
bovenliggende
commit
4b4e97ed10
6 gewijzigde bestanden met toevoegingen van 161 en 35 verwijderingen
  1. 2 1
      go.mod
  2. 1 0
      integration/utils_test.go
  3. 36 8
      server/modelpath.go
  4. 91 0
      server/modelpath_test.go
  5. 24 20
      server/routes_create_test.go
  6. 7 6
      server/routes_delete_test.go

+ 2 - 1
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
 )

+ 1 - 0
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)

+ 36 - 8
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
+			}
+
+			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(path, newpath); err != nil {
+			if err := os.Rename(source, target); err != nil {
 				return err
 			}
 		}

+ 91 - 0
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)
+		}
+	})
+}

+ 24 - 20
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{

+ 7 - 6
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)
 }