diff --git a/integration/test_batch_v4.py b/integration/test_batch_v4.py index 3646f1b1f..2d8e48b32 100644 --- a/integration/test_batch_v4.py +++ b/integration/test_batch_v4.py @@ -524,8 +524,10 @@ def test_add_1000_objects_with_async_indexing_and_wait( assert ret.total_count == nr_objects shards = client.collections.use(name).config.get_shards() - assert shards[0].status == "READY" - assert shards[0].vector_queue_size == 0 + if shards[0].status: + assert shards[0].status == "READY" + elif shards[0].per_node_status: + assert all(status == "READY" for status in shards[0].per_node_status.values()) @pytest.mark.skip("Difficult to find numbers that work reliably in the CI") @@ -545,8 +547,10 @@ def test_add_10000_objects_with_async_indexing_and_dont_wait( vector=[float((j + i) % nr_objects) / nr_objects for j in range(vec_length)], ) shard_status = old_client.schema.get_class_shards(name) - assert shard_status[0]["status"] == "INDEXING" - assert shard_status[0]["vectorQueueSize"] > 0 + if shard_status[0].status: + assert shard_status[0].status == "INDEXING" + elif shard_status[0].per_node_status: + assert all(status == "INDEXING" for status in shard_status[0].per_node_status.values()) assert len(client.batch.failed_objects) == 0 @@ -579,8 +583,10 @@ def test_add_1000_tenant_objects_with_async_indexing_and_wait_for_all( shards = client.collections.use(name).config.get_shards() for shard in shards: - assert shard.status == "READY" - assert shard.vector_queue_size == 0 + if shard.status: + assert shard.status == "READY" + elif shard.per_node_status: + assert all(status == "READY" for status in shard.per_node_status.values()) @pytest.mark.skip("Difficult to find numbers that work reliably in the CI") @@ -611,12 +617,11 @@ def test_add_1000_tenant_objects_with_async_indexing_and_wait_for_only_one( shards = client.collections.use(name).config.get_shards() for shard in shards: - if shard.name == tenants[0].name: - assert shard.status == "READY" - assert shard.vector_queue_size == 0 - else: - assert shard.status == "INDEXING" - assert shard.vector_queue_size > 0 + want = "READY" if shard.name == tenants[0].name else "INDEXING" + if shard.status: + assert shard.status == want + elif shard.per_node_status: + assert all(status == want for status in shard.per_node_status.values()) @pytest.mark.parametrize( diff --git a/integration/test_collection_config.py b/integration/test_collection_config.py index 786b486cd..897366541 100644 --- a/integration/test_collection_config.py +++ b/integration/test_collection_config.py @@ -787,8 +787,10 @@ def test_collection_config_get_shards(collection_factory: CollectionFactory) -> ) shards = collection.config.get_shards() assert len(shards) - assert shards[0].status == "READY" - assert shards[0].vector_queue_size == 0 + if shards[0].status: + assert shards[0].status == "READY" + elif shards[0].per_node_status: + assert all(status == "READY" for status in shards[0].per_node_status.values()) def test_collection_update_shards(collection_factory: CollectionFactory) -> None: @@ -836,11 +838,11 @@ def test_collection_config_get_shards_multi_tenancy(collection_factory: Collecti shards = collection.config.get_shards() assert len(shards) == 2 - assert shards[0].status == "READY" - assert shards[0].vector_queue_size == 0 - - assert shards[1].status == "READY" - assert shards[1].vector_queue_size == 0 + for shard in shards: + if shard.status: + assert shard.status == "READY" + elif shard.per_node_status: + assert all(status == "READY" for status in shard.per_node_status.values()) assert "tenant1" in [shard.name for shard in shards] assert "tenant2" in [shard.name for shard in shards] diff --git a/weaviate/collections/batch/batch_wrapper.py b/weaviate/collections/batch/batch_wrapper.py index a3a3598d6..1c31b5ab7 100644 --- a/weaviate/collections/batch/batch_wrapper.py +++ b/weaviate/collections/batch/batch_wrapper.py @@ -93,9 +93,16 @@ def __get_shards_readiness(self, shard: Shard) -> List[bool]: res = _decode_json_response_list(response, "Get shards' status") assert res is not None + return [ - (cast(str, shard.get("status")) == "READY") - & (cast(int, shard.get("vectorQueueSize")) == 0) + ( + all( + status == "READY" + for status in cast(dict[str, str], shard["per_node_status"]).values() + ) + if "per_node_status" in shard + else cast(str, shard["status"]) == "READY" + ) for shard in res ] @@ -195,8 +202,14 @@ async def __get_shards_readiness(self, shard: Shard) -> List[bool]: res = _decode_json_response_list(response, "Get shards' status") assert res is not None return [ - (cast(str, shard.get("status")) == "READY") - & (cast(int, shard.get("vectorQueueSize")) == 0) + ( + all( + status == "READY" + for status in cast(dict[str, str], shard["per_node_status"]).values() + ) + if "per_node_status" in shard + else cast(str, shard["status"]) == "READY" + ) for shard in res ] diff --git a/weaviate/collections/classes/config.py b/weaviate/collections/classes/config.py index 19396a7fc..4d7f586f6 100644 --- a/weaviate/collections/classes/config.py +++ b/weaviate/collections/classes/config.py @@ -2251,8 +2251,9 @@ class _CollectionConfigSimple(_ConfigBase): @dataclass class _ShardStatus: name: str - status: ShardTypes - vector_queue_size: int + status: Optional[ShardTypes] + vector_queue_size: Optional[int] + per_node_status: Optional[Dict[str, str]] ShardStatus = _ShardStatus diff --git a/weaviate/collections/config/executor.py b/weaviate/collections/config/executor.py index 103ab70ac..0ec1523dc 100644 --- a/weaviate/collections/config/executor.py +++ b/weaviate/collections/config/executor.py @@ -365,8 +365,9 @@ def resp(res: Response) -> List[ShardStatus]: return [ _ShardStatus( name=shard["name"], - status=shard["status"], - vector_queue_size=shard["vectorQueueSize"], + status=shard.get("status"), + vector_queue_size=None, + per_node_status=shard.get("per_node_status"), ) for shard in shards ]