mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
fix(registry): metadata capture gets its own timeout + panic isolation
warmFill/refreshProvider used to reuse listModels' own context for the subsequent captureModelMetadata call — a listModels attempt that took real, non-trivial time (or timed out outright) guaranteed the metadata fetch an already-spent, near-expired context, producing a spurious deadline error and log line every warm fill or refresh for an otherwise healthy provider. Both call sites now give captureModelMetadata its own fresh context.WithTimeout budget, computed after listModels returns. (review SHOULD-4) captureModelMetadata also gains its own deferred recover, isolating a panic inside fetchModelMetadata (or a future modelMetadataFetcher implementation) from refreshProvider's OUTER recover — which would otherwise overwrite the (already-nil, already-correct) err variable from a successful listModels call and cause finishRefresh to discard a freshly fetched, genuinely good model list over an unrelated bug in the metadata path. "Metadata must never fail discovery" is now enforced on the panic path, not just the plain-error-return one already handled. (review SHOULD-5) New regression tests: a fresh-timeout-budget check (a slow simulated listModels must not shrink the metadata fetch's own deadline) and a panic-in-fetchModelMetadata check (discovery must still succeed and record both models).
This commit is contained in:
+36
-4
@@ -401,12 +401,20 @@ func (m *modelRegistry) warmFill(ctx context.Context) {
|
||||
}
|
||||
fctx, cancel := context.WithTimeout(ctx, warmFillTimeout)
|
||||
ids, err := m.adapters[name].listModels(fctx)
|
||||
cancel()
|
||||
st.finishRefresh(m.now(), ids, err)
|
||||
if err != nil {
|
||||
m.log("%s", fmt.Sprintf("model registry: initial discovery for provider %q failed: %v", name, err))
|
||||
}
|
||||
m.captureModelMetadata(fctx, name, m.adapters[name], st)
|
||||
cancel()
|
||||
// Own timeout budget, review fix (SHOULD-4): fctx above is
|
||||
// spent by listModels — reusing it here would hand the metadata
|
||||
// fetch an already-near-expired (or fully expired, on a slow/
|
||||
// timed-out listModels) context every warm fill, guaranteeing a
|
||||
// spurious deadline error and log line even for a healthy
|
||||
// provider whose model-list call merely took a while.
|
||||
mctx, mcancel := context.WithTimeout(ctx, warmFillTimeout)
|
||||
m.captureModelMetadata(mctx, name, m.adapters[name], st)
|
||||
mcancel()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -448,7 +456,24 @@ type modelMetadataFetcher interface {
|
||||
// captured" ruling, it must never fail discovery itself, which
|
||||
// listModels (warmFill/refreshProvider's own caller) already accounts
|
||||
// for independently.
|
||||
//
|
||||
// The deferred recover (review fix, SHOULD-5) is this function's OWN,
|
||||
// separate from refreshProvider's outer recover: without it, a panic
|
||||
// inside fetchModelMetadata would unwind into refreshProvider's deferred
|
||||
// recover instead, which sets its OWN err variable and would then report
|
||||
// listModels' own, already-successful discovery fetch as failed —
|
||||
// discarding a good, freshly-fetched model list over an unrelated bug in
|
||||
// the metadata path. Catching it here, before it can ever reach that
|
||||
// outer recover, is what actually enforces "metadata must never fail
|
||||
// discovery" on the panic path, not just the plain-error-return one
|
||||
// already handled below.
|
||||
func (m *modelRegistry) captureModelMetadata(ctx context.Context, name string, adapter providerAdapter, st *providerState) {
|
||||
defer func() {
|
||||
if rec := recover(); rec != nil {
|
||||
m.log("%s", fmt.Sprintf("model registry: metadata capture for provider %q panicked: %v", name, rec))
|
||||
}
|
||||
}()
|
||||
|
||||
mf, ok := adapter.(modelMetadataFetcher)
|
||||
if !ok {
|
||||
return
|
||||
@@ -523,9 +548,16 @@ func (m *modelRegistry) refreshProvider(name string, st *providerState, adapter
|
||||
}()
|
||||
|
||||
fctx, cancel := context.WithTimeout(context.Background(), backgroundRefreshTimeout)
|
||||
defer cancel()
|
||||
ids, err = adapter.listModels(fctx)
|
||||
m.captureModelMetadata(fctx, name, adapter, st)
|
||||
cancel()
|
||||
|
||||
// Own timeout budget, review fix (SHOULD-4) — see warmFill's
|
||||
// identical comment: reusing fctx here would hand the metadata fetch
|
||||
// an already-spent context on every refresh whose listModels call
|
||||
// ran long, not just a slow or timed-out one.
|
||||
mctx, mcancel := context.WithTimeout(context.Background(), backgroundRefreshTimeout)
|
||||
m.captureModelMetadata(mctx, name, adapter, st)
|
||||
mcancel()
|
||||
}
|
||||
|
||||
// splitConfiguredProvider reports whether id has "prefix/rest" form where
|
||||
|
||||
@@ -288,3 +288,90 @@ func TestModelRegistry_ResolveMetaForAliasName(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestModelRegistry_WarmFill_MetadataGetsFreshTimeoutBudget is the
|
||||
// review fix (SHOULD-4) regression: the metadata fetch must get its OWN
|
||||
// context.WithTimeout budget, computed AFTER listModels finishes — not a
|
||||
// reuse of listModels' own, already-partially-spent context. A
|
||||
// listModels call that takes real, non-trivial time must not leave the
|
||||
// metadata fetch with a shrunken remaining budget.
|
||||
func TestModelRegistry_WarmFill_MetadataGetsFreshTimeoutBudget(t *testing.T) {
|
||||
t.Parallel()
|
||||
const simulatedListModelsDelay = 800 * time.Millisecond
|
||||
const slack = 400 * time.Millisecond // generous vs. the 800ms delay — see the assertion's own comment
|
||||
|
||||
var metadataCtxDeadline time.Time
|
||||
var sawDeadline bool
|
||||
|
||||
fa := newFakeMetaAdapter("lmstudio")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) {
|
||||
time.Sleep(simulatedListModelsDelay)
|
||||
return []string{"m1"}, nil
|
||||
}
|
||||
fa.fetchModelMetadataFn = func(ctx context.Context) (map[string]int, error) {
|
||||
metadataCtxDeadline, sawDeadline = ctx.Deadline()
|
||||
return map[string]int{"m1": 4096}, nil
|
||||
}
|
||||
adapters := map[string]providerAdapter{"lmstudio": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"lmstudio": {Discovery: true, MetadataPath: "/api/v0/models"}}}
|
||||
reg, err := newModelRegistry(adapters, cfg, func(string, ...any) {})
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
before := time.Now()
|
||||
reg.warmFill(context.Background())
|
||||
if !sawDeadline {
|
||||
t.Fatal("fetchModelMetadataFn never ran, or its context carried no deadline")
|
||||
}
|
||||
|
||||
// Old (buggy) behavior would reuse listModels' own context, whose
|
||||
// warmFillTimeout budget started BEFORE the simulated delay — its
|
||||
// deadline would land ~800ms EARLIER than a freshly computed one.
|
||||
// The fixed behavior computes a new deadline AFTER listModels
|
||||
// returns, so it lands close to (start-of-warmFill + delay +
|
||||
// warmFillTimeout), not (start-of-warmFill + warmFillTimeout).
|
||||
wantEarliest := before.Add(warmFillTimeout - slack)
|
||||
if metadataCtxDeadline.Before(wantEarliest) {
|
||||
t.Errorf("metadata fetch context deadline = %v, want at or after %v — it must get a fresh warmFillTimeout budget, not listModels' spent one", metadataCtxDeadline, wantEarliest)
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelRegistry_RefreshProvider_MetadataPanicDoesNotFailDiscovery is
|
||||
// the review fix (SHOULD-5) regression: a panic inside fetchModelMetadata
|
||||
// must never discard a listModels call that already succeeded.
|
||||
// captureModelMetadata's own recover (registry.go) must catch it before
|
||||
// it ever reaches refreshProvider's outer recover, which would otherwise
|
||||
// overwrite the (already-nil, already-correct) err variable and cause
|
||||
// finishRefresh to treat a genuinely successful discovery as failed.
|
||||
func TestModelRegistry_RefreshProvider_MetadataPanicDoesNotFailDiscovery(t *testing.T) {
|
||||
t.Parallel()
|
||||
fa := newFakeMetaAdapter("lmstudio")
|
||||
fa.listModelsFn = func(context.Context) ([]string, error) {
|
||||
return []string{"m1", "m2"}, nil
|
||||
}
|
||||
fa.fetchModelMetadataFn = func(context.Context) (map[string]int, error) {
|
||||
panic("simulated bug in fetchModelMetadata")
|
||||
}
|
||||
adapters := map[string]providerAdapter{"lmstudio": fa}
|
||||
cfg := &Config{Providers: map[string]*ProviderConfig{"lmstudio": {Discovery: true, DiscoveryInterval: "1h", MetadataPath: "/api/v0/models"}}}
|
||||
log := &recordingLog{}
|
||||
reg, err := newModelRegistry(adapters, cfg, log.fn)
|
||||
if err != nil {
|
||||
t.Fatalf("newModelRegistry: %v", err)
|
||||
}
|
||||
|
||||
reg.maybeRefresh(context.Background())
|
||||
waitUntil(t, time.Second, func() bool { return reg.states["lmstudio"].hasModel("m1") })
|
||||
|
||||
if !reg.states["lmstudio"].hasModel("m2") {
|
||||
t.Error("a panic in metadata capture must not discard the successful discovery result")
|
||||
}
|
||||
_, _, lastErr := reg.states["lmstudio"].snapshot()
|
||||
if lastErr != "" {
|
||||
t.Errorf("lastErr = %q, want empty — discovery itself succeeded; only metadata capture panicked", lastErr)
|
||||
}
|
||||
if log.count() == 0 {
|
||||
t.Error("want at least one log line recording the metadata-capture panic")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user