Add ingest gRPC methods to create and delete vectorsets (#2202)

* Add new ingest gRPC methods to create and delete vectorsets

* Ingest gRPC uses ORM KnowledgeBox methods

* Add methods in cluster manager

* Partial implementation on KnowledgeBox

* Mark vectorset to delete

* Implement indexes create/delete vectorset

nucliadb/nucliadb/common/cluster/manager.py

@@ -38,6 +38,7 @@
 )
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb_protos import (
+    knowledgebox_pb2,
     nodereader_pb2,
     noderesources_pb2,
     nodewriter_pb2,
@@ -411,6 +412,38 @@ async def maybe_create_new_shard(
             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_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_read_replica_nodes=False
+        )
+
 
 class StandaloneKBShardManager(KBShardManager):
     max_ops_before_checks = 200

nucliadb/nucliadb/ingest/orm/exceptions.py

@@ -48,3 +48,7 @@ class ResourceNotIndexable(Exception):
 
 class EntityManagementException(Exception):
     pass
+
+
+class VectorSetConflict(Exception):
+    pass

nucliadb/nucliadb/ingest/orm/knowledgebox.py

@@ -38,15 +38,15 @@
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.maindb.driver import Driver, Transaction
 from nucliadb.ingest import SERVICE_NAME, logger
-from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict
+from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict, VectorSetConflict
 from nucliadb.ingest.orm.resource import (
     KB_RESOURCE_SLUG,
     KB_RESOURCE_SLUG_BASE,
     Resource,
 )
 from nucliadb.ingest.orm.utils import choose_matryoshka_dimension, compute_paragraph_key
 from nucliadb.migrator.utils import get_latest_version
-from nucliadb_protos import writer_pb2
+from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.utilities import get_audit, get_storage
 
@@ -61,6 +61,9 @@
 KB_TO_DELETE = f"{KB_TO_DELETE_BASE}{{kbid}}"
 KB_TO_DELETE_STORAGE = f"{KB_TO_DELETE_STORAGE_BASE}{{kbid}}"
 
+KB_VECTORSET_TO_DELETE_BASE = "/kbvectorsettodelete"
+KB_VECTORSET_TO_DELETE = f"{KB_VECTORSET_TO_DELETE_BASE}/{{kbid}}"
+
 
 class KnowledgeBox:
     def __init__(self, txn: Transaction, storage: Storage, kbid: str):
@@ -418,6 +421,24 @@ async def iterate_resources(self) -> AsyncGenerator[Resource, None]:
                     disable_vectors=False,
                 )
 
+    async def create_vectorset(self, config: knowledgebox_pb2.VectorSetConfig):
+        if await datamanagers.vectorsets.exists(
+            self.txn, kbid=self.kbid, vectorset_id=config.vectorset_id
+        ):
+            raise VectorSetConflict(f"Vectorset {config.vectorset_id} already exists")
+        await datamanagers.vectorsets.set(self.txn, kbid=self.kbid, config=config)
+        shard_manager = get_shard_manager()
+        await shard_manager.create_vectorset(self.kbid, config)
+
+    async def delete_vectorset(self, vectorset_id: str):
+        await datamanagers.vectorsets.delete(
+            self.txn, kbid=self.kbid, vectorset_id=vectorset_id
+        )
+        # mark vectorset for async deletion
+        await self.txn.set(KB_VECTORSET_TO_DELETE.format(kbid=self.kbid), b"")
+        shard_manager = get_shard_manager()
+        await shard_manager.delete_vectorset(self.kbid, vectorset_id)
+
 
 def chunker(seq: Sequence, size: int):
     return (seq[pos : pos + size] for pos in range(0, len(seq), size))

nucliadb/nucliadb/ingest/service/writer.py

@@ -33,13 +33,16 @@
     NewKnowledgeBoxResponse,
     SemanticModelMetadata,
     UpdateKnowledgeBoxResponse,
