nucliadb 6.4.2.post4376__py3-none-any.whl → 6.4.2.post4378__py3-none-any.whl

nucliadb/train/generators/field_streaming.py

@@ -18,14 +18,16 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import AsyncGenerator, Optional
+import asyncio
+from typing import AsyncGenerator, AsyncIterable, Optional
 
-from nidx_protos.nodereader_pb2 import StreamRequest
+from nidx_protos.nodereader_pb2 import DocumentItem, StreamRequest
 
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
+from nucliadb.train.settings import settings
 from nucliadb_protos.dataset_pb2 import (
     FieldSplitData,
     FieldStreamingBatch,
@@ -50,82 +52,38 @@ async def generate_field_streaming_payloads(
     trainset: TrainSet,
     shard_replica_id: str,
 ) -> AsyncGenerator[FieldSplitData, None]:
-    # Query how many resources has each label
     request = StreamRequest()
     request.shard_id.id = shard_replica_id
 
     for label in trainset.filter.labels:
         request.filter.labels.append(f"/l/{label}")
-
     for path in trainset.filter.paths:
         request.filter.labels.append(f"/p/{path}")
-
     for metadata in trainset.filter.metadata:
         request.filter.labels.append(f"/m/{metadata}")
-
     for entity in trainset.filter.entities:
         request.filter.labels.append(f"/e/{entity}")
-
     for field in trainset.filter.fields:
         request.filter.labels.append(f"/f/{field}")
-
     for status in trainset.filter.status:
         request.filter.labels.append(f"/n/s/{status}")
 
     resources = set()
     fields = set()
 
-    async for document_item in get_nidx_searcher_client().Documents(request):
-        text_labels = []
-        for label in document_item.labels:
-            text_labels.append(label)
-
-        field_id = f"{document_item.uuid}{document_item.field}"
-        resources.add(document_item.uuid)
-
-        field_parts = document_item.field.split("/")
-        if len(field_parts) == 3:
-            _, field_type, field = field_parts
-            split = "0"
-        elif len(field_parts) == 4:
-            _, field_type, field, split = field_parts
-        else:
-            raise Exception(f"Invalid field definition {document_item.field}")
-
-        tl = FieldSplitData()
-        rid, field_type, field = field_id.split("/")
-        tl.rid = document_item.uuid
-        tl.field = field
-        tl.field_type = field_type
-        tl.split = split
-
-        field_unique_key = f"{rid}/{field_type}/{field}/{split}"
+    async for fsd in iter_field_split_data(
+        request, kbid, trainset, max_parallel=settings.field_streaming_parallelisation
+    ):
+        resources.add(fsd.rid)
+        field_unique_key = f"{fsd.rid}/{fsd.field_type}/{fsd.field}/{fsd.split}"
         if field_unique_key in fields:
             # This field has already been yielded. This can happen as we are streaming directly from nidx
             # and field deletions may not be reflected immediately in the index.
             logger.warning(f"Duplicated field found {field_unique_key}. Skipping.", extra={"kbid": kbid})
             continue
-
         fields.add(field_unique_key)
 
-        if trainset.exclude_text:
-            tl.text.text = ""
-        else:
-            extracted = await get_field_text(kbid, rid, field, field_type)
-            if extracted is not None:
-                tl.text.CopyFrom(extracted)
-
-        metadata_obj = await get_field_metadata(kbid, rid, field, field_type)
-        if metadata_obj is not None:
-            tl.metadata.CopyFrom(metadata_obj)
-
-        basic = await get_field_basic(kbid, rid, field, field_type)
-        if basic is not None:
-            tl.basic.CopyFrom(basic)
-
-        tl.labels.extend(text_labels)
-
-        yield tl
+        yield fsd
 
         if len(fields) % 1000 == 0:
             logger.info(
@@ -149,6 +107,72 @@ async def generate_field_streaming_payloads(
     )
 
 
+async def iter_field_split_data(
+    request: StreamRequest, kbid: str, trainset: TrainSet, max_parallel: int = 5
+) -> AsyncIterable[FieldSplitData]:
+    tasks: list[asyncio.Task] = []
+    async for document_item in get_nidx_searcher_client().Documents(request):
+        if len(tasks) >= max_parallel:
+            results = await asyncio.gather(*tasks)
+            for fsd in results:
+                yield fsd
+            tasks.clear()
+        tasks.append(asyncio.create_task(fetch_field_split_data(document_item, kbid, trainset)))
+    if len(tasks):
+        results = await asyncio.gather(*tasks)
+        for fsd in results:
+            yield fsd
+        tasks.clear()
+
+
+async def fetch_field_split_data(
+    document_item: DocumentItem, kbid: str, trainset: TrainSet
+) -> FieldSplitData:
+    field_id = f"{document_item.uuid}{document_item.field}"
+    field_parts = document_item.field.split("/")
+    if len(field_parts) == 3:
+        _, field_type, field = field_parts
+        split = "0"
+    elif len(field_parts) == 4:
+        _, field_type, field, split = field_parts
+    else:
+        raise Exception(f"Invalid field definition {document_item.field}")
+    _, field_type, field = field_id.split("/")
+    fsd = FieldSplitData()
+    fsd.rid = document_item.uuid
+    fsd.field = field
+    fsd.field_type = field_type
+    fsd.split = split
+    tasks = []
+    if trainset.exclude_text:
+        fsd.text.text = ""
+    else:
+        tasks.append(asyncio.create_task(_fetch_field_extracted_text(kbid, fsd)))
+    tasks.append(asyncio.create_task(_fetch_field_metadata(kbid, fsd)))
+    tasks.append(asyncio.create_task(_fetch_basic(kbid, fsd)))
+    await asyncio.gather(*tasks)
+    fsd.labels.extend(document_item.labels)
+    return fsd
+
+
+async def _fetch_field_extracted_text(kbid: str, fsd: FieldSplitData):
+    extracted = await get_field_text(kbid, fsd.rid, fsd.field, fsd.field_type)
+    if extracted is not None:
+        fsd.text.CopyFrom(extracted)
+
+
+async def _fetch_field_metadata(kbid: str, fsd: FieldSplitData):
+    metadata_obj = await get_field_metadata(kbid, fsd.rid, fsd.field, fsd.field_type)
+    if metadata_obj is not None:
+        fsd.metadata.CopyFrom(metadata_obj)
+
+
+async def _fetch_basic(kbid: str, fsd: FieldSplitData):
+    basic = await get_field_basic(kbid, fsd.rid, fsd.field, fsd.field_type)
+    if basic is not None:
+        fsd.basic.CopyFrom(basic)
+
+
 async def get_field_text(kbid: str, rid: str, field: str, field_type: str) -> Optional[ExtractedText]:
     orm_resource = await get_resource_from_cache_or_db(kbid, rid)
 

nucliadb/train/settings.py

@@ -34,6 +34,7 @@ class Settings(DriverSettings):
     internal_search_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/search"
 
     resource_cache_size: int = 2
+    field_streaming_parallelisation: int = 5
 
 
 settings = Settings()

{nucliadb-6.4.2.post4376.dist-info → nucliadb-6.4.2.post4378.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.2.post4376
+Version: 6.4.2.post4378
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.4.2.post4376
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.2.post4376
-Requires-Dist: nucliadb-protos>=6.4.2.post4376
-Requires-Dist: nucliadb-models>=6.4.2.post4376
-Requires-Dist: nidx-protos>=6.4.2.post4376
+Requires-Dist: nucliadb-telemetry[all]>=6.4.2.post4378
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.2.post4378
+Requires-Dist: nucliadb-protos>=6.4.2.post4378
+Requires-Dist: nucliadb-models>=6.4.2.post4378
+Requires-Dist: nidx-protos>=6.4.2.post4378
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.4.2.post4376.dist-info → nucliadb-6.4.2.post4378.dist-info}/RECORD

@@ -309,7 +309,7 @@ nucliadb/train/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/train/resource.py,sha256=3qQ_9Zdt5JAbtD-wpmt7OeDGRNKS-fQdKAuIQfznZm0,16219
 nucliadb/train/run.py,sha256=evz6CKVfJOzkbHMoaYz2mTMlKjJnNOb1O8zBBWMpeBw,1400
 nucliadb/train/servicer.py,sha256=scbmq8FriKsJGkOcoZB2Fg_IyIExn9Ux4W30mGDlkJQ,5728
-nucliadb/train/settings.py,sha256=8_-XCO_nrE98cMJpe0fYkXeF2bkWKff1VX-2qdwcsjs,1417
+nucliadb/train/settings.py,sha256=Vz-bQxwxYg6Qhc8Vnap95AwlYyCE1LF7NCPlLBfToXI,1462
 nucliadb/train/types.py,sha256=xyVYy8kHipAWoDb7Pn7dCYQ_efHPzDW_3AXg5M-aV28,1519
 nucliadb/train/upload.py,sha256=fTjH1KEL-0ogf3LV0T6ODO0QdPGwdZShSUtFUCAcUlA,3256
 nucliadb/train/uploader.py,sha256=xdLGz1ToDue9Q_M8A-_KYkO-V6fWKYOZQ6IGM4FuwWA,6424
@@ -322,7 +322,7 @@ nucliadb/train/api/v1/shards.py,sha256=GJRnQe8P-7_VTIN1oxVmxlrDA08qVN7opEZdbF4Wx
 nucliadb/train/api/v1/trainset.py,sha256=kpnpDgiMWr1FKHZJgwH7hue5kzilA8-i9X0YHlNeHuU,2113
 nucliadb/train/generators/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/train/generators/field_classifier.py,sha256=xUA10o9CtBtilbP3uc-8Wn_zQ0oK3BrqYGqZgxh4ZLk,3428
-nucliadb/train/generators/field_streaming.py,sha256=nje317SutX8QmHq-xwUphzUiozmzpCRfPXxhF_jFzdg,6441
+nucliadb/train/generators/field_streaming.py,sha256=fq4XxHf5fPLccPjO722tA9Xcw6smmceVqSme0fY2_NA,7631
 nucliadb/train/generators/image_classifier.py,sha256=BDXgyd5TGZRnzDnVRvp-qsRCuoTbTYwui3JiDIjuiDc,1736
 nucliadb/train/generators/paragraph_classifier.py,sha256=4sH3IQc7yJrlDs1C76SxFzL9N5mXWRZzJzoiF7y4dSQ,2703
 nucliadb/train/generators/paragraph_streaming.py,sha256=1xsc_IqP-1M0TzYTqu5qCvWBNp_J3Kyvnx8HVbToXmQ,3532
@@ -368,8 +368,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.4.2.post4376.dist-info/METADATA,sha256=FWD9EeIPDvZWgRX4R33CaHbOrlk2YYzGCtrnMGaqR0s,4152
-nucliadb-6.4.2.post4376.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.4.2.post4376.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.2.post4376.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.2.post4376.dist-info/RECORD,,
+nucliadb-6.4.2.post4378.dist-info/METADATA,sha256=W_SC_iY4NnfaM04p0SUiwwdGvfgm4eDKNVNaWXU2mp8,4152
+nucliadb-6.4.2.post4378.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.4.2.post4378.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.2.post4378.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.2.post4378.dist-info/RECORD,,