fix(inferflow): rebuild ComponentProviderHandler.componentMap per RegisterComponent call#382
fix(inferflow): rebuild ComponentProviderHandler.componentMap per RegisterComponent call#382anujGoel124 wants to merge 1 commit into
Conversation
…isterComponent call
`RegisterComponent` is registered as an etcd config-change callback
(cmd/inferflow/main.go:27 — via ReloadModelConfigMapAndRegisterComponents),
so it runs once per config event for the pod's lifetime. The previous
implementation APPENDED to `cp.componentMap` without ever rebuilding or
pruning. Entries from removed/renamed models leaked, and because etcd
config events arrive throughout the pod's lifetime, the map grew
unboundedly.
The fix is the same atomic-swap pattern already used in this codebase for
`featureSchemaCache` (handlers/config/config_schema.go:32-42): build a fresh
map locally from the supplied ModelConfig, then publish it under the write
lock. The old map becomes GC-able. Readers (GetComponent) are unaffected —
they continue to take the read lock and see either the old or the new map,
never a torn intermediate state.
The function's external contract is unchanged: after `RegisterComponent(cfg)`
returns, `cp.componentMap` reflects exactly the components in `cfg` (plus the
always-present `feature_initializer`). The difference is that prior calls
no longer accumulate.
Added focused tests in provider_test.go:
- rebuild-not-accumulate: two calls with different configs leave only the
second config's components in the map
- stable-size-across-repeated-calls: 100 identical RegisterComponent calls
keep map size constant (the simplest regression for the leak)
- non-ModelConfig input is a no-op
- empty ModelConfig retains only the feature_initializer entry
go test ./... in the inferflow module is green.
🟢 LOW — safe to fast-track · confidence 0.88This fixes an internal memory leak by rebuilding the component registry on each config reload instead of endlessly appending, without changing any external API. Top risks
Check before approving
8 tool(s) ran · DRS agentic v2 · #382 |
![m-agentic-review[bot]](https://avatars.githubusercontent.com/u/19436680?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 Code Review — PR #382
Language: Go | Files reviewed: 2
| Severity | Count | |
|---|---|---|
| BLOCKER | 🚨 | 1 |
| NITPICK | 🔧 | 1 |
| SUGGESTION | 💡 | 1 |
| WARNING | 5 |
| } | ||
|
|
||
| cp.mapMutex.Lock() | ||
| cp.componentMap = newMap |
There was a problem hiding this comment.
🚨 BLOCKER — Replacing the entire component map from each callback payload will drop components for other models if the etcd callback supplies only the changed model rather than a full snapshot.
The new code builds newMap only from modelConfig.ConfigMap in the current RegisterComponent call and then assigns it over cp.componentMap. The old append behavior preserved components registered by prior callbacks, so if callbacks are per-key or partial updates, this change removes still-live components that were not present in the latest event payload.
Also flagged on this line:
⚠️ WARNING — A nil or empty*ModelConfignow wipes all existing registered components instead of leaving the previous map intact.
| } | ||
| } | ||
|
|
||
| cp.mapMutex.Lock() |
There was a problem hiding this comment.
The write lock now only protects the final assignment, so two concurrent callbacks can build maps in parallel and publish in completion order rather than invocation/event order. A slow earlier callback can swap its stale map after a later callback has already published the current config.
Also flagged on this line:
⚠️ WARNING — Delaying the write lock until after the rebuild allows concurrent registrations to publish stale configs out of order.
| // Build a fresh map locally; never mutate cp.componentMap while readers | ||
| // (GetComponent) may be holding the read lock. The final atomic swap | ||
| // publishes the new map; the old map becomes GC-able. | ||
| newMap := make(map[string]dag.AbstractComponent) |
There was a problem hiding this comment.
With the old code, once RegisterComponent entered the write-locked section, GetComponent calls blocked until the newly registered components were present. The new code intentionally keeps readers unblocked while rebuilding, so requests can continue using the old component map during the rebuild window, which can be inconsistent if the DAG/model config has already been updated elsewhere.
Also flagged on this line:
⚠️ WARNING — Consider serializing or coalescing concurrentRegisterComponentrebuilds so config-event bursts do not build multiple full component maps at once.- 💡 SUGGESTION — Pre-size
newMapto the expected component count before filling it on each config event.
| } | ||
| } | ||
|
|
||
| // TestRegisterComponent_NilModelConfigStillWritesInitializer covers the edge |
There was a problem hiding this comment.
🔧 NITPICK — Rename this test comment to match TestRegisterComponent_EmptyModelConfigKeepsInitializer or change the test to actually pass a typed nil *config.ModelConfig.
The comment mentions a nil ModelConfig, but the test passes a non-nil ModelConfig with a nil ConfigMap, which makes the covered edge case unclear to future maintainers.
|
Closing this PR after a more rigorous audit. The
The dominant source of the +50 MB/h slope isn't visible to static
Shipping a 1% fix to another team's repo without confidence in the If the inferflow team is interested in collaborating on the actual |
1 participant
Context
ComponentProviderHandler.RegisterComponentis registered as an etcdconfig-change callback (see
cmd/inferflow/main.go:27—etcd.RegisterWatchPathCallback("", ReloadModelConfigMapAndRegisterComponents)).Every config event invokes
RegisterComponent(updatedConfig).The previous implementation appended to
cp.componentMapon every callwithout rebuilding or pruning. Over the pod's lifetime the map accumulated
entries from every model config that had ever been seen — even after those
models were removed or renamed in etcd. The result is a slow-growth memory
leak driven by config-event rate.
Production telemetry that prompted the fix
Container memory on the inferflow deployments shows monotonic per-pod growth
across multiple pod generations (lifetime ≥ 48 h pods, 7-day window):
prd-inferflow-pdp-adprd-inferflow-rv-ctThese pods aren't in a warm-up ramp — they're well past warm-up steady-state
and still growing.
What this PR fixes
handlers/inferflow/provider.go—RegisterComponentnow builds a freshlocal map from the supplied
*config.ModelConfigand publishes it under thewrite lock. The old map becomes GC-able. Same atomic-swap pattern already
used in this codebase for
featureSchemaCache(
handlers/config/config_schema.go:32-42).func (cp *ComponentProviderHandler) RegisterComponent(request interface{}) { - modelConfig, ok := request.(*config.ModelConfig) - if ok { - cp.mapMutex.Lock() - defer cp.mapMutex.Unlock() - cp.componentMap[featureInitComponent] = &components.FeatureInitComponent{...} - ... - for _, k := range fCompMap.Keys() { - cp.componentMap[k.(string)] = &components.FeatureComponent{...} // unbounded ADD - } - ... - } + modelConfig, ok := request.(*config.ModelConfig) + if !ok { + return + } + newMap := make(map[string]dag.AbstractComponent) + newMap[featureInitComponent] = &components.FeatureInitComponent{...} + ... + for _, k := range fCompMap.Keys() { + newMap[k.(string)] = &components.FeatureComponent{...} + } + ... + cp.mapMutex.Lock() + cp.componentMap = newMap // atomic swap; old map GC-able + cp.mapMutex.Unlock() }Expected runtime impact — honest scoping
This PR fixes one real correctness bug, but is unlikely to flatten the
slope on its own. Magnitude estimate vs. observed slope:
componentMapis roughly a pointer to a Component struct (~200 B).alone is ~0.06 to 0.60 MB/h per pod.
prd-inferflow-pdp-ad.So this PR explains ~1% of the observed slope. The remaining ~99% is
elsewhere — see Likely additional sources below. Validating via Grafana
over 48-72h post-merge is the honest check: if slope drops to a residual
in the 10s-of-MB/h range, this fix is doing its small-but-real part and
the rest needs separate attention.
Tests
TestRegisterComponent_RebuildsNotAccumulates— two calls withdifferent configs leave only the second config's components in the map.
TestRegisterComponent_StableSizeAcrossRepeatedSameCall— 100identical calls keep the map size constant.
*ModelConfiginput is a no-op.ConfigMapkeeps only thefeature_initializerentry.go test ./...in the inferflow module is green.go vet ./...is clean.Likely additional sources (not addressed here — flagging for owners)
Audit of the inferflow source identified these as likely contributors to the
remaining ~99% of the slope. Listed in order of suspected magnitude:
1.
pkg/inmemorycache/init.go:22— Cache initialized without visible size capA constant
IN_MEMORY_CACHE_SIZE_IN_BYTESis defined ininmemory_cache.go:6but I don't see it being passed to
newV1InMemoryCachein the init path. Ifthe underlying cache implementation enforces the size cap via env var
(
IN_MEMORY_CACHE_SIZE_IN_BYTES), confirming the deployment sets that env varwould close this hypothesis. If not, this is likely the dominant remaining
source — high-cardinality cache without eviction.
Suggested followup: verify
IN_MEMORY_CACHE_SIZE_IN_BYTESis set indeployment manifests, or wire it through the init path explicitly.
2.
handlers/external/featurestore/onfs.go:180-212— Per-batch allocation without sync.PoolAt high feature-store RPS, per-batch slice allocations create heavy GC
pressure. Not a "leak" in the unbounded-growth sense, but contributes to
memory growth via heap retention and GC overhead.
Suggested followup: introduce a
sync.Poolof pre-sized buffer slices;reset and reuse rather than allocate per call.
3.
pkg/etcd/v1.go:90-128— WatchPrefix goroutine lifecycleThe watcher goroutine has no shutdown signal; the V1 struct has no
Stop()method. Single long-lived goroutine living for the pod's lifetime — minor
hygiene, not a runtime leak driver, but worth fixing alongside the cache
work.
4.
pkg/etcd/v1.go:136-137—V1.dataMap/V1.metaMappartial-evictionSame accumulator shape as A8, but with partial eviction via the
DELETEhandler. Whether it actually leaks depends on the etcd usage pattern: if
model configs are added (PUT) but old keys are never explicitly
DELETEd(only overwritten),
dataMapaccumulates. Worth confirming the operationalpractice for config additions/removals.
Verification plan
Post-merge, watch per-pod memory slope on
prd-inferflow-pdp-adandprd-inferflow-rv-ctover 48-72h:excellent.
remaining sources (likely items 1-2 above) need separate fixes.
is item 1 or 2 above. Recommend
pprofheap profile on a long-lived podto nail it.
Happy to follow up with separate PRs for items 1-4 if useful.