Coverage for agentos/tools/memory_optimizer.py: 32%
212 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 11:37 +0800
1"""
2Memory Optimization Tools for AgentOS.
3Object pooling, LRU caching, memory monitoring, and smart caching with TTL.
4"""
6import threading
7import time
8from collections import OrderedDict
9from dataclasses import dataclass
10from typing import Any, Callable, Dict, Generic, List, Optional, TypeVar
12T = TypeVar("T")
15# ============================================================================
16# ObjectPool
17# ============================================================================
19class ObjectPool(Generic[T]):
20 """Thread-safe object pool with auto-expiry and size limits.
22 Reuses pre-allocated objects instead of creating/destroying them repeatedly.
23 """
25 def __init__(
26 self,
27 factory: Callable[[], T],
28 max_size: int = 100,
29 max_idle: int = 30,
30 idle_timeout: float = 300.0,
31 ):
32 self._factory = factory
33 self._max_size = max_size
34 self._max_idle = max_idle
35 self._idle_timeout = idle_timeout
36 self._pool: List[_PooledItem[T]] = []
37 self._lock = threading.Lock()
38 self._created: int = 0
39 self._borrowed: int = 0
40 self._returned: int = 0
42 def acquire(self) -> T:
43 """Borrow an object from the pool or create a new one."""
44 with self._lock:
45 now = time.monotonic()
46 self._evict_expired(now)
48 if self._pool:
49 item = self._pool.pop()
50 item.idle = False
51 self._borrowed += 1
52 return item.obj
54 if self._created < self._max_size:
55 self._created += 1
56 self._borrowed += 1
57 return self._factory()
59 # Pool fully allocated; create temporary object outside pool
60 self._borrowed += 1
61 return self._factory()
63 def release(self, obj: T) -> None:
64 """Return an object to the pool for reuse."""
65 with self._lock:
66 self._returned += 1
67 self._evict_expired(time.monotonic())
69 if len(self._pool) < self._max_idle:
70 self._pool.append(_PooledItem(obj=obj, idle=True, acquired_at=time.monotonic()))
71 # else: discard excess objects
73 def _evict_expired(self, now: float) -> None:
74 self._pool[:] = [item for item in self._pool if now - item.acquired_at < self._idle_timeout]
76 @property
77 def stats(self) -> Dict[str, Any]:
78 with self._lock:
79 return {
80 "created": self._created,
81 "borrowed": self._borrowed,
82 "returned": self._returned,
83 "idle": len(self._pool),
84 "active": self._borrowed - self._returned,
85 }
87 def __len__(self) -> int:
88 return len(self._pool)
91@dataclass
92class _PooledItem(Generic[T]):
93 obj: T
94 idle: bool
95 acquired_at: float
98# ============================================================================
99# LRUCache
100# ============================================================================
102class LRUCache(Generic[T]):
103 """Thread-safe LRU cache with capacity limit and optional TTL."""
105 def __init__(self, capacity: int = 1024, ttl: Optional[float] = None):
106 self._capacity = capacity
107 self._ttl = ttl
108 self._cache: OrderedDict[str, _CacheEntry[T]] = OrderedDict()
109 self._lock = threading.Lock()
110 self._hits: int = 0
111 self._misses: int = 0
112 self._evictions: int = 0
114 def get(self, key: str) -> Optional[T]:
115 with self._lock:
116 entry = self._cache.get(key)
117 if entry is None:
118 self._misses += 1
119 return None
120 if self._ttl and time.monotonic() - entry.timestamp > self._ttl:
121 del self._cache[key]
122 self._misses += 1
123 self._evictions += 1
124 return None
125 self._cache.move_to_end(key)
126 self._hits += 1
127 return entry.value
129 def put(self, key: str, value: T) -> None:
130 with self._lock:
131 if key in self._cache:
132 self._cache.move_to_end(key)
133 self._cache[key] = _CacheEntry(value=value, timestamp=time.monotonic())
134 return
135 if len(self._cache) >= self._capacity:
136 self._cache.popitem(last=False)
137 self._evictions += 1
138 self._cache[key] = _CacheEntry(value=value, timestamp=time.monotonic())
140 def remove(self, key: str) -> bool:
141 with self._lock:
142 if key in self._cache:
143 del self._cache[key]
144 return True
145 return False
147 def clear(self) -> None:
148 with self._lock:
149 self._cache.clear()
151 @property
152 def hit_rate(self) -> float:
153 total = self._hits + self._misses
154 return self._hits / total if total > 0 else 0.0
156 @property
157 def stats(self) -> Dict[str, Any]:
158 with self._lock:
159 return {
160 "size": len(self._cache),
161 "capacity": self._capacity,
162 "hits": self._hits,
163 "misses": self._misses,
164 "hit_rate": round(self.hit_rate, 4),
165 "evictions": self._evictions,
166 }
168 def __len__(self) -> int:
169 return len(self._cache)
171 def __contains__(self, key: str) -> bool:
172 return key in self._cache
175@dataclass
176class _CacheEntry(Generic[T]):
177 value: T
178 timestamp: float
181# ============================================================================
182# SmartCache
183# ============================================================================
185class SmartCache(Generic[T]):
186 """Multi-tier cache with compute-on-miss and automatic invalidation."""
188 def __init__(self, compute: Callable[[str], T], capacity: int = 1024, ttl: float = 300.0):
189 self._lru = LRUCache[T](capacity=capacity, ttl=ttl)
190 self._compute = compute
191 self._lock = threading.Lock()
193 def get(self, key: str) -> T:
194 """Get from cache or compute and cache on miss."""
195 cached = self._lru.get(key)
196 if cached is not None:
197 return cached
198 with self._lock:
199 # Double-check after acquiring lock
200 cached = self._lru.get(key)
201 if cached is not None:
202 return cached
203 value = self._compute(key)
204 self._lru.put(key, value)
205 return value
207 def prefetch(self, keys: List[str]) -> int:
208 """Pre-compute and cache values for a list of keys. Returns count cached."""
209 count = 0
210 for key in keys:
211 if key not in self._lru:
212 try:
213 self.get(key)
214 count += 1
215 except Exception:
216 pass
217 return count
219 def invalidate(self, key: str) -> bool:
220 return self._lru.remove(key)
222 def invalidate_pattern(self, pattern: str) -> int:
223 """Invalidate all keys containing pattern substring. Returns count removed."""
224 count = 0
225 for key in list(self._lru._cache.keys()):
226 if pattern in key:
227 if self._lru.remove(key):
228 count += 1
229 return count
231 def clear(self) -> None:
232 self._lru.clear()
234 @property
235 def stats(self) -> Dict[str, Any]:
236 return self._lru.stats
238 def __len__(self) -> int:
239 return len(self._lru)
242# ============================================================================
243# MemoryMonitor
244# ============================================================================
246class MemoryMonitor:
247 """Monitor per-component memory usage with high-water mark tracking."""
249 _singleton = None
250 _lock = threading.Lock()
252 def __new__(cls):
253 if cls._singleton is None:
254 with cls._lock:
255 if cls._singleton is None:
256 cls._singleton = super().__new__(cls)
257 cls._singleton._initialized = False
258 return cls._singleton
260 def __init__(self):
261 if self._initialized:
262 return
263 self._initialized = True
264 self._components: Dict[str, _ComponentMetrics] = {}
265 self._lock = threading.Lock()
267 def register(self, name: str) -> None:
268 with self._lock:
269 if name not in self._components:
270 self._components[name] = _ComponentMetrics(name=name)
272 def record_alloc(self, name: str, size_bytes: int) -> None:
273 with self._lock:
274 comp = self._components.get(name)
275 if comp:
276 comp.current_bytes += size_bytes
277 comp.total_allocations += 1
278 comp.peak_bytes = max(comp.peak_bytes, comp.current_bytes)
280 def record_free(self, name: str, size_bytes: int) -> None:
281 with self._lock:
282 comp = self._components.get(name)
283 if comp:
284 comp.current_bytes = max(0, comp.current_bytes - size_bytes)
286 def snapshot(self) -> Dict[str, Dict[str, Any]]:
287 with self._lock:
288 return {name: comp.to_dict() for name, comp in self._components.items()}
290 def alert(self, name: str, threshold_bytes: int) -> bool:
291 """Check if a component exceeds memory threshold."""
292 with self._lock:
293 comp = self._components.get(name)
294 if comp:
295 return comp.current_bytes > threshold_bytes
296 return False
298 @property
299 def total_current(self) -> int:
300 with self._lock:
301 return sum(c.current_bytes for c in self._components.values())
304@dataclass
305class _ComponentMetrics:
306 name: str
307 current_bytes: int = 0
308 peak_bytes: int = 0
309 total_allocations: int = 0
311 def to_dict(self) -> Dict[str, Any]:
312 return {
313 "name": self.name,
314 "current_bytes": self.current_bytes,
315 "peak_bytes": self.peak_bytes,
316 "total_allocations": self.total_allocations,
317 }
320# ============================================================================
321# Convenience Functions
322# ============================================================================
324def create_object_pool(
325 factory: Callable[[], T],
326 max_size: int = 100,
327 max_idle: int = 30,
328 idle_timeout: float = 300.0,
329) -> ObjectPool[T]:
330 """Create a thread-safe object pool."""
331 return ObjectPool(factory, max_size=max_size, max_idle=max_idle, idle_timeout=idle_timeout)
334def create_lru_cache(capacity: int = 1024, ttl: Optional[float] = None) -> LRUCache[Any]:
335 """Create a thread-safe LRU cache."""
336 return LRUCache(capacity=capacity, ttl=ttl)
339def create_smart_cache(
340 compute: Callable[[str], T],
341 capacity: int = 1024,
342 ttl: float = 300.0,
343) -> SmartCache[T]:
344 """Create a smart cache with compute-on-miss."""
345 return SmartCache(compute, capacity=capacity, ttl=ttl)
348def get_memory_monitor() -> MemoryMonitor:
349 """Get the singleton memory monitor."""
350 return MemoryMonitor()