[Improvements] Manage segment cache and memory by nicolasgere · Pull Request #1670 · chroma-core/chroma

nicolasgere · 2024-01-23T01:45:29Z

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.

github-actions · 2024-01-23T01:45:44Z

HammadB · 2024-01-23T01:48:08Z

chromadb/config.py

    is_persistent: bool = False
    persist_directory: str = "./chroma"

+    chroma_memory_limit: int = 0


nit add units to this chroma_memory_limit_gb etc

HammadB

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

nicolasgere · 2024-01-25T21:00:32Z

chromadb/config.py

@@ -116,6 +116,7 @@ class Settings(BaseSettings):  # type: ignore
    is_persistent: bool = False
    persist_directory: str = "./chroma"



I would update the documentation.

HammadB · 2024-01-25T23:12:11Z

chromadb/config.py

    is_persistent: bool = False
    persist_directory: str = "./chroma"

+    chroma_memory_limit_bytes: int = 0


How do I turn this capability on and off? Is 0 implicitly off?

yes. 0 is unlimited

HammadB · 2024-01-26T18:45:36Z

chromadb/segment/impl/manager/local.py

 }


+def get_directory_size(directory: str) -> int:


can we move this into utils

HammadB · 2024-01-26T18:52:34Z

chromadb/segment/impl/manager/local.py

+        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():


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

HammadB · 2024-01-26T18:53:22Z

chromadb/segment/impl/manager/local.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)


nit: cleanup_segment could be renamed a bit more explicitly / when we have a proper lru cache abstraction for this

chromadb/test/property/test_segment_manager.py

HammadB · 2024-01-26T18:55:48Z

chromadb/test/property/test_segment_manager.py

+
+    @invariant()
+    @precondition(lambda self: self.system.settings.is_persistent is True)
+    def cache_should_not_be_bigger_than_settings(self):


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?

Yes, it this what we want? We could add a message in log when a collection got evicted for memory constraint

Yeah lets log

HammadB · 2024-01-26T18:57:05Z

chromadb/test/property/test_segment_manager.py

+
+
+@patch('chromadb.segment.impl.manager.local.get_directory_size', SegmentManagerStateMachine.mock_directory_size)
+def test_segment_manager(caplog: pytest.LogCaptureFixture, system: System) -> None:


HammadB · 2024-01-26T18:58:35Z

chromadb/config.py

    persist_directory: str = "./chroma"

+    chroma_memory_limit_bytes: int = 0
    chroma_server_host: Optional[str] = None


Can we introduce a config - called segment_manager_cache_policy and make this one of many types?

HammadB · 2024-01-30T00:32:59Z

chromadb/segment/impl/manager/local.py

        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:


should use require() and catch the exception since its optional

HammadB · 2024-01-30T00:33:40Z

chromadb/segment/impl/manager/local.py

            instance.reset_state()
        self._instances = {}
-        self._segment_cache = defaultdict(dict)
+        self.segment_cache[SegmentScope.METADATA].reset()


we don't need to to cache metadata segments with LRU, since its all one segment in sqlite

HammadB · 2024-01-30T00:34:37Z

chromadb/segment/impl/manager/local.py

+        if len(segments) == 0:
+            return 0
+        size = get_directory_size(
+            os.path.join(self._system.settings.require("persist_directory"), str(segments[0]["id"])))


nit: leave comment stating assumption of one vector segment for this line - otherwise hardcoded [0] is confusing

HammadB

Some minor changes, but otherwise looks great!

MrZoidberg · 2024-01-31T11:33:27Z

I think the memory management should use docker limits if present as a default limit.

nicolasgere · 2024-01-31T19:18:53Z

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.

[inprovements] Manage segment cache and memory

dd2331d

nicolasgere requested a review from HammadB January 23, 2024 01:45

HammadB reviewed Jan 23, 2024

View reviewed changes

HammadB requested changes Jan 23, 2024

View reviewed changes

nicolas added 4 commits January 24, 2024 19:17

add testing

877d712

cleanup test

cb15ae4

fix test

48b9cfe

fix test

db78fa1

nicolasgere commented Jan 25, 2024

View reviewed changes

nicolas added 2 commits January 25, 2024 13:40

improve test

06e7a84

fix type

1f30cbc

nicolasgere changed the title ~~WIP: [inprovements] Manage segment cache and memory~~ [Improvements] Manage segment cache and memory Jan 25, 2024

nicolasgere requested a review from HammadB January 25, 2024 23:02

HammadB reviewed Jan 25, 2024

View reviewed changes

HammadB reviewed Jan 26, 2024

View reviewed changes

chromadb/segment/impl/manager/local.py Outdated

}

def get_directory_size(directory: str) -> int:

Copy link

Collaborator

HammadB Jan 26, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we move this into utils

HammadB reviewed Jan 26, 2024

View reviewed changes

chromadb/test/property/test_segment_manager.py Show resolved Hide resolved

HammadB reviewed Jan 26, 2024

View reviewed changes

nicolas added 5 commits January 29, 2024 11:25

updata based on review

9e909ef

add comment

8a6f537

fix cache

097cc51

fix fmt

42fbc6d

fmt more code

2386cdd

nicolasgere requested a review from HammadB January 30, 2024 00:24

HammadB reviewed Jan 30, 2024

View reviewed changes

HammadB approved these changes Jan 30, 2024

View reviewed changes

nicolas added 4 commits January 30, 2024 11:57

updagte based on comment

ba45af9

updagte based on comment

9724918

fix test

034940c

add comment

15cc717

nicolasgere merged commit e5751fd into main Jan 31, 2024

tazarov pushed a commit to amikos-tech/chroma-core that referenced this pull request Feb 10, 2024

[Improvements] Manage segment cache and memory (chroma-core#1670)

c50ad16

Add configuration to limit the memory use by segment cache.

		@@ -116,6 +116,7 @@ class Settings(BaseSettings): # type: ignore
		is_persistent: bool = False
		persist_directory: str = "./chroma"



		@patch('chromadb.segment.impl.manager.local.get_directory_size', SegmentManagerStateMachine.mock_directory_size)
		def test_segment_manager(caplog: pytest.LogCaptureFixture, system: System) -> None:

Conversation

nicolasgere commented Jan 23, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jan 23, 2024

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HammadB left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HammadB left a comment

Choose a reason for hiding this comment

Uh oh!

MrZoidberg commented Jan 31, 2024

Uh oh!

nicolasgere commented Jan 31, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nicolasgere commented Jan 23, 2024 •

edited

Loading