[Improvements] Manage segment cache and memory#1670
Conversation
Reviewer ChecklistPlease leverage this checklist to ensure your code review is thorough before approving Testing, Bugs, Errors, Logs, Documentation
System Compatibility
Quality
|
| is_persistent: bool = False | ||
| persist_directory: str = "./chroma" | ||
|
|
||
| chroma_memory_limit: int = 0 |
There was a problem hiding this comment.
nit add units to this chroma_memory_limit_gb etc
HammadB
left a comment
HammadB
left a comment
There was a problem hiding this comment.
Can we add a test for this? We could do a property-based state machine test or basic unit tests with a couple cases. The former is preferable as we can also test edge cases like partial sizes/best-effort exhaustively etc
| @@ -116,6 +116,7 @@ class Settings(BaseSettings): # type: ignore | |||
| is_persistent: bool = False | |||
| persist_directory: str = "./chroma" | |||
|
|
|||
There was a problem hiding this comment.
I would update the documentation.
nicolasgere
changed the title
| is_persistent: bool = False | ||
| persist_directory: str = "./chroma" | ||
|
|
||
| chroma_memory_limit_bytes: int = 0 |
There was a problem hiding this comment.
How do I turn this capability on and off? Is 0 implicitly off?
There was a problem hiding this comment.
yes. 0 is unlimited
| } | ||
|
|
||
|
|
||
| def get_directory_size(directory: str) -> int: |
There was a problem hiding this comment.
can we move this into utils
| total_size = sum(segment_sizes.values()) | ||
| new_segment_size = self._get_segment_disk_size(collection_id) | ||
|
|
||
| while total_size + new_segment_size >= target_size and self._segment_cache.keys(): |
There was a problem hiding this comment.
Can we move this LRU logic into a cleaner abstraction with eviction handles - like https://github.com/chroma-core/chroma/blob/main/chromadb/utils/lru_cache.py
| if scope not in self._segment_cache[collection_id]: | ||
| memory_limit = self._system.settings.require("chroma_memory_limit_bytes") | ||
| if type == VectorReader and self._system.settings.require("is_persistent") and memory_limit > 0: | ||
| self._cleanup_segment(collection_id, memory_limit) |
There was a problem hiding this comment.
nit: cleanup_segment could be renamed a bit more explicitly / when we have a proper lru cache abstraction for this
|
|
||
| @invariant() | ||
| @precondition(lambda self: self.system.settings.is_persistent is True) | ||
| def cache_should_not_be_bigger_than_settings(self): |
There was a problem hiding this comment.
whats the behavior for boundary conditions? Eg. if the limit is 10GB and we have two files - 6GB and 7GB, will we always only allow one?
There was a problem hiding this comment.
Yes, it this what we want? We could add a message in log when a collection got evicted for memory constraint
|
|
||
|
|
||
| @patch('chromadb.segment.impl.manager.local.get_directory_size', SegmentManagerStateMachine.mock_directory_size) | ||
| def test_segment_manager(caplog: pytest.LogCaptureFixture, system: System) -> None: |
| persist_directory: str = "./chroma" | ||
|
|
||
| chroma_memory_limit_bytes: int = 0 | ||
| chroma_server_host: Optional[str] = None |
There was a problem hiding this comment.
Can we introduce a config - called segment_manager_cache_policy and make this one of many types?
| self._instances = {} | ||
| self._segment_cache = defaultdict(dict) | ||
| self.segment_cache: Dict[SegmentScope, SegmentCache] = {SegmentScope.METADATA: BasicCache(), SegmentScope.VECTOR: BasicCache()} | ||
| if system.settings.chroma_segment_cache_policy is "LRU" and system.settings.chroma_memory_limit_bytes > 0: |
There was a problem hiding this comment.
should use require() and catch the exception since its optional
| instance.reset_state() | ||
| self._instances = {} | ||
| self._segment_cache = defaultdict(dict) | ||
| self.segment_cache[SegmentScope.METADATA].reset() |
There was a problem hiding this comment.
we don't need to to cache metadata segments with LRU, since its all one segment in sqlite
| if len(segments) == 0: | ||
| return 0 | ||
| size = get_directory_size( | ||
| os.path.join(self._system.settings.require("persist_directory"), str(segments[0]["id"]))) |
There was a problem hiding this comment.
nit: leave comment stating assumption of one vector segment for this line - otherwise hardcoded [0] is confusing
HammadB
left a comment
HammadB
left a comment
There was a problem hiding this comment.
Some minor changes, but otherwise looks great!
|
I think the memory management should use docker limits if present as a default limit. |
|
It can be possible, even if for now, our limit is not an hard limit. We are adding documentation to explain more about how to use it. chroma-core/docs#211 |
Add configuration to limit the memory use by segment cache.
3 participants
The goal is too limit the memory usage by setting a limit that Chroma will do the best effort to not use more memory than the threshold by removing from cache segment.
How?
When a segment need to be load from the disk, we will check what is the size of all the already loaded index in memory, and will free some memory using LRU kind algo.