nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/common/cluster/manager.py

@@ -23,7 +23,6 @@ import uuid
 from typing import Any, Awaitable, Callable, Optional
 
 import backoff
-from nucliadb_protos.nodewriter_pb2 import IndexMessage, IndexMessageSource, TypeMessage
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.base import AbstractIndexNode
@@ -37,12 +36,15 @@ from nucliadb.common.cluster.exceptions import (
     ShardsNotFound,
 )
 from nucliadb.common.maindb.driver import Transaction
+from nucliadb.common.nidx import NIDX_ENABLED, get_nidx, get_nidx_api_client, get_nidx_fake_node
 from nucliadb_protos import (
+    knowledgebox_pb2,
     nodereader_pb2,
     noderesources_pb2,
     nodewriter_pb2,
     writer_pb2,
 )
+from nucliadb_protos.nodewriter_pb2 import IndexMessage, IndexMessageSource, NewShardRequest, TypeMessage
 from nucliadb_telemetry import errors
 from nucliadb_utils.utilities import get_indexing, get_storage
 
@@ -124,7 +126,7 @@ def remove_index_node(node_id: str, primary_id: Optional[str] = None) -> None:
 class KBShardManager:
     # TODO: move to data manager
     async def get_shards_by_kbid_inner(self, kbid: str) -> writer_pb2.Shards:
-        async with datamanagers.with_transaction(read_only=True) as txn:
+        async with datamanagers.with_ro_transaction() as txn:
             result = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
             if result is None:
                 # could be None because /shards doesn't exist, or beacause the
@@ -142,6 +144,8 @@ class KBShardManager:
         kbid: str,
         aw: Callable[[AbstractIndexNode, str], Awaitable[Any]],
         timeout: float,
+        *,
+        use_nidx: bool,
         use_read_replica_nodes: bool = False,
     ) -> list[Any]:
         shards = await self.get_shards_by_kbid(kbid)
@@ -149,7 +153,7 @@ class KBShardManager:
 
         for shard_obj in shards:
             node, shard_id = choose_node(
-                shard_obj, use_read_replica_nodes=use_read_replica_nodes
+                shard_obj, use_nidx=use_nidx, use_read_replica_nodes=use_read_replica_nodes
             )
             if shard_id is None:
                 raise ShardNotFound("Found a node but not a shard")
@@ -158,7 +162,7 @@ class KBShardManager:
 
         try:
             results = await asyncio.wait_for(
-                asyncio.gather(*ops, return_exceptions=True),  # type: ignore
+                asyncio.gather(*ops, return_exceptions=True),
                 timeout=timeout,
             )
         except asyncio.TimeoutError as exc:
@@ -171,7 +175,7 @@ class KBShardManager:
     async def get_current_active_shard(
         self, txn: Transaction, kbid: str
     ) -> Optional[writer_pb2.ShardObject]:
-        kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
+        kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=False)
         if kb_shards is None:
             return None
 
@@ -195,23 +199,25 @@ class KBShardManager:
             )
             raise
 
-        kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
+        kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
         if kb_shards is None:
-            msg = (
-                "Attempting to create a shard for a KB when it has no stored shards in maindb",
-            )
+            msg = ("Attempting to create a shard for a KB when it has no stored shards in maindb",)
             logger.error(msg, extra={"kbid": kbid})
             raise ShardsNotFound(msg)
 
-        existing_kb_nodes = [
-            replica.node for shard in kb_shards.shards for replica in shard.replicas
-        ]
+        existing_kb_nodes = [replica.node for shard in kb_shards.shards for replica in shard.replicas]
         nodes = sorted_primary_nodes(
             avoid_nodes=existing_kb_nodes,
             ignore_nodes=settings.drain_nodes,
         )
 
+        vectorsets = {
+            vectorset_id: vectorset_config.vectorset_index_config
+            async for vectorset_id, vectorset_config in datamanagers.vectorsets.iter(txn, kbid=kbid)
+        }
+
         shard_uuid = uuid.uuid4().hex
+
         shard = writer_pb2.ShardObject(shard=shard_uuid, read_only=False)
         try:
             # Attempt to create configured number of replicas
@@ -228,28 +234,56 @@ class KBShardManager:
                 if node is None:
                     logger.error(f"Node {node_id} is not found or not available")
                     continue
-                is_matryoshka = len(kb_shards.model.matryoshka_dimensions) > 0
+
                 try:
