[Improvements] Manage segment cache and memory #1670

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

@@ -116,6 +116,7 @@ class Settings(BaseSettings):  # type: ignore
    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

@@ -116,6 +116,7 @@ class Settings(BaseSettings):  # type: ignore
    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

 SEGMENT_TYPE_IMPLS = {
    SegmentType.SQLITE: "chromadb.segment.impl.metadata.sqlite.SqliteMetadataSegment",
    SegmentType.HNSW_LOCAL_MEMORY: "chromadb.segment.impl.vector.local_hnsw.LocalHnswSegment",
    SegmentType.HNSW_LOCAL_PERSISTED: "chromadb.segment.impl.vector.local_persistent_hnsw.PersistentLocalHnswSegment",
 }


+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

@@ -150,6 +197,9 @@ def get_segment(self, collection_id: UUID, type: Type[S]) -> S:
            raise ValueError(f"Invalid segment type: {type}")

        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

@@ -116,6 +116,7 @@ class Settings(BaseSettings):  # type: ignore
    is_persistent: bool = False
    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

@@ -60,7 +59,12 @@ def __init__(self, system: System):
        self._system = system
        self._opentelemetry_client = system.require(OpenTelemetryClient)
        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

@@ -97,7 +106,9 @@ def reset_state(self) -> None:
            instance.stop()
            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/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
98 checks passed

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.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Improvements] Manage segment cache and memory #1670

[Improvements] Manage segment cache and memory #1670

nicolasgere commented Jan 23, 2024 •

edited

github-actions bot commented Jan 23, 2024

HammadB Jan 23, 2024

HammadB left a comment

nicolasgere Jan 25, 2024

HammadB Jan 25, 2024

nicolasgere Jan 26, 2024

HammadB Jan 26, 2024

HammadB Jan 26, 2024

HammadB Jan 26, 2024

HammadB Jan 26, 2024

nicolasgere Jan 29, 2024

HammadB Jan 30, 2024

HammadB Jan 26, 2024

HammadB Jan 26, 2024

HammadB Jan 30, 2024

HammadB Jan 30, 2024

HammadB Jan 30, 2024

HammadB left a comment

MrZoidberg commented Jan 31, 2024

nicolasgere commented Jan 31, 2024

		@@ -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:

[Improvements] Manage segment cache and memory #1670

[Improvements] Manage segment cache and memory #1670

Conversation

nicolasgere commented Jan 23, 2024 • edited

github-actions bot commented Jan 23, 2024

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Choose a reason for hiding this comment

HammadB left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HammadB left a comment

Choose a reason for hiding this comment

MrZoidberg commented Jan 31, 2024

nicolasgere commented Jan 31, 2024

nicolasgere commented Jan 23, 2024 •

edited