MemorySaver that evicts oldest threads when exceeding max_threads.
| 28 | |
| 29 | |
| 30 | class BoundedMemorySaver(MemorySaver): |
| 31 | """MemorySaver that evicts oldest threads when exceeding max_threads.""" |
| 32 | |
| 33 | def __init__(self, max_threads: int = 200): |
| 34 | super().__init__() |
| 35 | self.max_threads = max_threads |
| 36 | self._insertion_order: OrderedDict[str, None] = OrderedDict() |
| 37 | |
| 38 | def put(self, config, checkpoint, metadata, new_versions): |
| 39 | thread_id = config["configurable"]["thread_id"] |
| 40 | # Move to end if already tracked, otherwise insert |
| 41 | self._insertion_order[thread_id] = None |
| 42 | self._insertion_order.move_to_end(thread_id) |
| 43 | |
| 44 | result = super().put(config, checkpoint, metadata, new_versions) |
| 45 | |
| 46 | while len(self.storage) > self.max_threads and self._insertion_order: |
| 47 | oldest_thread, _ = self._insertion_order.popitem(last=False) |
| 48 | if oldest_thread in self.storage: |
| 49 | logger.info( |
| 50 | "BoundedMemorySaver: evicting thread %s (%d threads stored)", |
| 51 | oldest_thread, |
| 52 | len(self.storage), |
| 53 | ) |
| 54 | del self.storage[oldest_thread] |
| 55 | return result |