nucliadb 6.3.4.post3732__py3-none-any.whl → 6.3.4.post3740__py3-none-any.whl

nucliadb/common/cluster/rollover.py

@@ -31,6 +31,7 @@ from nucliadb.common.external_index_providers.manager import (
     get_external_index_manager,
 )
 from nucliadb.common.nidx import get_nidx_fake_node
+from nucliadb.migrator.settings import settings
 from nucliadb_protos import writer_pb2
 from nucliadb_telemetry import errors
 
@@ -43,6 +44,8 @@ from .utils import (
 
 logger = logging.getLogger(__name__)
 
+resource_index_semaphore = asyncio.Semaphore(settings.max_concurrent_rollover_resources)
+
 
 class UnexpectedRolloverError(Exception):
     pass
@@ -234,10 +237,33 @@ async def index_to_rollover_index(
     # now index on all new shards only
     while True:
         async with datamanagers.with_transaction() as txn:
-            resource_id = await datamanagers.rollover.get_to_index(txn, kbid=kbid)
-        if resource_id is None:
+            resource_ids = await datamanagers.rollover.get_to_index(
+                txn, kbid=kbid, count=settings.max_concurrent_rollover_resources
+            )
+        if resource_ids is None:
             break
 
+        batch = [
+            _index_resource_to_rollover_index(app_context, rollover_shards, kbid, rid, external)
+            for rid in resource_ids
+        ]
+        await asyncio.gather(*batch)
+
+    async with datamanagers.with_transaction() as txn:
+        state.resources_indexed = True
+        await datamanagers.rollover.set_rollover_state(txn, kbid=kbid, state=state)
+        await datamanagers.rollover.update_kb_rollover_shards(txn, kbid=kbid, kb_shards=rollover_shards)
+        await txn.commit()
+
+
+async def _index_resource_to_rollover_index(
+    app_context: ApplicationContext,
+    rollover_shards: writer_pb2.Shards,
+    kbid: str,
+    resource_id: str,
+    external: Optional[ExternalIndexManager] = None,
+) -> None:
+    async with resource_index_semaphore:
         async with datamanagers.with_transaction() as txn:
             shard_id = await datamanagers.resources.get_resource_shard_id(
                 txn, kbid=kbid, rid=resource_id
@@ -250,7 +276,7 @@ async def index_to_rollover_index(
             async with datamanagers.with_transaction() as txn:
                 await datamanagers.rollover.remove_to_index(txn, kbid=kbid, resource=resource_id)
                 await txn.commit()
-            continue
+            return
 
         shard = _get_shard(rollover_shards, shard_id)
         if shard is None:  # pragma: no cover
@@ -268,7 +294,7 @@ async def index_to_rollover_index(
             async with datamanagers.with_transaction() as txn:
                 await datamanagers.rollover.remove_to_index(txn, kbid=kbid, resource=resource_id)
                 await txn.commit()
-            continue
+            return
 
         if external is not None:
             await external.index_resource(resource_id, index_message, to_rollover_indexes=True)
@@ -287,12 +313,6 @@ async def index_to_rollover_index(
             )
             await txn.commit()
 
-    async with datamanagers.with_transaction() as txn:
-        state.resources_indexed = True
-        await datamanagers.rollover.set_rollover_state(txn, kbid=kbid, state=state)
-        await datamanagers.rollover.update_kb_rollover_shards(txn, kbid=kbid, kb_shards=rollover_shards)
-        await txn.commit()
-
 
 async def cutover_index(
     app_context: ApplicationContext, kbid: str, external: Optional[ExternalIndexManager] = None

nucliadb/common/datamanagers/rollover.py

@@ -91,11 +91,11 @@ async def add_batch_to_index(txn: Transaction, *, kbid: str, batch: list[str]) -
         await txn.set(key, b"")
 
 
-async def get_to_index(txn: Transaction, *, kbid: str) -> Optional[str]:
+async def get_to_index(txn: Transaction, *, kbid: str, count: int) -> Optional[list[str]]:
     key = KB_ROLLOVER_RESOURCES_TO_INDEX.format(kbid=kbid, resource="")
-    found = [key async for key in txn.keys(key, count=1)]
+    found = [key async for key in txn.keys(key, count=count)]
     if found:
-        return found[0].split("/")[-1]
+        return [f.split("/")[-1] for f in found]
     return None
 
 

nucliadb/migrator/settings.py

@@ -29,3 +29,10 @@ class Settings(pydantic_settings.BaseSettings):
         default=5,
         description="Maximum number of concurrent KB migrations allowed.",
     )
+    max_concurrent_rollover_resources: int = pydantic.Field(
+        default=5,
+        description="Maximum number of concurrent index operations during rollover.",
+    )
+
+
+settings = Settings()

{nucliadb-6.3.4.post3732.dist-info → nucliadb-6.3.4.post3740.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.3.4.post3732
+Version: 6.3.4.post3740
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,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.3.4.post3732
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.4.post3732
-Requires-Dist: nucliadb-protos>=6.3.4.post3732
-Requires-Dist: nucliadb-models>=6.3.4.post3732
-Requires-Dist: nidx-protos>=6.3.4.post3732
+Requires-Dist: nucliadb-telemetry[all]>=6.3.4.post3740
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.4.post3740
+Requires-Dist: nucliadb-protos>=6.3.4.post3740
+Requires-Dist: nucliadb-models>=6.3.4.post3740
+Requires-Dist: nidx-protos>=6.3.4.post3740
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn

{nucliadb-6.3.4.post3732.dist-info → nucliadb-6.3.4.post3740.dist-info}/RECORD

@@ -62,7 +62,7 @@ nucliadb/common/cluster/exceptions.py,sha256=t7v_l93t44l2tQpdQXgO_w-c4YZRcaayOz1
 nucliadb/common/cluster/grpc_node_dummy.py,sha256=LxONv0mhDFhx7mI91qqGfQlQ-R0qOGDYaxhXoBHLXaE,3548
 nucliadb/common/cluster/manager.py,sha256=KIzqAYGgdVK3GicJ9LdLoei8arWZ7H60imbc32USPj4,12754
 nucliadb/common/cluster/rebalance.py,sha256=cLUlR08SsqmnoA_9GDflV6k2tXmkAPpyFxZErzp45vo,8754
-nucliadb/common/cluster/rollover.py,sha256=iTJ9EQmHbzXL34foNFto-hqdC0Kq1pF1mNxqv0jqhBs,25362
+nucliadb/common/cluster/rollover.py,sha256=QCvns28pr81geSQvC37dGKa-nYGu4_IfZwJxyzKg6z4,26042
 nucliadb/common/cluster/settings.py,sha256=JPwV_0U_i618Tn66GWUq6qCKNjy4TWkGEGld9GwH5uk,2048
 nucliadb/common/cluster/utils.py,sha256=7nQvnVFxM4XV7J560R8hUA-GPzrgD19UlQxHrl4mZUc,4687
 nucliadb/common/cluster/standalone/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
@@ -79,7 +79,7 @@ nucliadb/common/datamanagers/kb.py,sha256=P7EhF4tApIUG2jw_HH1oMufTKG9__kuOLKnrCN
 nucliadb/common/datamanagers/labels.py,sha256=Zm0GQpSPoGXEEysUY7VsDIcyKSIIQsMVphj23IyM9_c,4502
 nucliadb/common/datamanagers/processing.py,sha256=ByxdZzdbAfJGqC6__mY-zryjk040TyQfcUq3rxujeoY,1587
 nucliadb/common/datamanagers/resources.py,sha256=cuwcVL-GEjS1VHigJtz5SG_dRhl09UADQ9MSYOOqs70,10779
-nucliadb/common/datamanagers/rollover.py,sha256=c_DE3jtZusNL_9aOVjHOB9PV5OSVg7GJ5J-Ny0goHBE,7833
+nucliadb/common/datamanagers/rollover.py,sha256=BM1hJ2cEU91xekM5PtmnA0SN3i3w0WmodiyTpO8YZZs,7865
 nucliadb/common/datamanagers/search_configurations.py,sha256=O-8eW43CE46GcxO6TB5hpi27NBguv4BL4SI1vLlN8os,2463
 nucliadb/common/datamanagers/synonyms.py,sha256=zk3GEH38KF5vV_VcuL6DCg-2JwgXJfQl7Io6VPqv2cw,1566
 nucliadb/common/datamanagers/utils.py,sha256=McHlXvE4P3x-bBY3pr0n8djbTDQvI1G5WusJrnRdhLA,1827
@@ -164,7 +164,7 @@ nucliadb/migrator/datamanager.py,sha256=sPrFvgn0aOyWKUUkwaaUmAqqn3jKlVe1IpIXrkoK
 nucliadb/migrator/exceptions.py,sha256=jTj3YhKmFwUyjjgoKUNoCAiGrpEbB64X1Um212nSNQ8,889
 nucliadb/migrator/migrator.py,sha256=tpdKvqqB0p1klcX2LjUD0Br6brcxxfJMVMMoKABmPD4,10834
 nucliadb/migrator/models.py,sha256=3PJkL2PGvKgIG0KIBv4H5XCsOVmwWMlRV3m0ntDj10Q,1145
-nucliadb/migrator/settings.py,sha256=jOUX0ZMunCXN8HpF9xXN0aunJYRhu4Vdr_ffjRIqwtw,1144
+nucliadb/migrator/settings.py,sha256=3eK6PIwqqtoCea9V7-RMjVfZC_0BC5DqPPlfo9XMyNE,1340
 nucliadb/migrator/utils.py,sha256=NgUreUvON8_nWEzTxELBMWlfV7E6-6qi-g0DMEbVEz4,2885
 nucliadb/models/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/models/responses.py,sha256=qnuOoc7TrVSUnpikfTwHLKez47_DE4mSFzpxrwtqijA,1599
@@ -354,8 +354,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.3.4.post3732.dist-info/METADATA,sha256=RR7gEyd40e4MNsR8vtWeJ7EcTRfLRPL22oE4FVF9LUI,4291
-nucliadb-6.3.4.post3732.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-nucliadb-6.3.4.post3732.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.4.post3732.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.4.post3732.dist-info/RECORD,,
+nucliadb-6.3.4.post3740.dist-info/METADATA,sha256=kXv8CG9jprFzMIeB6zi4A3R9gUN7RQ9J2CiTGBolbu0,4291
+nucliadb-6.3.4.post3740.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+nucliadb-6.3.4.post3740.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.3.4.post3740.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.3.4.post3740.dist-info/RECORD,,