fix(tools): prevent OOM with nested $defs in tool schemas#19112
Conversation
) Remove the "flatten defs" loop that pre-expanded each $def using unpack_defs(). When definitions reference each other (common pattern), this caused exponential memory growth because each subsequent call copied already-expanded content. Root cause: The loop was calling unpack_defs(value, defs_copy) for each def, but since defs often reference each other, each call would copy increasingly large expanded structures via deepcopy. Fix: Call unpack_defs() only on the parameters, letting it handle nested refs recursively with proper circular reference detection. This resolves the same refs without the exponential memory explosion. Testing showed: - Before fix: 65KB schema → 2.8GB (35,000x expansion, then OOM) - After fix: 65KB schema → 17MB (reasonable 260x expansion) Changes: - factory.py: Remove flatten loop in _bedrock_tools_pt() - vertex_ai/common_utils.py: Remove flatten loop in _build_vertex_schema() - Add regression test for OOM with nested refs
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Hey! I'm the author of #19098. Your fix will definitely help, and I have another suggestion - memoize the unpack_defs function so it caches refs and doesn't keep copying them into memory. If the schema is significantly nested and self-referential, memoizing saves a ton of memory |
|
@rsp2k @DLakin01 - we're trying to do a better job with reducing OOM issues. Would you be open to working with us/giving us advice on how we can get this under control ? If yes - I'd love to get your advice, sharing a link to my cal for your convenience |
|
@rsp2k is this ready for review? your pr status shows draft |
|
In reply to @DLakin01's memoization suggestion: Great suggestion! I explored this in a separate branch ( TL;DRThe OOM fix (removing the flatten loop) is the right solution. Memoization doesn't provide additional benefit due to how Why Memoization Is Challenging HereI tried several memoization approaches, but they all hit the same fundamental issue: 1. Context-Dependent ExpansionThe expansion result depends on the If we cache the "expanded Expression", it contains a 2. Each Reference Needs Its Own CopyEven if we cache expanded defs, we still need 3. Benchmarks Confirm No Benefit
The two-phase memoization approach (pre-expand all defs, then resolve from cache) actually does more work because it deepcopies twice. The Real FixThe OOM was caused by this loop that pre-expanded defs into defs: # PROBLEMATIC - caused O(n!) memory growth:
for _, value in defs_copy.items():
unpack_defs(value, defs_copy) # Each iteration expands already-expanded content!When defs reference each other (A→B→C), each iteration deepcopied already-expanded content, causing exponential growth. Removing this loop fixes the issue completely:
Future Optimization IdeasIf further optimization is ever needed, possibilities include:
But with the current fix, performance is already good (~120KB peak, <1ms per schema), so these aren't necessary. Thanks again for the suggestion - it led to a deeper understanding of why the current approach works well! 🙏 |
krrishdholakia
changed the base branch from
main
to
litellm_staging_01_16_2026
krrishdholakia
merged commit 5a9f6e9
into
BerriAI:litellm_staging_01_16_2026
4 participants
Summary
unpack_defs()already handles nested refs recursively with circular detectionTest Plan
test_bedrock_tools_unpack_defs_no_oom_with_nested_refsCloses #19098