-                    shard_created = await node.new_shard(
-                        kbid,
-                        similarity=kb_shards.similarity,
-                        release_channel=kb_shards.release_channel,
-                        normalize_vectors=is_matryoshka,
+                    if not vectorsets:
+                        # bw/c KBs without vectorsets
+                        is_matryoshka = len(kb_shards.model.matryoshka_dimensions) > 0
+                        vector_index_config = nodewriter_pb2.VectorIndexConfig(
+                            similarity=kb_shards.similarity,
+                            vector_type=nodewriter_pb2.VectorType.DENSE_F32,
+                            vector_dimension=kb_shards.model.vector_dimension,
+                            normalize_vectors=is_matryoshka,
+                        )
+
+                        shard_created = await node.new_shard(
+                            kbid,
+                            vector_index_config=vector_index_config,
+                        )
+
+                    else:
+                        shard_created = await node.new_shard_with_vectorsets(
+                            kbid,
+                            vectorsets_configs=vectorsets,
+                        )
+
+                except Exception as exc:
+                    errors.capture_exception(exc)
+                    logger.exception(
+                        f"Error creating new shard for KB", extra={"kbid": kbid, "node_id": node}
                     )
-                except Exception as e:
-                    errors.capture_exception(e)
-                    logger.exception(f"Error creating new shard at {node}: {e}")
                     continue
 
                 replica = writer_pb2.ShardReplica(node=str(node_id))
                 replica.shard.CopyFrom(shard_created)
                 shard.replicas.append(replica)
                 replicas_created += 1
-        except Exception as e:
-            errors.capture_exception(e)
-            logger.error(f"Unexpected error creating new shard: {e}")
+
+                nidx_api = get_nidx_api_client()
+                if nidx_api:
+                    req = NewShardRequest(
+                        kbid=kbid,
+                        vectorsets_configs=vectorsets,
+                    )
+
+                    resp = await nidx_api.NewShard(req)  # type: ignore
+                    shard.nidx_shard_id = resp.id
+
+        except Exception as exc:
+            errors.capture_exception(exc)
+            logger.exception(f"Unexpected error creating new shard for KB", extra={"kbid": kbid})
             await self.rollback_shard(shard)
-            raise e
+            raise exc
 
         # set previous shard as read only, we only have one writable shard at a
         # time
@@ -284,6 +318,17 @@ class KBShardManager:
                         exc_info=True,
                     )
 
