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

nucliadb/common/cluster/manager.py

@@ -19,13 +19,10 @@
 #
 import asyncio
 import logging
-import random
 import uuid
 from typing import Any, Awaitable, Callable, Optional
 
 import backoff
-from nucliadb_protos.knowledgebox_pb2 import SemanticModelMetadata  # type: ignore
-from nucliadb_protos.nodewriter_pb2 import IndexMessage, IndexMessageSource, TypeMessage
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.base import AbstractIndexNode
@@ -39,13 +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,
-    utils_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
 
@@ -71,6 +70,11 @@ def get_index_node(node_id: str) -> Optional[AbstractIndexNode]:
     return INDEX_NODES.get(node_id)
 
 
+def clear_index_nodes():
+    INDEX_NODES.clear()
+    READ_REPLICA_INDEX_NODES.clear()
+
+
 def get_read_replica_node_ids(node_id: str) -> list[str]:
     return list(READ_REPLICA_INDEX_NODES.get(node_id, set()))
 
@@ -122,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
@@ -140,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)
@@ -147,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")
@@ -156,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:
@@ -169,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
 
@@ -183,8 +189,6 @@ class KBShardManager:
         self,
         txn: Transaction,
         kbid: str,
-        semantic_model: SemanticModelMetadata,
-        release_channel: utils_pb2.ReleaseChannel.ValueType,
     ) -> writer_pb2.ShardObject:
         try:
             check_enough_nodes()
@@ -195,26 +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:
-            # First logic shard on the index
-            kb_shards = writer_pb2.Shards()
-            kb_shards.kbid = kbid
-            # B/c with Shards.actual
-            kb_shards.actual = -1
-            kb_shards.similarity = semantic_model.similarity_function
-            kb_shards.model.CopyFrom(semantic_model)
-        else:
-            # New logic shard on an existing index
-            pass
+            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]
+        nodes = sorted_primary_nodes(
+            avoid_nodes=existing_kb_nodes,
+            ignore_nodes=settings.drain_nodes,
+        )
 
-        kb_shards.release_channel = release_channel
-        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)
+        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
@@ -231,26 +234,56 @@ class KBShardManager:
                 if node is None:
                     logger.error(f"Node {node_id} is not found or not available")
                     continue
+
                 try:
-                    shard_created = await node.new_shard(
-                        kbid,
-                        similarity=kb_shards.similarity,
-                        release_channel=kb_shards.release_channel,
+                    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
@@ -259,8 +292,8 @@ class KBShardManager:
 
         # Append the created shard and make `actual` point to it.
         kb_shards.shards.append(shard)
-        # B/c with Shards.actual
-        kb_shards.actual += 1
+        # B/c with Shards.actual - we only use last created shard
+        kb_shards.actual = len(kb_shards.shards) - 1
 
         await datamanagers.cluster.update_kb_shards(txn, kbid=kbid, shards=kb_shards)
 
@@ -273,7 +306,7 @@ class KBShardManager:
             node = get_index_node(node_id)
             if node is not None:
                 try:
-                    logger.warning(
+                    logger.info(
                         "Deleting shard replica",
                         extra={"shard": replica_id, "node": node_id},
                     )
@@ -285,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
@@ -304,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()
@@ -320,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,
@@ -330,6 +380,9 @@ class KBShardManager:
         reindex_id: Optional[str] = None,
         source: IndexMessageSource.ValueType = IndexMessageSource.PROCESSOR,
     ) -> None:
+        """
+        Stores the Resource object in the object storage and sends an IndexMessage to the indexing Nats stream.
+        """
         if txid == -1 and reindex_id is None:
             # This means we are injecting a complete resource via ingest gRPC
             # outside of a transaction. We need to treat this as a reindex operation.
@@ -337,7 +390,7 @@ class KBShardManager:
 
         storage = await get_storage()
         indexing = get_indexing()
-
+        nidx = get_nidx()
         indexpb = IndexMessage()
 
         if reindex_id is not None:
@@ -364,34 +417,65 @@ class KBShardManager:
             indexpb.shard = replica_id
             await indexing.index(indexpb, node_id)
 
-    def should_create_new_shard(self, num_paragraphs: int, num_fields: int) -> bool:
-        return (
-            num_paragraphs > settings.max_shard_paragraphs
-            or num_fields > settings.max_shard_fields
-        )
+        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
 
     async def maybe_create_new_shard(
         self,
         kbid: str,
         num_paragraphs: int,
-        num_fields: int,
-        release_channel: utils_pb2.ReleaseChannel.ValueType = utils_pb2.ReleaseChannel.STABLE,
     ):
-        if not self.should_create_new_shard(num_paragraphs, num_fields):
+        if not self.should_create_new_shard(num_paragraphs):
             return
 
-        logger.warning({"message": "Adding shard", "kbid": kbid})
+        logger.info({"message": "Adding shard", "kbid": kbid})
 
         async with datamanagers.with_transaction() as txn:
-            model = await datamanagers.kb.get_model_metadata(txn, kbid=kbid)
-            await self.create_shard_by_kbid(
-                txn,
-                kbid,
-                semantic_model=model,
-                release_channel=release_channel,
-            )
+            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
@@ -399,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
@@ -416,19 +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,
-                shard_info.fields,
-                shard_info.metadata.release_channel,
             )
