feature: generic object pool (#14702)

* add: generic object pool & tests

Introduced a reusable object pool that can be applied across the codebase.
Note: memory growth is managed via eviction settings—using a hard cap could
reduce performance, so eviction is the preferred safeguard.

* fix: simpler tests

Author: AlexsanderHamir

litellm/litellm_core_utils/object_pooling.py

@@ -0,0 +1,136 @@
+"""
+Generic object pooling utilities for LiteLLM.
+
+This module provides a flexible object pooling system that can be used
+to pool any type of object, reducing memory allocation overhead and
+improving performance for frequently created/destroyed objects.
+
+Memory Management Strategy:
+- Balanced eviction-based memory control to optimize reuse ratio
+- Moderate eviction frequency (300s) to maintain high object reuse
+- Conservative eviction weight (0.3) to avoid destroying useful objects
+- Lower pre-warm count (5) to reduce initial memory footprint
+- Always keeps at least one object available for high availability
+- Unlimited pools when maxsize is not specified (eviction controls actual usage)
+"""
+
+from typing import Type, TypeVar, Optional, Callable
+from pond import Pond, PooledObjectFactory, PooledObject
+
+T = TypeVar('T')
+
+class GenericPooledObjectFactory(PooledObjectFactory):
+    """Generic factory class for creating pooled objects of any type."""
+    
+    def __init__(
+        self, 
+        object_class: Type[T], 
+        pooled_maxsize: Optional[int] = None,  # None = unlimited pool with eviction-based memory control
+        least_one: bool = True,  # Always keep at least one for high concurrency
+        initializer: Optional[Callable[[T], None]] = None
+    ):
+        # Only pass maxsize to Pond if user specified it - otherwise let Pond handle unlimited pools
+        if pooled_maxsize is not None:
+            super().__init__(pooled_maxsize=pooled_maxsize, least_one=least_one)
+        else:
+            super().__init__(least_one=least_one)
+        self.object_class = object_class
+        self.initializer = initializer
+        self._user_maxsize = pooled_maxsize  # Store original user preference
+    
+    def createInstance(self) -> PooledObject:
+        """Create a new instance wrapped in a PooledObject."""
+        # Create a properly initialized instance
+        obj = self.object_class()
+        return PooledObject(obj)
+    
+    def destroy(self, pooled_object: PooledObject):
+        """Destroy the pooled object."""
+        if hasattr(pooled_object.keeped_object, '__dict__'):
+            pooled_object.keeped_object.__dict__.clear()
+        del pooled_object
+    
+    def reset(self, pooled_object: PooledObject) -> PooledObject:
+        """Reset the pooled object to a clean state."""
+        obj = pooled_object.keeped_object
+        # Reset the object by calling its reset method if it exists
+        if hasattr(obj, 'reset') and callable(getattr(obj, 'reset')):
+            obj.reset()
+        else:
+            # Fallback: clear all attributes to reset the object
+            if hasattr(obj, '__dict__'):
+                obj.__dict__.clear()
+        return pooled_object
+    
+    def validate(self, pooled_object: PooledObject) -> bool:
+        """Validate if the pooled object is still usable."""
+        return pooled_object.keeped_object is not None
+
+# Global pond instances
+_pools: dict[str, Pond] = {}
+
+def get_object_pool(
+    pool_name: str,
+    object_class: Type[T],
+    pooled_maxsize: Optional[int] = None,  # None = unlimited pool with eviction-based memory control
+    least_one: bool = True,  # Always keep at least one
+    borrowed_timeout: int = 10,  # Longer timeout for high concurrency
+    time_between_eviction_runs: int = 300,  # Less frequent eviction to maintain high reuse ratio
+    eviction_weight: float = 0.3,  # Less aggressive eviction for better reuse
+    prewarm_count: int = 5  # Lower pre-warm count to reduce initial memory usage
+) -> Pond:
+    """Get or create a global object pool instance with balanced eviction-based memory control.
+    
+    Memory is controlled through moderate eviction to balance reuse ratio and memory usage:
+    - Moderate eviction frequency (300s) to maintain high object reuse ratio
+    - Conservative eviction weight (0.3) to avoid destroying useful objects
+    - Lower pre-warm count (5) to reduce initial memory footprint
+    
+    Args:
+        pool_name: Unique name for the pool
+        object_class: The class type to pool
+        pooled_maxsize: Maximum number of objects in the pool (None = truly unlimited)
+        least_one: Whether to keep at least one object in the pool (default: True)
+        borrowed_timeout: Timeout for borrowing objects (seconds, default: 10)
+        time_between_eviction_runs: Time between eviction runs (seconds, default: 300)
+        eviction_weight: Weight for eviction algorithm (default: 0.3, conservative)
+        prewarm_count: Number of objects to pre-warm the pool with (default: 5)
+    
+    Returns:
+        Pond instance for the specified object type
+    """
+    
+    if pool_name in _pools:
+        return _pools[pool_name]
+    
+    # Create new pond
+    pond = Pond(
+        borrowed_timeout=borrowed_timeout,
+        time_between_eviction_runs=time_between_eviction_runs,
+        thread_daemon=True,
+        eviction_weight=eviction_weight
+    )
+    
+    # Register the factory with user's maxsize preference
+    factory = GenericPooledObjectFactory(
+        object_class=object_class,
+        pooled_maxsize=pooled_maxsize,
+        least_one=least_one
+    )
+    pond.register(factory, name=f"{pool_name}Factory")
+    
+    # Pre-warm the pool
+    _prewarm_pool(pond, pool_name, prewarm_count)
+    
+    _pools[pool_name] = pond
+    return pond
+
+def _prewarm_pool(pond: Pond, pool_name: str, prewarm_count: int = 20) -> None:
+    """Pre-warm the pool with initial objects for high concurrency."""
+    for _ in range(prewarm_count):
+        try:
+            pooled_obj = pond.borrow(name=f"{pool_name}Factory")
+            pond.recycle(pooled_obj, name=f"{pool_name}Factory")
+        except Exception:
+            # If pre-warming fails, just continue
+            break

tests/litellm_utils_tests/test_object_pooling.py

@@ -0,0 +1,106 @@
+"""
+Simplified tests for object pooling utilities in litellm.
+"""
+
+import pytest
+
+from litellm.litellm_core_utils.object_pooling import (
+    get_object_pool,
+    _pools
+)
+
+
+class SimpleObject:
+    """Simple test object with internal reset tracking."""
+    def __init__(self):
+        self.data = {}
+        self.reset_count = 0
+        self.creation_id = id(self)
+    
+    def reset(self):
+        """Reset method that tracks how many times it's called."""
+        self.data.clear()
+        self.reset_count += 1
+    
+    def set_data(self, key, value):
+        """Set data to verify reset works."""
+        self.data[key] = value
+
+
+class SimpleObjectNoReset:
+    """Test object without reset method."""
+    def __init__(self):
+        self.data = {}
+        self.creation_id = id(self)
+
+
+class TestObjectPooling:
+    """Simplified test suite for object pooling."""
+    
+    def setup_method(self):
+        """Clear pools before each test."""
+        _pools.clear()
+    
+    def test_reset_method_works(self):
+        """Test that reset method is called when recycling objects."""
+        pool_name = "reset_test"
+        pool = get_object_pool(pool_name, SimpleObject, pooled_maxsize=1, prewarm_count=0)
+        
+        # Get an object and modify it
+        obj = pool.borrow(name=f"{pool_name}Factory")
+        obj.keeped_object.set_data("test", "value")
+        initial_reset_count = obj.keeped_object.reset_count
+        
+        # Return to pool (should trigger reset)
+        pool.recycle(obj, name=f"{pool_name}Factory")
+        
+        # Get the same object back
+        obj2 = pool.borrow(name=f"{pool_name}Factory")
+        
+        # Verify reset was called
+        assert obj2.keeped_object.reset_count == initial_reset_count + 1
+        assert obj2.keeped_object.data == {}  # Data should be cleared
+        assert obj.keeped_object.creation_id == obj2.keeped_object.creation_id  # Same object
+    
+    def test_fallback_reset_works(self):
+        """Test fallback reset when no reset method exists."""
+        pool_name = "fallback_test"
+        pool = get_object_pool(pool_name, SimpleObjectNoReset, pooled_maxsize=1, prewarm_count=0)
+        
+        # Get an object and modify it
+        obj = pool.borrow(name=f"{pool_name}Factory")
+        obj.keeped_object.data["test"] = "value"
+        
+        # Return to pool (should trigger fallback reset)
+        pool.recycle(obj, name=f"{pool_name}Factory")
+        
+        # Get the same object back
+        obj2 = pool.borrow(name=f"{pool_name}Factory")
+        
+        # Verify fallback reset worked - all attributes should be cleared by __dict__.clear()
+        assert obj2.keeped_object.__dict__ == {}, "All attributes should be cleared by fallback reset"
+        assert obj is obj2, "Should be the same pooled object instance"
+    
+    def test_pool_reuses_objects(self):
+        """Test that pool actually reuses objects instead of creating new ones."""
+        pool_name = "reuse_test"
+        pool = get_object_pool(pool_name, SimpleObject, pooled_maxsize=1, prewarm_count=0)
+        
+        # Get first object
+        obj1 = pool.borrow(name=f"{pool_name}Factory")
+        creation_id1 = obj1.keeped_object.creation_id
+        
+        # Return it
+        pool.recycle(obj1, name=f"{pool_name}Factory")
+        
+        # Get second object
+        obj2 = pool.borrow(name=f"{pool_name}Factory")
+        creation_id2 = obj2.keeped_object.creation_id
+        
+        # Should be the same object (reused)
+        assert creation_id1 == creation_id2, "Pool should reuse objects"
+        assert obj1.keeped_object is obj2.keeped_object, "Should be same object instance"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])