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 (
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

@ -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
}
}

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"
"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)
}