diff --git a/manifest/manifest.go b/manifest/manifest.go index be54ce12a..61c808324 100644 --- a/manifest/manifest.go +++ b/manifest/manifest.go @@ -1564,6 +1564,26 @@ func checkBlobDigestReader(r io.Reader, digest string) error { return nil } +// FindName returns the stored spelling of a case-insensitive name match, or +// os.ErrNotExist. Like Manifests, it ignores invalid manifests and prefers v2 +// entries, but only reads manifest contents for matching names. +func FindName(n model.Name) (model.Name, error) { + refs, err := namedManifestRefs(true) + if err != nil { + return model.Name{}, err + } + + for name, ref := range refs { + if !name.EqualFold(n) { + continue + } + if _, err := parseManifestFile(name, ref.path, ref.root, runnerPreferences()); err == nil { + return name, nil + } + } + return model.Name{}, os.ErrNotExist +} + func Manifests(continueOnError bool) (map[model.Name]*Manifest, error) { refs, err := namedManifestRefs(continueOnError) if err != nil { diff --git a/manifest/manifest_test.go b/manifest/manifest_test.go index 765a6795d..0fa72b515 100644 --- a/manifest/manifest_test.go +++ b/manifest/manifest_test.go @@ -146,6 +146,9 @@ func TestV2ManifestPathEscapesHostPort(t *testing.T) { if err := WriteManifestData(name, data); err != nil { t.Fatal(err) } + if got, err := FindName(name); err != nil || got.String() != name.String() { + t.Fatalf("FindName = %v, %v; want %v, nil", got, err, name) + } manifestPath, err := V2PathForName(name) if err != nil { @@ -730,6 +733,9 @@ func TestPartialManifestListTracksPresentAndMissingChildren(t *testing.T) { if err := WriteManifestData(name, parentData); err != nil { t.Fatal(err) } + if got, err := FindName(name); err != nil || got.String() != name.String() { + t.Fatalf("FindName = %v, %v; want %v, nil", got, err, name) + } parentSum := sha256.Sum256(parentData) parentDigest := fmt.Sprintf("sha256:%x", parentSum) @@ -1051,6 +1057,9 @@ func TestParseNamedManifestRejectsUnsafeSymlinks(t *testing.T) { if err == nil || !strings.Contains(err.Error(), "not a sha256 blob") { t.Fatalf("err = %v, want not a sha256 blob", err) } + if _, err := FindName(name); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("FindName error = %v, want os.ErrNotExist", err) + } }) t.Run("blob basename outside blob store", func(t *testing.T) { @@ -1071,6 +1080,9 @@ func TestParseNamedManifestRejectsUnsafeSymlinks(t *testing.T) { if err == nil || !strings.Contains(err.Error(), "does not match blob") { t.Fatalf("err = %v, want does not match blob", err) } + if _, err := FindName(name); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("FindName error = %v, want os.ErrNotExist", err) + } }) } @@ -1138,6 +1150,62 @@ func TestManifestsV2ShadowsLegacy(t *testing.T) { } } +func TestFindName(t *testing.T) { + for _, tc := range []struct { + name string + legacy string + v2 string + found bool + }{ + {name: "missing"}, + {name: "legacy", legacy: `{}`, found: true}, + {name: "v2", v2: `{}`, found: true}, + {name: "v2 shadows corrupt legacy", legacy: `{`, v2: `{}`, found: true}, + {name: "corrupt legacy", legacy: `{`}, + {name: "corrupt v2 shadows legacy", legacy: `{}`, v2: `{`}, + { + name: "manifest list", found: true, + v2: `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp","config":{"digest":"sha256:abc"}}]}`, + }, + { + name: "manifest list with missing child", + v2: `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp","digest":"sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}]}`, + }, + { + name: "unsupported runner", + v2: `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"unsupported","config":{"digest":"sha256:abc"}}]}`, + }, + { + name: "nested manifest list", + v2: `{"mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp","mediaType":"application/vnd.ollama.manifest.list.v2+json","manifests":[{"runner":"llamacpp"}]}]}`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + name := model.ParseName("MyOrg/MyModel:Q8") + if tc.legacy != "" { + if err := WriteLegacyManifestData(name, []byte(tc.legacy)); err != nil { + t.Fatal(err) + } + } + if tc.v2 != "" { + digest := writeManifestBlobForTest(t, []byte(tc.v2)) + if err := linkManifest(name, digest); err != nil { + t.Fatal(err) + } + } + got, err := FindName(model.ParseName("myorg/mymodel:q8")) + if tc.found { + if err != nil || got.String() != name.String() { + t.Fatalf("FindName = %v, %v; want %v, nil", got, err, name) + } + } else if !errors.Is(err, os.ErrNotExist) { + t.Fatalf("FindName error = %v, want os.ErrNotExist", err) + } + }) + } +} + func TestManifests(t *testing.T) { cases := map[string]struct { ps []string diff --git a/mlx/stream.go b/mlx/stream.go index ec5e81f1f..082d7679e 100644 --- a/mlx/stream.go +++ b/mlx/stream.go @@ -60,6 +60,11 @@ type Stream struct { ctx C.mlx_stream } +// Synchronize waits for submitted work and its completion handlers on the stream. +func (s Stream) Synchronize() { + mlxCheck(C.mlx_synchronize(s.ctx)) +} + func (s Stream) LogValue() slog.Value { str := mlxCheck(C.mlx_string_new()) mlxCheck(C.mlx_stream_tostring(&str, s.ctx)) diff --git a/mlxrunner/score.go b/mlxrunner/score.go index b6cdef4e1..744b82fe9 100644 --- a/mlxrunner/score.go +++ b/mlxrunner/score.go @@ -22,6 +22,10 @@ type scoreRow struct { candidates []int32 } +// Retain enough scratch space for short decision requests without keeping the +// temporary working set of a large prompt resident between requests. +const scoreScratchLimit = 512 << 20 + // Models may project just the requested vocabulary rows at higher precision. type candidateUnembedder interface { UnembedCandidates(hidden, candidates *mlx.Array) *mlx.Array @@ -34,8 +38,22 @@ func (r *Runner) scoreHandler(w http.ResponseWriter, req *http.Request) { return } result, err := mlxthread.Call(req.Context(), r.mlxThread, func() (llm.ScoreResponse, error) { - defer mlx.ClearCache() - return r.score(req.Context(), input) + if !mlx.MetalIsAvailable() { + defer mlx.ClearCache() + return r.score(req.Context(), input) + } + keepScratch := false + defer func() { + // Keep small reusable buffers, but release large working sets while idle. + if !keepScratch || mlx.CacheMemory() > scoreScratchLimit { + mlx.ClearCache() + } + }() + result, err := r.score(req.Context(), input) + // Eval waits for results, but completion handlers can still free buffers. + mlx.DefaultStream().Synchronize() + keepScratch = err == nil + return result, err }) if err != nil { status := http.StatusInternalServerError diff --git a/mlxrunner/score_test.go b/mlxrunner/score_test.go index 725a6ac68..3867f8d43 100644 --- a/mlxrunner/score_test.go +++ b/mlxrunner/score_test.go @@ -200,7 +200,7 @@ func TestScoreValidation(t *testing.T) { } } -func TestScoreHandlerClearsScratchCache(t *testing.T) { +func TestScoreHandlerScratchCache(t *testing.T) { worker := mlxtest.Worker(t) if err := worker.Do(context.Background(), func() error { if !mlx.GPUIsAvailable() { @@ -211,11 +211,19 @@ func TestScoreHandlerClearsScratchCache(t *testing.T) { }); err != nil { t.Skipf("MLX GPU not available: %v", err) } + t.Cleanup(func() { + if err := worker.Do(context.Background(), func() error { + mlx.ClearCache() + return nil + }); err != nil { + t.Error(err) + } + }) r := &Runner{Tokenizer: newTestTokenizer(t, []int32{7}), contextLength: 512, mlxThread: worker} defer worker.Do(context.Background(), func() error { r.Close(); return nil }) for _, n := range []int{16, 128, 512, 16} { for _, cancelled := range []bool{false, true} { - // Model prior requests leaving differently sized scratch buffers. + // Start with reusable scratch from an earlier request. // These CPU-filled arrays are freed synchronously into MLX's cache. const scratchBytes = 4 << 20 if err := worker.Do(context.Background(), func() error { @@ -250,14 +258,17 @@ func TestScoreHandlerClearsScratchCache(t *testing.T) { t.Fatalf("tokens=%d cancelled=%t: status %d: %s", n, cancelled, response.Code, response.Body) } memory, err := mlxthread.Call(context.Background(), worker, func() (int, error) { - return mlx.CacheMemory(), nil + memory := mlx.CacheMemory() + if !cancelled && mlx.MetalIsAvailable() && memory == 0 { + return 0, errors.New("successful scoring discarded all reusable scratch buffers") + } + return memory, nil }) if err != nil { t.Fatal(err) } - // The injected scratch buffer must be cleared. Active allocations - // belong to live model caches or other tests on the shared worker. - if memory >= scratchBytes { + // Successful requests may reuse scratch; cancelled requests release it. + if memory > scoreScratchLimit || (cancelled && memory >= scratchBytes) { t.Fatalf("tokens=%d cancelled=%t: retained %d scratch bytes", n, cancelled, memory) } } diff --git a/server/routes.go b/server/routes.go index 73b9100f1..d4ed69897 100644 --- a/server/routes.go +++ b/server/routes.go @@ -1445,22 +1445,19 @@ func getExistingName(n model.Name) (model.Name, error) { return n, nil } + // The usual request names an existing model. Avoid decoding every installed + // manifest just to resolve its spelling; partial matches below are for new names. + if name, err := manifest.FindName(n); err == nil { + return name, nil + } else if !errors.Is(err, os.ErrNotExist) { + return zero, err + } + existing, err := manifest.Manifests(true) if err != nil { return zero, err } - // First pass: look for a full case-insensitive match across all four - // parts. If found, return the on-disk canonical name directly. - for e := range existing { - if strings.EqualFold(e.Host, n.Host) && - strings.EqualFold(e.Namespace, n.Namespace) && - strings.EqualFold(e.Model, n.Model) && - strings.EqualFold(e.Tag, n.Tag) { - return e, nil - } - } - - // Second pass: find the single manifest with the longest consecutive + // Find the single manifest with the longest consecutive // case-insensitive prefix match (host -> namespace -> model) and copy // only the matching prefix parts from that manifest. The tag is left // as-is so that an unrelated manifest with a matching tag cannot diff --git a/server/routes_test.go b/server/routes_test.go index 32b480619..1d82124cf 100644 --- a/server/routes_test.go +++ b/server/routes_test.go @@ -34,6 +34,72 @@ import ( "github.com/ollama/ollama/version" ) +func TestGetExistingName(t *testing.T) { + t.Setenv("OLLAMA_MODELS", t.TempDir()) + for _, name := range []string{"MyOrg/MyModel:Q4", "OtherOrg/OtherModel:Q8", "DefaultModel:latest", "localhost:12345/MyOrg/MyModel:Q4"} { + if err := manifest.WriteManifestData(model.ParseName(name), []byte(`{}`)); err != nil { + t.Fatal(err) + } + } + for _, tc := range []struct{ input, want string }{ + {"MyOrg/MyModel:Q4", "MyOrg/MyModel:Q4"}, + {"MYORG/MYMODEL:q4", "MyOrg/MyModel:Q4"}, + {"myorg/mymodel:q8", "MyOrg/MyModel:q8"}, + {"myorg/newmodel:q8", "MyOrg/newmodel:q8"}, + {"neworg/newmodel:q8", "neworg/newmodel:q8"}, + {"defaultmodel", "DefaultModel:latest"}, + {"registry.ollama.ai/myorg/mymodel:q4", "MyOrg/MyModel:Q4"}, + {"ollama.com/myorg/mymodel:q4", "ollama.com/myorg/mymodel:q4"}, + {"registry.ollama.com/myorg/mymodel:q4", "registry.ollama.com/myorg/mymodel:q4"}, + {"LOCALHOST:12345/myorg/mymodel:q4", "localhost:12345/MyOrg/MyModel:Q4"}, + {"sha256-" + strings.Repeat("a", 64), "sha256-" + strings.Repeat("a", 64)}, + } { + t.Run(tc.input, func(t *testing.T) { + got, err := getExistingName(model.ParseName(tc.input)) + if err != nil { + t.Fatal(err) + } + if want := model.ParseName(tc.want); got.String() != want.String() { + t.Errorf("getExistingName = %v, want %v", got, want) + } + }) + } +} + +func BenchmarkGetExistingName(b *testing.B) { + // Tensor manifests are large; lookup should only decode the requested model. + layers := make([]manifest.Layer, 512) + for i := range layers { + layers[i] = manifest.Layer{ + MediaType: manifest.MediaTypeImageTensor, + Digest: fmt.Sprintf("sha256:%064x", i), + Size: 1024, + } + } + data, err := json.Marshal(manifest.Manifest{SchemaVersion: 2, Layers: layers}) + if err != nil { + b.Fatal(err) + } + for _, models := range []int{1, 10, 100} { + b.Run(fmt.Sprintf("models=%d", models), func(b *testing.B) { + b.Setenv("OLLAMA_MODELS", b.TempDir()) + for i := range models { + name := model.ParseName(fmt.Sprintf("Test/Model%d:Latest", i)) + if err := manifest.WriteManifestData(name, data); err != nil { + b.Fatal(err) + } + } + name := model.ParseName("test/model0:latest") + b.ReportAllocs() + for b.Loop() { + if _, err := getExistingName(name); err != nil { + b.Fatal(err) + } + } + }) + } +} + func TestPsHandlerUsesRunningManifestAndRunner(t *testing.T) { gin.SetMode(gin.TestMode)