Skip to content

Commit

Permalink
More Posthog batching (chroma-core#1342)
Browse files Browse the repository at this point in the history
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- Batch posthog requests much more aggressively. We really should do
batches of ever-increasing size for a given collection but I want to get
this fix out this week.
 - New functionality
	 - ...

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

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

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*
  • Loading branch information
beggers authored Nov 7, 2023
1 parent 69cd2ba commit ccfa914
Showing 1 changed file with 45 additions and 2 deletions.
47 changes: 45 additions & 2 deletions chromadb/telemetry/product/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def __init__(self, collection_uuid: str, embedding_function: str):


class CollectionAddEvent(ProductTelemetryEvent):
max_batch_size: ClassVar[int] = 100
max_batch_size: ClassVar[int] = 1000
batch_size: int
collection_uuid: str
add_amount: int
Expand Down Expand Up @@ -67,6 +67,8 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionAddEvent":


class CollectionUpdateEvent(ProductTelemetryEvent):
max_batch_size: ClassVar[int] = 100
batch_size: int
collection_uuid: str
update_amount: int
with_embeddings: int
Expand All @@ -80,17 +82,36 @@ def __init__(
with_embeddings: int,
with_metadata: int,
with_documents: int,
batch_size: int = 1,
):
super().__init__()
self.collection_uuid = collection_uuid
self.update_amount = update_amount
self.with_embeddings = with_embeddings
self.with_metadata = with_metadata
self.with_documents = with_documents
self.batch_size = batch_size

@property
def batch_key(self) -> str:
return self.collection_uuid + self.name

def batch(self, other: "ProductTelemetryEvent") -> "CollectionUpdateEvent":
if not self.batch_key == other.batch_key:
raise ValueError("Cannot batch events")
other = cast(CollectionUpdateEvent, other)
total_amount = self.update_amount + other.update_amount
return CollectionUpdateEvent(
collection_uuid=self.collection_uuid,
update_amount=total_amount,
with_documents=self.with_documents + other.with_documents,
with_metadata=self.with_metadata + other.with_metadata,
with_embeddings=self.with_embeddings + other.with_embeddings,
batch_size=self.batch_size + other.batch_size,
)

class CollectionQueryEvent(ProductTelemetryEvent):
max_batch_size: ClassVar[int] = 20
max_batch_size: ClassVar[int] = 1000
batch_size: int
collection_uuid: str
query_amount: int
Expand Down Expand Up @@ -147,6 +168,8 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionQueryEvent":


class CollectionGetEvent(ProductTelemetryEvent):
max_batch_size: ClassVar[int] = 100
batch_size: int
collection_uuid: str
ids_count: int
limit: int
Expand All @@ -160,13 +183,33 @@ def __init__(
limit: int,
include_metadata: int,
include_documents: int,
batch_size: int = 1,
):
super().__init__()
self.collection_uuid = collection_uuid
self.ids_count = ids_count
self.limit = limit
self.include_metadata = include_metadata
self.include_documents = include_documents
self.batch_size = batch_size

@property
def batch_key(self) -> str:
return self.collection_uuid + self.name + str(self.limit)

def batch(self, other: "ProductTelemetryEvent") -> "CollectionGetEvent":
if not self.batch_key == other.batch_key:
raise ValueError("Cannot batch events")
other = cast(CollectionGetEvent, other)
total_amount = self.ids_count + other.ids_count
return CollectionGetEvent(
collection_uuid=self.collection_uuid,
ids_count=total_amount,
limit=self.limit,
include_metadata=self.include_metadata + other.include_metadata,
include_documents=self.include_documents + other.include_documents,
batch_size=self.batch_size + other.batch_size,
)


class CollectionDeleteEvent(ProductTelemetryEvent):
Expand Down

0 comments on commit ccfa914

Please sign in to comment.