[ENH] Max batch size warning (#995)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- In v0.4.6 we introduced batching at the sqlite layer for the
embeddings_queue to maximize throughput through submit_embeddings calls.
SQlite has a max size of variables it can bind at once that is defined
at compile time. This changes makes the embeddings_queue introspect its
compile time flags to inform its max_batch_size. Max_batch_size is a new
property method added to the Producer class that indicates the maximum
batch size submit_embeddings() will accept. Right now we use the error
internal to submit_embeddings in the sqlite impl, but callers could use
this to pre-validate input and throw a message as well.
 - New functionality
	 - None

## Test plan
Added a test for the below and above batch size cases.

## Documentation Changes
We should update the usage guide to cover this caveat.

chromadb/db/mixins/embeddings_queue.py

@@ -68,9 +68,13 @@ def __init__(
             self.callback = callback
 
     _subscriptions: Dict[str, Set[Subscription]]
+    _max_batch_size: Optional[int]
+    # How many variables are in the insert statement for a single record
+    VARIABLES_PER_RECORD = 6
 
     def __init__(self, system: System):
         self._subscriptions = defaultdict(set)
+        self._max_batch_size = None
         super().__init__(system)
 
     @override
@@ -115,6 +119,15 @@ def submit_embeddings(
         if len(embeddings) == 0:
             return []
 
+        if len(embeddings) > self.max_batch_size:
+            raise ValueError(
+                f"""
+                Cannot submit more than {self.max_batch_size:,} embeddings at once.
+                Please submit your embeddings in batches of size
+                {self.max_batch_size:,} or less.
+                """
+            )
+
         t = Table("embeddings_queue")
         insert = (
             self.querybuilder()
@@ -208,6 +221,28 @@ def min_seqid(self) -> SeqId:
     def max_seqid(self) -> SeqId:
         return 2**63 - 1
 
+    @property
+    @override
+    def max_batch_size(self) -> int:
+        if self._max_batch_size is None:
+            with self.tx() as cur:
+                cur.execute("PRAGMA compile_options;")
+                compile_options = cur.fetchall()
+
+                for option in compile_options:
+                    if "MAX_VARIABLE_NUMBER" in option[0]:
+                        # The pragma returns a string like 'MAX_VARIABLE_NUMBER=999'
+                        self._max_batch_size = int(option[0].split("=")[1]) // (
+                            self.VARIABLES_PER_RECORD
+                        )
+
+                if self._max_batch_size is None:
+                    # This value is the default for sqlite3 versions < 3.32.0
+                    # It is the safest value to use if we can't find the pragma for some
+                    # reason
+                    self._max_batch_size = 999 // self.VARIABLES_PER_RECORD
+        return self._max_batch_size
+
     def _prepare_vector_encoding_metadata(
         self, embedding: SubmitEmbeddingRecord
     ) -> Tuple[Optional[bytes], Optional[str], Optional[str]]:

chromadb/ingest/__init__.py

@@ -59,7 +59,15 @@ def submit_embeddings(
         """Add a batch of embedding records to the given topic. Returns the SeqIDs of
         the records. The returned SeqIDs will be in the same order as the given
         SubmitEmbeddingRecords. However, it is not guaranteed that the SeqIDs will be
-        processed in the same order as the given SubmitEmbeddingRecords."""
+        processed in the same order as the given SubmitEmbeddingRecords. If the number
+        of records exceeds the maximum batch size, an exception will be thrown."""
+        pass
+
+    @property
+    @abstractmethod
+    def max_batch_size(self) -> int:
+        """Return the maximum number of records that can be submitted in a single call
+        to submit_embeddings."""
         pass
 
 

chromadb/test/ingest/test_producer_consumer.py

@@ -98,15 +98,15 @@ def __call__(self, embeddings: Sequence[EmbeddingRecord]) -> None:
             if len(self.embeddings) >= n:
                 event.set()
 
-    async def get(self, n: int) -> Sequence[EmbeddingRecord]:
+    async def get(self, n: int, timeout_secs: int = 10) -> Sequence[EmbeddingRecord]:
         "Wait until at least N embeddings are available, then return all embeddings"
         if len(self.embeddings) >= n:
             return self.embeddings[:n]
         else:
             event = Event()
             self.waiters.append((n, event))
             # timeout so we don't hang forever on failure
-            await wait_for(event.wait(), 10)
+            await wait_for(event.wait(), timeout_secs)
             return self.embeddings[:n]
 
 
@@ -322,3 +322,28 @@ async def test_multiple_topics_batch(
             recieved = await consume_fns[i].get(total_produced + PRODUCE_BATCH_SIZE)
             assert_records_match(embeddings_n[i], recieved)
         total_produced += PRODUCE_BATCH_SIZE
+
+
+@pytest.mark.asyncio
+async def test_max_batch_size(
+    producer_consumer: Tuple[Producer, Consumer],
+    sample_embeddings: Iterator[SubmitEmbeddingRecord],
+) -> None:
+    producer, consumer = producer_consumer
+    producer.reset_state()
+    max_batch_size = producer_consumer[0].max_batch_size
+    assert max_batch_size > 0
+
+    # Make sure that we can produce a batch of size max_batch_size
+    embeddings = [next(sample_embeddings) for _ in range(max_batch_size)]
+    consume_fn = CapturingConsumeFn()
+    consumer.subscribe("test_topic", consume_fn, start=consumer.min_seqid())
+    producer.submit_embeddings("test_topic", embeddings=embeddings)
+    received = await consume_fn.get(max_batch_size, timeout_secs=120)
+    assert_records_match(embeddings, received)
+
+    embeddings = [next(sample_embeddings) for _ in range(max_batch_size + 1)]
+    # Make sure that we can't produce a batch of size > max_batch_size
+    with pytest.raises(ValueError) as e:
+        producer.submit_embeddings("test_topic", embeddings=embeddings)
+    assert "Cannot submit more than" in str(e.value)