This commit is contained in:
Michael Yang 2024-05-20 14:04:27 -07:00
parent 8d62a65ca7
commit 4b4e97ed10
6 changed files with 161 additions and 35 deletions

3
go.mod
View File

@ -18,6 +18,7 @@ require (
require ( require (
github.com/agnivade/levenshtein v1.1.1 github.com/agnivade/levenshtein v1.1.1
github.com/d4l3k/go-bfloat16 v0.0.0-20211005043715-690c3bdd05f1 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/mattn/go-runewidth v0.0.14
github.com/nlpodyssey/gopickle v0.3.0 github.com/nlpodyssey/gopickle v0.3.0
github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c 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/net v0.25.0 // indirect
golang.org/x/sys v0.20.0 golang.org/x/sys v0.20.0
golang.org/x/term 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 google.golang.org/protobuf v1.34.1
gopkg.in/yaml.v3 v3.0.1 // indirect 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 { if tmp := os.Getenv("OLLAMA_HOST"); tmp != ollamaHost {
slog.Info("setting env", "OLLAMA_HOST", ollamaHost) slog.Info("setting env", "OLLAMA_HOST", ollamaHost)
t.Setenv("OLLAMA_HOST", ollamaHost) t.Setenv("OLLAMA_HOST", ollamaHost)
envconfig.LoadConfig()
} }
slog.Info("starting server", "url", ollamaHost) slog.Info("starting server", "url", ollamaHost)

View File

@ -160,22 +160,50 @@ func migrateRegistryDomain() error {
return err return err
} }
olddomainpath := filepath.Join(manifests, "registry.ollama.ai") targetDomain := filepath.Join(manifests, DefaultRegistry)
newdomainpath := 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 { sourceDomain := filepath.Join(manifests, "registry.ollama.ai")
if err != nil {
//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 return err
} }
if !info.IsDir() { if !info.IsDir() {
slog.Info("migrating registry domain", "path", path) slog.Info("migrating registry domain", "path", source)
newpath := filepath.Join(newdomainpath, strings.TrimPrefix(path, olddomainpath))
if err := os.MkdirAll(filepath.Dir(newpath), 0o755); err != nil { rel, err := filepath.Rel(sourceDomain, source)
if err != nil {
return err 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 return err
} }
} }

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

@ -14,6 +14,7 @@ import (
"testing" "testing"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/api" "github.com/ollama/ollama/api"
"github.com/ollama/ollama/envconfig" "github.com/ollama/ollama/envconfig"
"github.com/ollama/ollama/llm" "github.com/ollama/ollama/llm"
@ -79,8 +80,11 @@ func checkFileExists(t *testing.T, p string, expect []string) {
t.Fatal(err) t.Fatal(err)
} }
if !slices.Equal(actual, expect) { slices.Sort(actual)
t.Fatalf("expected slices to be equal %v", 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{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -127,7 +131,7 @@ func TestCreateFromModel(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ w = createRequest(t, s.CreateModelHandler, api.CreateRequest{
@ -141,8 +145,8 @@ func TestCreateFromModel(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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"),
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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -168,7 +172,7 @@ func TestCreateRemovesLayers(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -188,7 +192,7 @@ func TestCreateRemovesLayers(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -215,7 +219,7 @@ func TestCreateUnsetsSystem(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -235,7 +239,7 @@ func TestCreateUnsetsSystem(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -271,7 +275,7 @@ func TestCreateMergeParameters(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -292,8 +296,8 @@ func TestCreateMergeParameters(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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"),
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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -330,8 +334,8 @@ func TestCreateMergeParameters(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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"),
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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -374,7 +378,7 @@ func TestCreateReplacesMessages(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -394,8 +398,8 @@ func TestCreateReplacesMessages(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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"),
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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -450,7 +454,7 @@ func TestCreateTemplateSystem(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -496,7 +500,7 @@ func TestCreateLicenses(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ 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{ 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"),
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{ checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
@ -57,7 +57,7 @@ func TestDelete(t *testing.T) {
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{ 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{ 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) t.Fatalf("expected status code 200, actual %d", w.Code)
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{}) checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), nil)
checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{}) checkFileExists(t, filepath.Join(p, "blobs", "*"), nil)
} }
func TestDeleteDuplicateLayers(t *testing.T) { func TestDeleteDuplicateLayers(t *testing.T) {
p := t.TempDir() p := t.TempDir()
t.Setenv("OLLAMA_MODELS", p) t.Setenv("OLLAMA_MODELS", p)
envconfig.LoadConfig()
var s Server var s Server
n := model.ParseName("test") n := model.ParseName("test")
@ -103,5 +104,5 @@ func TestDeleteDuplicateLayers(t *testing.T) {
t.Errorf("expected status code 200, actual %d", w.Code) t.Errorf("expected status code 200, actual %d", w.Code)
} }
checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{}) checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), nil)
} }