+    VectorSetConfig,
 )
 from nucliadb_protos.resources_pb2 import CloudFile
 from nucliadb_protos.writer_pb2 import (
     BinaryData,
     BrokerMessage,
     DelEntitiesRequest,
     DelLabelsRequest,
+    DelVectorSetRequest,
+    DelVectorSetResponse,
     ExtractedVectorsWrapper,
     FileRequest,
     FileUploaded,
@@ -59,6 +62,8 @@
     ListMembersResponse,
     NewEntitiesGroupRequest,
     NewEntitiesGroupResponse,
+    NewVectorSetRequest,
+    NewVectorSetResponse,
     OpStatusWriter,
     SetEntitiesRequest,
     SetLabelsRequest,
@@ -80,12 +85,12 @@
 from nucliadb.common.maindb.utils import setup_driver
 from nucliadb.ingest import SERVICE_NAME, logger
 from nucliadb.ingest.orm.entities import EntitiesManager
-from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict
+from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict, VectorSetConflict
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
 from nucliadb.ingest.orm.processor import Processor, sequence_manager
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
 from nucliadb.ingest.settings import settings
-from nucliadb_protos import utils_pb2, writer_pb2, writer_pb2_grpc
+from nucliadb_protos import nodewriter_pb2, utils_pb2, writer_pb2, writer_pb2_grpc
 from nucliadb_telemetry import errors
 from nucliadb_utils import const
 from nucliadb_utils.settings import is_onprem_nucliadb, running_settings
@@ -687,6 +692,59 @@ async def generate_buffer(
         result = FileUploaded()
         return result
 
+    async def NewVectorSet(
+        self, request: NewVectorSetRequest, context=None
+    ) -> NewVectorSetResponse:
+        config = VectorSetConfig(
+            vectorset_id=request.vectorset_id,
+            vectorset_index_config=nodewriter_pb2.VectorIndexConfig(
+                similarity=request.similarity,
+                normalize_vectors=request.normalize_vectors,
+                vector_type=request.vector_type,
+                vector_dimension=request.vector_dimension,
+            ),
+            matryoshka_dimensions=request.matryoshka_dimensions,
+        )
+        response = NewVectorSetResponse()
+        try:
+            async with self.driver.transaction() as txn:
+                kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
+                await kbobj.create_vectorset(config)
+                await txn.commit()
+        except VectorSetConflict as exc:
+            response.status = NewVectorSetResponse.Status.ERROR
+            response.details = str(exc)
+        except Exception as exc:
+            errors.capture_exception(exc)
+            logger.error(
+                "Error in ingest gRPC while creating a vectorset", exc_info=True
+            )
+            response.status = NewVectorSetResponse.Status.ERROR
+            response.details = str(exc)
+        else:
+            response.status = NewVectorSetResponse.Status.OK
+        return response
+
+    async def DelVectorSet(
+        self, request: DelVectorSetRequest, context=None
+    ) -> DelVectorSetResponse:
+        response = DelVectorSetResponse()
+        try:
+            async with self.driver.transaction() as txn:
+                kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
+                await kbobj.delete_vectorset(request.vectorset_id)
+                await txn.commit()
+        except Exception as exc:
+            errors.capture_exception(exc)
+            logger.error(
+                "Error in ingest gRPC while deleting a vectorset", exc_info=True
+            )
+            response.status = DelVectorSetResponse.Status.ERROR
+            response.details = str(exc)
+        else:
+            response.status = DelVectorSetResponse.Status.OK
+        return response
+
 
 LEARNING_SIMILARITY_FUNCTION_TO_PROTO = {
     "cosine": utils_pb2.VectorSimilarity.COSINE,

nucliadb_protos/python/nucliadb_protos/train_pb2.pyi

@@ -127,6 +127,7 @@ from nucliadb_protos.writer_pb2 import (
     DelEntitiesRequest as DelEntitiesRequest,
     DelLabelsRequest as DelLabelsRequest,
     DelVectorSetRequest as DelVectorSetRequest,
+    DelVectorSetResponse as DelVectorSetResponse,
     Error as Error,
     FileRequest as FileRequest,
     FileUploaded as FileUploaded,
@@ -150,13 +151,14 @@ from nucliadb_protos.writer_pb2 import (
     MergeEntitiesRequest as MergeEntitiesRequest,
     NewEntitiesGroupRequest as NewEntitiesGroupRequest,
     NewEntitiesGroupResponse as NewEntitiesGroupResponse,
+    NewVectorSetRequest as NewVectorSetRequest,
+    NewVectorSetResponse as NewVectorSetResponse,
     Notification as Notification,
     NotificationSource as NotificationSource,
     OpStatusWriter as OpStatusWriter,
     PROCESSOR as PROCESSOR,
     SetEntitiesRequest as SetEntitiesRequest,
     SetLabelsRequest as SetLabelsRequest,
-    SetVectorSetRequest as SetVectorSetRequest,
     SetVectorsRequest as SetVectorsRequest,
     SetVectorsResponse as SetVectorsResponse,
     ShardObject as ShardObject,

nucliadb_protos/python/nucliadb_protos/train_pb2_grpc.pyi

@@ -125,6 +125,7 @@ from nucliadb_protos.writer_pb2 import (
     DelEntitiesRequest as DelEntitiesRequest,
     DelLabelsRequest as DelLabelsRequest,
     DelVectorSetRequest as DelVectorSetRequest,
+    DelVectorSetResponse as DelVectorSetResponse,
     Error as Error,
     FileRequest as FileRequest,
     FileUploaded as FileUploaded,
@@ -148,13 +149,14 @@ from nucliadb_protos.writer_pb2 import (
     MergeEntitiesRequest as MergeEntitiesRequest,
     NewEntitiesGroupRequest as NewEntitiesGroupRequest,
     NewEntitiesGroupResponse as NewEntitiesGroupResponse,
+    NewVectorSetRequest as NewVectorSetRequest,
+    NewVectorSetResponse as NewVectorSetResponse,
     Notification as Notification,
     NotificationSource as NotificationSource,
     OpStatusWriter as OpStatusWriter,
     PROCESSOR as PROCESSOR,
     SetEntitiesRequest as SetEntitiesRequest,
     SetLabelsRequest as SetLabelsRequest,
-    SetVectorSetRequest as SetVectorSetRequest,
     SetVectorsRequest as SetVectorsRequest,
     SetVectorsResponse as SetVectorsResponse,
     ShardObject as ShardObject,