feat: Check per-node shard status in wait_for_vector_indexing

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(

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]

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
         ]
 

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

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
             ]