-            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,
@@ -444,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,
@@ -467,21 +547,45 @@ class StandaloneKBShardManager(KBShardManager):
         reindex_id: Optional[str] = None,
         source: IndexMessageSource.ValueType = IndexMessageSource.PROCESSOR,
     ) -> None:
+        """
+        Calls the node writer's SetResource method directly to store the resource in the node.
+        There is no queuing for standalone nodes at the moment -- indexing is done synchronously.
+        """
         index_node = None
         for shardreplica in shard.replicas:
             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,
@@ -513,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]:
@@ -528,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)
@@ -550,7 +662,10 @@ def choose_node(
         ranked_nodes.setdefault(score, []).append((node, shard_replica_id))
 
     top = ranked_nodes[max(ranked_nodes)]
-    selected_node, shard_replica_id = random.choice(top)
+    # As shard replica ids are random numbers, we sort by shard replica id and choose its
+    # node to make sure we choose in deterministically but we don't favour any node in particular
+    top.sort(key=lambda x: x[1])
+    selected_node, shard_replica_id = top[0]
     return selected_node, shard_replica_id
 
 
@@ -558,17 +673,17 @@ def check_enough_nodes():
     """
     It raises an exception if it can't find enough nodes for the configured replicas.
     """
+    drain_nodes = settings.drain_nodes
     target_replicas = settings.node_replicas
     available_nodes = get_index_nodes()
+    available_nodes = [node for node in available_nodes if node.id not in drain_nodes]
     if len(available_nodes) < target_replicas:
         raise NodeClusterSmall(
             f"Not enough nodes. Total: {len(available_nodes)}, Required: {target_replicas}"
         )
     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(
@@ -576,26 +691,32 @@ def check_enough_nodes():
             )
 
 
-def sorted_primary_nodes(avoid_nodes: Optional[list[str]] = None) -> list[str]:
+def sorted_primary_nodes(
+    avoid_nodes: Optional[list[str]] = None,
+    ignore_nodes: Optional[list[str]] = None,
+) -> list[str]:
     """
     Returns the list of all primary node ids sorted by decreasing available
     disk space (from more to less available disk reported).
 
-    It will put the node ids in `avoid_nodes` at the tail of the list.
+    Nodes in `avoid_nodes` are placed at the tail of the list.
+    Nodes in `ignore_nodes` are ignored and never returned.
     """
     primary_nodes = get_index_nodes(include_secondary=False)
 
     # Sort by available disk
-    sorted_primary_nodes = sorted(
-        primary_nodes, key=lambda n: n.available_disk, reverse=True
-    )
-    available_node_ids = [node.id for node in sorted_primary_nodes]
+    sorted_nodes = sorted(primary_nodes, key=lambda n: n.available_disk, reverse=True)
+    available_node_ids = [node.id for node in sorted_nodes]
 
     avoid_nodes = avoid_nodes or []
-    # get preferred nodes first
+    ignore_nodes = ignore_nodes or []
+
+    # Get the non-avoided nodes first
     preferred_nodes = [nid for nid in available_node_ids if nid not in avoid_nodes]
-    # now, add to the end of the last nodes
-    preferred_node_order = preferred_nodes + [
-        nid for nid in available_node_ids if nid not in preferred_nodes
-    ]
-    return preferred_node_order
+
+    # 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]
+
+    # Remove ignore_nodes from the list
+    result_nodes = [nid for nid in result_nodes if nid not in ignore_nodes]
+    return result_nodes