codekingpro/portable-devtools

1# Add up to 200k records until the log-is-full message is seen.
2
3import grpc
4import math
5import random
6import time
7
8import numpy as np
9
10from chromadb.api import ClientAPI
11from chromadb.proto.logservice_pb2 import SealLogRequest, MigrateLogRequest
12from chromadb.proto.logservice_pb2_grpc import LogServiceStub
13from chromadb.test.conftest import (
14    reset,
15    skip_if_not_cluster,
16)
17from chromadb.test.property import invariants
18from chromadb.test.utils.wait_for_version_increase import wait_for_version_increase
19
20RECORDS = 2000000
21BATCH_SIZE = 100
22
23@skip_if_not_cluster()
24def test_log_backpressure(
25    client: ClientAPI,
26) -> None:
27    seed = time.time()
28    random.seed(seed)
29    print("Generating data with seed ", seed)
30    reset(client)
31    collection = client.create_collection(
32        name="test",
33        metadata={"hnsw:construction_ef": 128, "hnsw:search_ef": 128, "hnsw:M": 128},
34    )
35
36    time.sleep(1)
37
38    print('backpressuring for', collection.id)
39
40    excepted = False
41    # Add RECORDS records, where each embedding has 3 dimensions randomly generated between 0 and 1
42    for i in range(0, RECORDS, BATCH_SIZE):
43        ids = []
44        embeddings = []
45        ids.extend([str(x) for x in range(i, i + BATCH_SIZE)])
46        embeddings.extend([np.random.rand(1, 3)[0] for x in range(i, i + BATCH_SIZE)])
47        try:
48            collection.add(ids=ids, embeddings=embeddings)
49        except Exception as x:
50            print(f"Caught exception:\n{x}")
51            if 'log needs compaction before accepting more writes; please backoff exponentially and retry' in str(x):
52                excepted = True
53                break
54    assert excepted, "Expected an exception to be thrown."
55