codekingpro/portable-devtools

1import time
2from chromadb.api import ClientAPI
3from chromadb.test.conftest import COMPACTION_SLEEP
4
5TIMEOUT_INTERVAL = 1
6
7
8def get_collection_version(client: ClientAPI, collection_name: str) -> int:
9    coll = client.get_collection(collection_name)
10    return coll.get_model()["version"]
11
12
13def wait_for_version_increase(
14    client: ClientAPI,
15    collection_name: str,
16    initial_version: int,
17    additional_time: int = 0,
18) -> int:
19    timeout = COMPACTION_SLEEP
20    deadline = time.time() + timeout + additional_time
21    target_version = initial_version + 1
22
23    curr_version = get_collection_version(client, collection_name)
24    if curr_version == initial_version:
25        print(
26            "[wait_for_version_increase] "
27            f"collection={collection_name} "
28            f"waiting for version >= {target_version} "
29            f"(current={curr_version}, timeout={timeout + additional_time}s)"
30        )
31    while curr_version == initial_version:
32        time.sleep(TIMEOUT_INTERVAL)
33        if time.time() > deadline:
34            collection_id = client.get_collection(collection_name).id
35            raise TimeoutError(
36                "Model was not updated in time for "
37                f"{collection_id}; waited for version >= {target_version}, "
38                f"last seen version {curr_version}"
39            )
40        curr_version = get_collection_version(client, collection_name)
41
42    return curr_version
43