Skip to content

Commit

Permalink
[BUG]: Trace decorator fix - local segment manager (chroma-core#1978)
Browse files Browse the repository at this point in the history
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
	 - Fixed a wrong decorator in local segment manager
	 - Added a few more tracing decorators

## Test plan
*How are these changes tested?*

- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
N/A
  • Loading branch information
tazarov authored Apr 8, 2024
1 parent 9296d28 commit 97be123
Showing 1 changed file with 12 additions and 4 deletions.
16 changes: 12 additions & 4 deletions chromadb/segment/impl/manager/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,10 @@ def __init__(self, system: System):
segment_limit, callback=lambda _, v: v.close_persistent_index()
)

@trace_method(
"LocalSegmentManager.callback_cache_evict",
OpenTelemetryGranularity.OPERATION_AND_SEGMENT,
)
def callback_cache_evict(self, segment: Segment):
collection_id = segment["collection"]
self.logger.info(f"LRU cache evict collection {collection_id}")
Expand Down Expand Up @@ -163,10 +167,6 @@ def delete_segments(self, collection_id: UUID) -> Sequence[UUID]:
self.segment_cache[SegmentScope.METADATA].pop(collection_id)
return [s["id"] for s in segments]

@trace_method(
"LocalSegmentManager.get_segment",
OpenTelemetryGranularity.OPERATION_AND_SEGMENT,
)
def _get_segment_disk_size(self, collection_id: UUID) -> int:
segments = self._sysdb.get_segments(
collection=collection_id, scope=SegmentScope.VECTOR
Expand All @@ -182,13 +182,21 @@ def _get_segment_disk_size(self, collection_id: UUID) -> int:
)
return size

@trace_method(
"LocalSegmentManager._get_segment_sysdb",
OpenTelemetryGranularity.OPERATION_AND_SEGMENT,
)
def _get_segment_sysdb(self, collection_id: UUID, scope: SegmentScope):
segments = self._sysdb.get_segments(collection=collection_id, scope=scope)
known_types = set([k.value for k in SEGMENT_TYPE_IMPLS.keys()])
# Get the first segment of a known type
segment = next(filter(lambda s: s["type"] in known_types, segments))
return segment

@trace_method(
"LocalSegmentManager.get_segment",
OpenTelemetryGranularity.OPERATION_AND_SEGMENT,
)
@override
def get_segment(self, collection_id: UUID, type: Type[S]) -> S:
if type == MetadataReader:
Expand Down

0 comments on commit 97be123

Please sign in to comment.