From 97be123030ab23dd700fc863ea891598ca797712 Mon Sep 17 00:00:00 2001 From: Trayan Azarov Date: Tue, 9 Apr 2024 02:44:12 +0300 Subject: [PATCH] [BUG]: Trace decorator fix - local segment manager (#1978) ## 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 --- chromadb/segment/impl/manager/local.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/chromadb/segment/impl/manager/local.py b/chromadb/segment/impl/manager/local.py index 19710b921d8..ec8ab17383c 100644 --- a/chromadb/segment/impl/manager/local.py +++ b/chromadb/segment/impl/manager/local.py @@ -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}") @@ -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 @@ -182,6 +182,10 @@ 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()]) @@ -189,6 +193,10 @@ def _get_segment_sysdb(self, collection_id: UUID, scope: SegmentScope): 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: