From d9dec1509d71b49e17a4cb4f03641fe59aa37cf9 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Mon, 5 Oct 2026 17:34:41 -0700 Subject: [PATCH] Reduce model lookup and MLX decision request overhead (#18806) * Reduce model lookup and MLX decision request overhead Avoid decoding unrelated manifests when resolving existing model names. Reuse small Metal scratch buffers between decision requests, releasing larger working sets and clearing scratch after scoring errors. * address comments --- manifest/manifest.go | 20 ++++++++++++ manifest/manifest_test.go | 68 +++++++++++++++++++++++++++++++++++++++ mlx/stream.go | 5 +++ mlxrunner/score.go | 22 +++++++++++-- mlxrunner/score_test.go | 23 +++++++++---- server/routes.go | 21 ++++++------ server/routes_test.go | 66 +++++++++++++++++++++++++++++++++++++ 7 files changed, 205 insertions(+), 20 deletions(-) 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)