+        nidx_api = get_nidx_api_client()
+        if nidx_api and shard.nidx_shard_id:
+            try:
+                await nidx_api.DeleteShard(noderesources_pb2.ShardId(id=shard.nidx_shard_id))
+            except Exception as rollback_error:
+                errors.capture_exception(rollback_error)
+                logger.error(
+                    f"New shard rollback error. Nidx Shard: {shard.nidx_shard_id}",
+                    exc_info=True,
+                )
+
     def indexing_replicas(self, shard: writer_pb2.ShardObject) -> list[tuple[str, str]]:
         """
         Returns the replica ids and nodes for the shard replicas
@@ -303,10 +348,9 @@ class KBShardManager:
     ) -> None:
         indexing = get_indexing()
         storage = await get_storage()
+        nidx = get_nidx()
 
-        await storage.delete_indexing(
-            resource_uid=uuid, txid=txid, kb=kb, logical_shard=shard.shard
-        )
+        await storage.delete_indexing(resource_uid=uuid, txid=txid, kb=kb, logical_shard=shard.shard)
 
         for replica_id, node_id in self.indexing_replicas(shard):
             indexpb: nodewriter_pb2.IndexMessage = nodewriter_pb2.IndexMessage()
@@ -319,6 +363,13 @@ class KBShardManager:
             indexpb.kbid = kb
             await indexing.index(indexpb, node_id)
 
+        if nidx is not None and shard.nidx_shard_id:
+            nidxpb: nodewriter_pb2.IndexMessage = nodewriter_pb2.IndexMessage()
+            nidxpb.shard = shard.nidx_shard_id
+            nidxpb.resource = uuid
+            nidxpb.typemessage = nodewriter_pb2.TypeMessage.DELETION
+            await nidx.index(nidxpb)
+
     async def add_resource(
         self,
         shard: writer_pb2.ShardObject,
@@ -339,7 +390,7 @@ class KBShardManager:
 
         storage = await get_storage()
         indexing = get_indexing()
-
+        nidx = get_nidx()
         indexpb = IndexMessage()
 
         if reindex_id is not None:
@@ -366,6 +417,10 @@ class KBShardManager:
             indexpb.shard = replica_id
             await indexing.index(indexpb, node_id)
 
+        if nidx is not None and shard.nidx_shard_id:
+            indexpb.shard = shard.nidx_shard_id
+            await nidx.index(indexpb)
+
     def should_create_new_shard(self, num_paragraphs: int) -> bool:
         return num_paragraphs > settings.max_shard_paragraphs
 
@@ -383,6 +438,44 @@ class KBShardManager:
             await self.create_shard_by_kbid(txn, kbid)
             await txn.commit()
 
+    async def create_vectorset(self, kbid: str, config: knowledgebox_pb2.VectorSetConfig):
+        """Create a new vectorset in all KB shards."""
+
+        async def _create_vectorset(node: AbstractIndexNode, shard_id: str):
+            vectorset_id = config.vectorset_id
+            index_config = config.vectorset_index_config
+            result = await node.add_vectorset(shard_id, vectorset_id, index_config)
+            if result.status != result.Status.OK:
+                raise NodeError(
+                    f"Unable to create vectorset {vectorset_id} in kb {kbid} shard {shard_id}"
+                )
+
+        await self.apply_for_all_shards(
+            kbid, _create_vectorset, timeout=10, use_nidx=False, use_read_replica_nodes=False
+        )
+        if NIDX_ENABLED:
+            await self.apply_for_all_shards(
+                kbid, _create_vectorset, timeout=10, use_nidx=True, use_read_replica_nodes=False
+            )
+
+    async def delete_vectorset(self, kbid: str, vectorset_id: str):
+        """Delete a vectorset from all KB shards"""
+
+        async def _delete_vectorset(node: AbstractIndexNode, shard_id: str):
+            result = await node.remove_vectorset(shard_id, vectorset_id)
+            if result.status != result.Status.OK:
+                raise NodeError(
+                    f"Unable to delete vectorset {vectorset_id} in kb {kbid} shard {shard_id}"
+                )
+
+        await self.apply_for_all_shards(
+            kbid, _delete_vectorset, timeout=10, use_nidx=False, use_read_replica_nodes=False
+        )
+        if NIDX_ENABLED:
+            await self.apply_for_all_shards(
+                kbid, _delete_vectorset, timeout=10, use_nidx=True, use_read_replica_nodes=False
+            )
+
 
 class StandaloneKBShardManager(KBShardManager):
     max_ops_before_checks = 200
@@ -390,11 +483,9 @@ class StandaloneKBShardManager(KBShardManager):
     def __init__(self):
         super().__init__()
         self._lock = asyncio.Lock()
-        self._change_count: dict[tuple[str, str], int] = {}  # type: ignore
+        self._change_count: dict[tuple[str, str], int] = {}
 
-    async def _resource_change_event(
-        self, kbid: str, node_id: str, shard_id: str
-    ) -> None:
+    async def _resource_change_event(self, kbid: str, node_id: str, shard_id: str) -> None:
         if (node_id, shard_id) not in self._change_count:
             self._change_count[(node_id, shard_id)] = 0
         self._change_count[(node_id, shard_id)] += 1
@@ -407,17 +498,15 @@ class StandaloneKBShardManager(KBShardManager):
             if index_node is None:
                 return
             shard_info: noderesources_pb2.Shard = await index_node.reader.GetShard(
-                nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
+                nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))
             )
             await self.maybe_create_new_shard(
                 kbid,
                 shard_info.paragraphs,
             )
-            await index_node.writer.GC(noderesources_pb2.ShardId(id=shard_id))  # type: ignore
+            await index_node.writer.GC(noderesources_pb2.ShardId(id=shard_id))
 
-    @backoff.on_exception(
-        backoff.expo, NodesUnsync, jitter=backoff.random_jitter, max_tries=5
-    )
+    @backoff.on_exception(backoff.expo, NodesUnsync, jitter=backoff.random_jitter, max_tries=5)
     async def delete_resource(
         self,
         shard: writer_pb2.ShardObject,
@@ -433,19 +522,21 @@ class StandaloneKBShardManager(KBShardManager):
             req.shard_id = shardreplica.shard.id
             index_node = get_index_node(shardreplica.node)
             if index_node is None:  # pragma: no cover
-                raise NodesUnsync(
-                    f"Node {shardreplica.node} is not found or not available"
-                )
+                raise NodesUnsync(f"Node {shardreplica.node} is not found or not available")
             await index_node.writer.RemoveResource(req)  # type: ignore
             asyncio.create_task(
-                self._resource_change_event(
-                    kb, shardreplica.node, shardreplica.shard.id
-                )
+                self._resource_change_event(kb, shardreplica.node, shardreplica.shard.id)
             )
 
-    @backoff.on_exception(
-        backoff.expo, NodesUnsync, jitter=backoff.random_jitter, max_tries=5
-    )
+        nidx = get_nidx()
+        if nidx is not None and shard.nidx_shard_id:
+            indexpb: nodewriter_pb2.IndexMessage = nodewriter_pb2.IndexMessage()
+            indexpb.shard = shard.nidx_shard_id
+            indexpb.resource = uuid
+            indexpb.typemessage = nodewriter_pb2.TypeMessage.DELETION
+            await nidx.index(indexpb)
+
+    @backoff.on_exception(backoff.expo, NodesUnsync, jitter=backoff.random_jitter, max_tries=5)
     async def add_resource(
         self,
         shard: writer_pb2.ShardObject,
@@ -465,16 +556,36 @@ class StandaloneKBShardManager(KBShardManager):
             resource.shard_id = resource.resource.shard_id = shardreplica.shard.id
             index_node = get_index_node(shardreplica.node)
             if index_node is None:  # pragma: no cover
-                raise NodesUnsync(
-                    f"Node {shardreplica.node} is not found or not available"
-                )
+                raise NodesUnsync(f"Node {shardreplica.node} is not found or not available")
             await index_node.writer.SetResource(resource)  # type: ignore
             asyncio.create_task(
-                self._resource_change_event(
-                    kb, shardreplica.node, shardreplica.shard.id
-                )
+                self._resource_change_event(kb, shardreplica.node, shardreplica.shard.id)
             )
 
+        nidx = get_nidx()
+        if nidx is not None and shard.nidx_shard_id:
+            storage = await get_storage()
+            indexpb = IndexMessage()
+            storage_key = await storage.indexing(
+                resource, txid, partition, kb=kb, logical_shard=shard.shard
+            )
+
+            indexpb.typemessage = TypeMessage.CREATION
+            indexpb.storage_key = storage_key
+            indexpb.kbid = kb
+            indexpb.source = source
+            indexpb.resource = resource.resource.uuid
+            indexpb.shard = shard.nidx_shard_id
+
+            await nidx.index(indexpb)
+
+            # Delete indexing message (no longer needed)
+            try:
+                if storage.indexing_bucket:
+                    await storage.delete_upload(storage_key, storage.indexing_bucket)
+            except Exception:
+                pass
+
 
 def get_all_shard_nodes(
     shard: writer_pb2.ShardObject,
@@ -506,6 +617,7 @@ def get_all_shard_nodes(
 def choose_node(
     shard: writer_pb2.ShardObject,
     *,
+    use_nidx: bool,
     target_shard_replicas: Optional[list[str]] = None,
     use_read_replica_nodes: bool = False,
 ) -> tuple[AbstractIndexNode, str]:
@@ -521,6 +633,13 @@ def choose_node(
     `target_shard_replicas` is the least preferent.
 
     """
