[CHORE] Add instrumentation to a slow Python call.

chroma-core · Dec 18, 2024 · a15c780 · a15c780
1 parent 9ed5b07
commit a15c780
Showing 1 changed file with 14 additions and 11 deletions.
diff --git a/chromadb/server/fastapi/__init__.py b/chromadb/server/fastapi/__init__.py
@@ -710,17 +710,20 @@ async def get_collection(
 
         add_attributes_to_current_span({"tenant": tenant})
 
-        api_collection_model = cast(
-            CollectionModel,
-            await to_thread.run_sync(
-                self._api.get_collection,
-                collection_name,
-                tenant,
-                database_name,
-                limiter=self._capacity_limiter,
-            ),
-        )
-        return api_collection_model
+        @trace_method("Outside the pool", OpenTelemetryGranularity.OPERATION)
+        def inner():
+            api_collection_model = cast(
+                CollectionModel,
+                await to_thread.run_sync(
+                    self._api.get_collection,
+                    collection_name,
+                    tenant,
+                    database_name,
+                    limiter=self._capacity_limiter,
+                ),
+            )
+            return api_collection_model
+        return inner()
 
     @trace_method("FastAPI.update_collection", OpenTelemetryGranularity.OPERATION)
     async def update_collection(