mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-19 05:47:48 +00:00
perf(registry): cache the encoded /v1/models response per group
GET /v1/models cost 15.9ms, 21.3MB, 552,884 allocs per call at the
live 1,033-model catalog (registry.go's listFor plus routes_unified.go's
json.Encoder), uncached — and OpenAI SDK clients poll this endpoint per
session, so an unchanged catalog was being re-walked and re-marshaled
on every request.
Add modelRegistry.modelsJSON(grp), caching the fully encoded
{object:list,data:[...]} body per (*group, modelsGen). modelsGen
is an atomic counter bumped by a new modelRegistry.finishRefresh
wrapper around providerState.finishRefresh — the only two places that
ever mutate a provider's discovered set or discovery-captured metadata
(warmFill, refreshProvider) already funnel through providerState.
finishRefresh once per refresh cycle, so wrapping it is sufficient to
invalidate the cache on every possible change, including a
metadata-only update on an otherwise-failed refresh (captureModelMetadata
runs unconditionally alongside finishRefresh in both call sites).
Bumped unconditionally rather than only on success, since a failed
listModels fetch can still pair with a successful metadata capture.
Cached per *group pointer: authStore builds every group once, at
construction, and never rebuilds or replaces that map afterward — a
config reload constructs an entirely new Gateway (and modelRegistry)
with its own empty cache, so a *group key can never alias a different
group's catalog either within one registry's lifetime or across a
reload.
routes_unified.go's handleModels call site shrinks to a single write of
modelsJSON's result, replacing the previous inline listFor+json.Encode;
routes_unified.go is otherwise untouched (owned by a sibling change in
this same effort).
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -874,6 +875,117 @@ func TestModelRegistry_ListFor_Collision_LosingProviderStillReachableViaPrefixed
|
||||
}
|
||||
}
|
||||
|
||||
// --- modelsJSON: cached encoded /v1/models body (perf finding 3) ---
|
||||
|
||||
// TestModelRegistry_ModelsJSON_CachesUntilFinishRefresh proves the perf
|
||||
// fix: a group's encoded body is the SAME []byte (not merely
|
||||
// byte-equal — literally not re-marshaled) across repeated calls, until
|
||||
// finishRefresh actually records a refresh attempt for some provider
|
||||
// (modelsGen bumps), at which point the next call reflects the change.
|
||||
func TestModelRegistry_ModelsJSON_CachesUntilFinishRefresh(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"gpt-test"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grp := &group{name: "all"}
|
||||
first := reg.modelsJSON(grp)
|
||||
if len(first) == 0 {
|
||||
t.Fatal("modelsJSON returned an empty body")
|
||||
}
|
||||
second := reg.modelsJSON(grp)
|
||||
if len(first) == 0 || len(second) == 0 || &first[0] != &second[0] {
|
||||
t.Error("two calls between refreshes returned different backing arrays, want the identical cached []byte")
|
||||
}
|
||||
if !strings.Contains(string(first), "gpt-test") {
|
||||
t.Errorf("body = %s, want it to contain %q", first, "gpt-test")
|
||||
}
|
||||
|
||||
// A discovery refresh — even for a provider whose own discovery is
|
||||
// disabled here, driven directly via the registry's own finishRefresh
|
||||
// wrapper rather than a real background goroutine — must invalidate
|
||||
// the cache.
|
||||
reg.finishRefresh(reg.states["openai"], reg.now(), []string{"gpt-test", "gpt-new"}, nil)
|
||||
|
||||
third := reg.modelsJSON(grp)
|
||||
if string(third) == string(second) {
|
||||
t.Error("modelsJSON body unchanged after finishRefresh added a new model, want it to reflect the new catalog")
|
||||
}
|
||||
if !strings.Contains(string(third), "gpt-new") {
|
||||
t.Errorf("body after refresh = %s, want it to contain the newly discovered %q", third, "gpt-new")
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelRegistry_ModelsJSON_PerGroupIsolation proves the cache never
|
||||
// leaks one group's catalog to another: two groups authorized for
|
||||
// disjoint providers must each get their own encoded body, keyed by
|
||||
// *group pointer identity (modelsJSON's own doc comment).
|
||||
func TestModelRegistry_ModelsJSON_PerGroupIsolation(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{
|
||||
"openai": newFakeAdapter("openai"),
|
||||
"anthropic": newFakeAdapter("anthropic"),
|
||||
}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{
|
||||
"openai": {Models: []string{"gpt-test"}},
|
||||
"anthropic": {Models: []string{"claude-x"}},
|
||||
}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grpOpenAI := &group{name: "openai-only", providers: []string{"openai"}}
|
||||
grpAnthropic := &group{name: "anthropic-only", providers: []string{"anthropic"}}
|
||||
|
||||
bodyOpenAI := reg.modelsJSON(grpOpenAI)
|
||||
bodyAnthropic := reg.modelsJSON(grpAnthropic)
|
||||
|
||||
if strings.Contains(string(bodyOpenAI), "claude-x") {
|
||||
t.Errorf("openai-only group's body leaked anthropic's model: %s", bodyOpenAI)
|
||||
}
|
||||
if strings.Contains(string(bodyAnthropic), "gpt-test") {
|
||||
t.Errorf("anthropic-only group's body leaked openai's model: %s", bodyAnthropic)
|
||||
}
|
||||
if !strings.Contains(string(bodyOpenAI), "gpt-test") {
|
||||
t.Errorf("openai-only group's body = %s, want it to contain gpt-test", bodyOpenAI)
|
||||
}
|
||||
if !strings.Contains(string(bodyAnthropic), "claude-x") {
|
||||
t.Errorf("anthropic-only group's body = %s, want it to contain claude-x", bodyAnthropic)
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelRegistry_ModelsJSON_MatchesUncachedListFor proves modelsJSON's
|
||||
// output is byte-identical to the OLD, pre-cache call site
|
||||
// (json.NewEncoder(w).Encode(...), routes_unified.go before perf finding
|
||||
// 3) for the same envelope over listFor's own result — the cache changes
|
||||
// WHEN the body is computed, never WHAT it computes, including
|
||||
// json.Encoder.Encode's own trailing '\n' that a bare json.Marshal does
|
||||
// not add (review fix, Should-Fix 5 — modelsJSON appends it explicitly).
|
||||
func TestModelRegistry_ModelsJSON_MatchesUncachedListFor(t *testing.T) {
|
||||
t.Parallel()
|
||||
adapters := map[string]providerAdapter{"openai": newFakeAdapter("openai")}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"openai": {Models: []string{"zeta", "alpha"}}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
grp := allowAllGroup()
|
||||
got := reg.modelsJSON(grp)
|
||||
|
||||
var want bytes.Buffer
|
||||
if err := json.NewEncoder(&want).Encode(map[string]any{"object": "list", "data": reg.listFor(grp)}); err != nil {
|
||||
t.Fatalf("json.NewEncoder.Encode: %v", err)
|
||||
}
|
||||
if string(got) != want.String() {
|
||||
t.Errorf("modelsJSON = %q, want %q (byte-identical to the old json.Encoder.Encode call site)", got, want.String())
|
||||
}
|
||||
}
|
||||
|
||||
// --- end-to-end: GET /v1/models through Gateway.ServeHTTP ---
|
||||
|
||||
func TestGateway_ServeHTTP_GetV1Models_ReturnsGroupFilteredList(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user