+
+    # Use nidx if requested and enabled, fallback to node
+    if shard.nidx_shard_id and use_nidx:
+        fake_node = get_nidx_fake_node()
+        if fake_node:
+            return fake_node, shard.nidx_shard_id
+
     target_shard_replicas = target_shard_replicas or []
 
     shard_nodes = get_all_shard_nodes(shard, use_read_replicas=use_read_replica_nodes)
@@ -564,9 +683,7 @@ def check_enough_nodes():
         )
     if settings.max_node_replicas >= 0:
         available_nodes = list(
-            filter(
-                lambda n: n.shard_count < settings.max_node_replicas, available_nodes  # type: ignore
-            )
+            filter(lambda n: n.shard_count < settings.max_node_replicas, available_nodes)
         )
         if len(available_nodes) < target_replicas:
             raise NodeClusterSmall(
@@ -598,9 +715,7 @@ def sorted_primary_nodes(
     preferred_nodes = [nid for nid in available_node_ids if nid not in avoid_nodes]
 
     # Add avoid_nodes to the end of the last nodes
-    result_nodes = preferred_nodes + [
-        nid for nid in available_node_ids if nid not in preferred_nodes
-    ]
+    result_nodes = preferred_nodes + [nid for nid in available_node_ids if nid not in preferred_nodes]
 
     # Remove ignore_nodes from the list
     result_nodes = [nid for nid in result_nodes if nid not in ignore_nodes]

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
@@ -43,28 +44,26 @@ async def get_shards_paragraphs(kbid: str) -> list[tuple[str, int]]:
     """
     Ordered shard -> num paragraph by number of paragraphs
     """
-    async with datamanagers.with_transaction(read_only=True) 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(read_only=True) 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
@@ -89,12 +88,10 @@ 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.info(
@@ -105,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,
@@ -122,13 +119,11 @@ async def move_set_of_kb_resources(
             async with (
                 datamanagers.with_transaction() as txn,
                 locking.distributed_lock(
-                    locking.RESOURCE_INDEX_LOCK.format(
-                        kbid=kbid, resource_id=resource_id
-                    )
+                    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
@@ -175,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]
@@ -198,13 +191,13 @@ async def rebalance_kb(context: ApplicationContext, kbid: str) -> None:
 async def run(context: ApplicationContext) -> None:
     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
-            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)
-                    ):
-                        await rebalance_kb(context, kbid)
+            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.")
@@ -222,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():