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

migrations/0027_rollover_texts3.py

@@ -27,7 +27,6 @@ import logging
 
 from nucliadb import learning_proxy
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.rollover import rollover_kb_index
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -38,7 +37,10 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     await maybe_fix_vector_dimensions(context, kbid)
-    await rollover_kb_index(context, kbid)
+
+    # We only need 1 rollover migration defined at a time; otherwise, we will
+    # possibly run many for a kb when we only ever need to run one
+    # await rollover_kb_index(context, kbid)
 
 
 async def maybe_fix_vector_dimensions(context: ExecutionContext, kbid: str) -> None:

migrations/0033_rollover_nidx_relation_2.py

@@ -0,0 +1,39 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+
+"""Migration #33
+
+Rollover to update to nidx_relation v2
+
+"""
+
+import logging
+
+from nucliadb.common.cluster.rollover import rollover_kb_index
+from nucliadb.migrator.context import ExecutionContext
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    await rollover_kb_index(context, kbid)

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.post3729.dist-info → nucliadb-6.3.4.post3740.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.3.4.post3729
+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.post3729
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.4.post3729
-Requires-Dist: nucliadb-protos>=6.3.4.post3729
-Requires-Dist: nucliadb-models>=6.3.4.post3729
-Requires-Dist: nidx-protos>=6.3.4.post3729
+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.post3729.dist-info → nucliadb-6.3.4.post3740.dist-info}/RECORD

@@ -22,12 +22,13 @@ migrations/0022_fix_paragraph_deletion_bug.py,sha256=-tH342VXF-8xwc_h3P1cYaUtTT1
 migrations/0023_backfill_pg_catalog.py,sha256=gw22pU5cAtg2a7n7xVaVqT2itjAoDMNtzWwTllwqIvg,2993
 migrations/0025_assign_models_to_kbs_v2.py,sha256=QC6nDF2Wyc6zQMqNoKzvz-3507UpDyJztlbIsvlwHss,4678
 migrations/0026_fix_high_cardinality_content_types.py,sha256=BsbBkvZDzjRHQfoouZNNtHA1xMxTKm8wOVnp_WAS9j4,2322
-migrations/0027_rollover_texts3.py,sha256=UQDaMOayVuqDisf82NDrPStoEVveHvdjkSmzbIcU9o4,2730
+migrations/0027_rollover_texts3.py,sha256=J7tRVr-OooLuZCFkhYquOPHFRMiHNQT1klT8r0d8a_8,2815
 migrations/0028_extracted_vectors_reference.py,sha256=49DHCIlBpjofU8cYVHTdWv0EBIlnPTWV2WCezf0rJUo,2392
 migrations/0029_backfill_field_status.py,sha256=QWF69n1da9lpRnbEpgbqPjSQ-Wfn6rMC7Enz6bBYGt4,5663
 migrations/0030_label_deduplication.py,sha256=y14TxtCMi3-TBMz_eZoyyPDHNlZb29taJujlDuHumsA,2008
 migrations/0031_languages_deduplication.py,sha256=o6va6lP3oTRT1uSzp5MIhHHBFbhCxSZ-oNlXXpiAdUo,2340
 migrations/0032_remove_old_relations.py,sha256=ZaswhmRRsLgw6DVYVdT7cP-gdBf4X3PL9fklUXrmSD0,2318
+migrations/0033_rollover_nidx_relation_2.py,sha256=LLhw7LknMVl1mu_mQlsijVUYDe7F87eZx2N8IyeZUK8,1244
 migrations/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 migrations/pg/0001_bootstrap.py,sha256=Fsqkeof50m7fKiJN05kmNEMwiKDlOrAgcAS5sLLkutA,1256
 migrations/pg/0002_catalog.py,sha256=Rsleecu351Ty19kYZgOpqX5G3MEAY8nMxCJrAeuS2Mw,1690
@@ -61,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
@@ -78,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
@@ -163,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
@@ -353,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.post3729.dist-info/METADATA,sha256=7WsHIY6R-JkCVTyVFbfCWk-avqWWNy_lJEYRNtow_F4,4291
-nucliadb-6.3.4.post3729.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
-nucliadb-6.3.4.post3729.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.4.post3729.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.4.post3729.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,,