nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/cluster/rebalance.py

@@ -25,6 +25,7 @@ from nucliadb.common.cluster.manager import choose_node
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.context import ApplicationContext
 from nucliadb_protos import nodereader_pb2, noderesources_pb2
+from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
 from nucliadb_telemetry.utils import setup_telemetry
 from nucliadb_utils import const
@@ -36,33 +37,33 @@ from .utils import delete_resource_from_shard, index_resource_to_shard, wait_for
 
 logger = logging.getLogger(__name__)
 
+REBALANCE_LOCK = "rebalance"
+
 
 async def get_shards_paragraphs(kbid: str) -> list[tuple[str, int]]:
     """
     Ordered shard -> num paragraph by number of paragraphs
     """
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
     if kb_shards is None:
         return []
 
     results = {}
     for shard_meta in kb_shards.shards:
-        node, shard_id = choose_node(shard_meta)
+        # Rebalance using node as source of truth. But it will rebalance nidx
+        node, shard_id = choose_node(shard_meta, use_nidx=False)
         shard_data: nodereader_pb2.Shard = await node.reader.GetShard(
             nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
         )
         results[shard_meta.shard] = shard_data.paragraphs
 
-    return [
-        (shard, paragraphs)
-        for shard, paragraphs in sorted(results.items(), key=lambda x: x[1])
-    ]
+    return [(shard, paragraphs) for shard, paragraphs in sorted(results.items(), key=lambda x: x[1])]
 
 
 async def maybe_add_shard(kbid: str) -> None:
     async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=kbid)):
-        async with datamanagers.with_transaction() as txn:
+        async with datamanagers.with_ro_transaction() as txn:
             kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
         if kb_shards is None:
             return
@@ -75,14 +76,8 @@ async def maybe_add_shard(kbid: str) -> None:
         ):
             # create new shard
             async with datamanagers.with_transaction() as txn:
-                model = await datamanagers.kb.get_model_metadata(txn, kbid=kbid)
                 sm = get_shard_manager()
-                await sm.create_shard_by_kbid(
-                    txn,
-                    kbid,
-                    semantic_model=model,
-                    release_channel=kb_shards.release_channel,
-                )
+                await sm.create_shard_by_kbid(txn, kbid)
                 await txn.commit()
 
 
@@ -93,15 +88,13 @@ async def move_set_of_kb_resources(
     to_shard_id: str,
     count: int = 20,
 ) -> None:
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
     if kb_shards is None:  # pragma: no cover
-        logger.warning(
-            "No shards found for kb. This should not happen.", extra={"kbid": kbid}
-        )
+        logger.warning("No shards found for kb. This should not happen.", extra={"kbid": kbid})
         return
 
-    logger.warning(
+    logger.info(
         "Rebalancing kb shards",
         extra={"kbid": kbid, "from": from_shard_id, "to": to_shard_id, "count": count},
     )
@@ -109,7 +102,7 @@ async def move_set_of_kb_resources(
     from_shard = [s for s in kb_shards.shards if s.shard == from_shard_id][0]
     to_shard = [s for s in kb_shards.shards if s.shard == to_shard_id][0]
 
-    from_node, from_shard_replica_id = choose_node(from_shard)
+    from_node, from_shard_replica_id = choose_node(from_shard, use_nidx=False)
     search_response: nodereader_pb2.SearchResponse = await from_node.reader.Search(  # type: ignore
         nodereader_pb2.SearchRequest(
             shard=from_shard_replica_id,
@@ -123,11 +116,14 @@ async def move_set_of_kb_resources(
     for result in search_response.document.results:
         resource_id = result.uuid
         try:
-            async with datamanagers.with_transaction() as txn, locking.distributed_lock(
-                locking.RESOURCE_INDEX_LOCK.format(kbid=kbid, resource_id=resource_id)
+            async with (
+                datamanagers.with_transaction() as txn,
+                locking.distributed_lock(
+                    locking.RESOURCE_INDEX_LOCK.format(kbid=kbid, resource_id=resource_id)
+                ),
             ):
                 found_shard_id = await datamanagers.resources.get_resource_shard_id(
-                    txn, kbid=kbid, rid=resource_id
+                    txn, kbid=kbid, rid=resource_id, for_update=True
                 )
                 if found_shard_id is None:
                     # resource deleted
@@ -174,9 +170,7 @@ async def rebalance_kb(context: ApplicationContext, kbid: str) -> None:
 
     shard_paragraphs = await get_shards_paragraphs(kbid)
     rebalanced_shards = set()
-    while any(
-        paragraphs > settings.max_shard_paragraphs for _, paragraphs in shard_paragraphs
-    ):
+    while any(paragraphs > settings.max_shard_paragraphs for _, paragraphs in shard_paragraphs):
         # find the shard with the least/most paragraphs
         smallest_shard = shard_paragraphs[0][0]
         largest_shard = shard_paragraphs[-1][0]
@@ -195,14 +189,20 @@ async def rebalance_kb(context: ApplicationContext, kbid: str) -> None:
 
 
 async def run(context: ApplicationContext) -> None:
-    async with locking.distributed_lock("rebalance"):
-        # go through each kb and see if shards need to be reduced in size
-        async with datamanagers.with_transaction() as txn:
-            async for kbid, _ in datamanagers.kb.get_kbs(txn):
-                async with locking.distributed_lock(
-                    locking.KB_SHARDS_LOCK.format(kbid=kbid)
-                ):
+    try:
+        async with locking.distributed_lock(REBALANCE_LOCK):
+            # get all kb ids
+            async with datamanagers.with_ro_transaction() as txn:
+                kbids = [kbid async for kbid, _ in datamanagers.kb.get_kbs(txn)]
+            # go through each kb and see if shards need to be reduced in size
+            for kbid in kbids:
+                async with locking.distributed_lock(locking.KB_SHARDS_LOCK.format(kbid=kbid)):
                     await rebalance_kb(context, kbid)
+    except locking.ResourceLocked as exc:
+        if exc.key == REBALANCE_LOCK:
+            logger.warning("Another rebalance process is already running.")
+            return
+        raise
 
 
 async def run_command(context: ApplicationContext) -> None:
@@ -215,13 +215,18 @@ async def run_command(context: ApplicationContext) -> None:
 
     try:
         await run(context)
-    except (asyncio.CancelledError, RuntimeError):
+    except (asyncio.CancelledError, RuntimeError):  # pragma: no cover
         return
-    except Exception:
+    except Exception as ex:  # pragma: no cover
         logger.exception("Failed to run rebalancing.")
+        errors.capture_exception(ex)
     finally:
-        await context.finalize()
-        await metrics_server.shutdown()
+        try:
+            await metrics_server.shutdown()
+            await context.finalize()
+        except Exception:  # pragma: no cover
+            logger.exception("Error tearing down utilities on rebalance command")
+            pass
 
 
 def main():