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

migrations/0025_assign_models_to_kbs_v2.py

@@ -0,0 +1,113 @@
+# 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 #25 (Fixed migration 24)
+
+Vectorsets are coming and we need to be ready at nucliadb. Vector index config
+shouldn't be stored anymore in the `Shards` protobuffer, we need to migrate to
+the new vectorsets config.
+
+This migration asks learning_config for each KB configuration and saves the
+model name as the vectorset_id. Creates a vectorset configuration for each model
+and deprecates the vectors index config from the `Shards` protobuffer.
+
+This migration should work for onprem and hosted deployments, as
+learning_proxy handles which API is used (internal or external)
+
+"""
+
+import logging
+
+from nucliadb import learning_proxy
+from nucliadb.common import datamanagers
+from nucliadb.migrator.context import ExecutionContext
+from nucliadb_protos import (
+    knowledgebox_pb2,
+    nodewriter_pb2,
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    async with context.kv_driver.transaction(read_only=True) as txn:
+        vectorsets_count = len([vs async for vs in datamanagers.vectorsets.iter(txn, kbid=kbid)])
+    if vectorsets_count > 0:
+        logger.info("Skipping KB with vectorsets already populated", extra={"kbid": kbid})
+        return
+
+    learning_config = await learning_proxy.get_configuration(kbid)
+    if learning_config is None:
+        logger.warning(f"KB has no learning config", extra={"kbid": kbid})
+        return None
+
+    vectorset_id = learning_config.semantic_model
+    learning_model_metadata = learning_config.into_semantic_model_metadata()
+    learning_similarity = learning_model_metadata.similarity_function
+    learning_vector_dimension = learning_model_metadata.vector_dimension
+    learning_matryoshka_dimensions = learning_model_metadata.matryoshka_dimensions
+    learning_normalize_vectors = len(learning_matryoshka_dimensions) > 0
+
+    async with context.kv_driver.transaction(read_only=True) as txn:
+        semantic_model = await datamanagers.kb.get_model_metadata(txn, kbid=kbid)
+
+        maindb_similarity = semantic_model.similarity_function
+
+        maindb_vector_dimension = None
+        if semantic_model.vector_dimension:
+            maindb_vector_dimension = semantic_model.vector_dimension
+
+        maindb_matryoshka_dimensions: list[int] = []
+        if len(semantic_model.matryoshka_dimensions) > 0:
+            maindb_matryoshka_dimensions.extend(semantic_model.matryoshka_dimensions)
+
+        maindb_normalize_vectors = len(maindb_matryoshka_dimensions) > 0
+
+    if (
+        maindb_similarity != learning_similarity
+        or (maindb_vector_dimension is not None and maindb_vector_dimension != learning_vector_dimension)
+        or set(maindb_matryoshka_dimensions) != set(learning_matryoshka_dimensions)
+        or maindb_normalize_vectors != learning_normalize_vectors
+    ):
+        logger.error(
+            "KB has mismatched data between nucliadb and learning_config! Please, review manually",
+            extra={"kbid": kbid},
+        )
+        return None
+
+    default_vectorset = knowledgebox_pb2.VectorSetConfig(
+        vectorset_id=vectorset_id,
+        vectorset_index_config=nodewriter_pb2.VectorIndexConfig(
+            vector_dimension=maindb_vector_dimension,
+            similarity=maindb_similarity,
+            vector_type=nodewriter_pb2.VectorType.DENSE_F32,  # we only support this for now
+            normalize_vectors=maindb_normalize_vectors,
+        ),
+        matryoshka_dimensions=maindb_matryoshka_dimensions,
+    )
+
+    async with context.kv_driver.transaction() as txn:
+        # Populate KB vectorsets with data from learning. We are skipping KBs
+        # with this key already set, so we can set here safely
+        await datamanagers.vectorsets.set(txn, kbid=kbid, config=default_vectorset)
+        await txn.commit()

migrations/0026_fix_high_cardinality_content_types.py

@@ -0,0 +1,61 @@
+# 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 #26
+
+Previously, there was no validation on content types added by users on upload. This caused that in some KBs,
+there were content types that included random uuids, which caused high cardinality in the content type field.
+
+This migration will fix those invalid content types.
+"""
+
+import logging
+
+from nucliadb.common import datamanagers
+from nucliadb.migrator.context import ExecutionContext
+
+logger = logging.getLogger(__name__)
+
+
+AFFECTED_KBS = [
+    "78d289e0-dd4d-448c-84b5-8ef0b01a5aba",
+]
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    if kbid not in AFFECTED_KBS:
+        return
+    async for rid in datamanagers.resources.iterate_resource_ids(kbid=kbid):
+        async with datamanagers.with_rw_transaction() as txn:
+            basic = await datamanagers.resources.get_basic(txn, kbid=kbid, rid=rid)
+            if not basic or not basic.icon:
+                continue
+            # We're aiming to fix content types like "multipart/form-data; boundary={uuid}"
+            if "multipart/form-data" not in basic.icon:
+                continue
+            if "boundary=" not in basic.icon:
+                continue
+            logger.info("Fixing content type for resource", extra={"kbid": kbid, "rid": rid})
+            basic.icon = "multipart/form-data"
+            await datamanagers.resources.set_basic(txn, kbid=kbid, rid=rid, basic=basic)
+            await txn.commit()

migrations/0027_rollover_texts3.py

@@ -0,0 +1,73 @@
+# 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 #27
+
+Rollover for nucliadb_texts3
+"""
+
+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__)
+
+
+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)
+
+
+async def maybe_fix_vector_dimensions(context: ExecutionContext, kbid: str) -> None:
+    learning_config = await learning_proxy.get_configuration(kbid)
+    if learning_config is None:
+        logger.warning(f"KB has no learning config", extra={"kbid": kbid})
+        return
+
+    async with context.kv_driver.transaction() as txn:
+        vectorsets = [vs async for vs in datamanagers.vectorsets.iter(txn, kbid=kbid)]
+        if len(vectorsets) != 1:
+            # If multiple vectorsets, they are new shards created correctly, we can safely skip it
+            logger.warning(f"KB has {len(vectorsets)} vectorsets, skipping...", extra={"kbid": kbid})
+            return
+        vectorset = vectorsets[0][1]
+
+        # Correct value, skip
+        if vectorset.vectorset_index_config.vector_dimension != 0:
+            return
+
+        learning_model_metadata = learning_config.into_semantic_model_metadata()
+        logger.info(
+            f"Fixing KB vectorset dimension",
+            extra={
+                "kbid": kbid,
+                "from": vectorset.vectorset_index_config.vector_dimension,
+                "to": learning_model_metadata.vector_dimension,
+            },
+        )
+        vectorset.vectorset_index_config.vector_dimension = learning_model_metadata.vector_dimension
+
+        await datamanagers.vectorsets.set(txn, kbid=kbid, config=vectorset)

nucliadb/ingest/fields/date.py → migrations/pg/0001_bootstrap.py

@@ -17,18 +17,16 @@
 # 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/>.
 #
-from nucliadb_protos.resources_pb2 import FieldDatetime
 
-from nucliadb.ingest.fields.base import Field
+from nucliadb.common.maindb.pg import PGTransaction
 
 
-class Datetime(Field):
-    pbklass = FieldDatetime
-    value: FieldDatetime
-    type: str = "d"
-
-    async def set_value(self, payload: FieldDatetime):
-        await self.db_set_value(payload)
-
-    async def get_value(self) -> FieldDatetime:
-        return await self.db_get_value()
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        # IF NOT EXISTS just for compatibility with older install predating the migration system
+        await cur.execute("""
+            CREATE TABLE IF NOT EXISTS resources (
+                key TEXT PRIMARY KEY,
+                value BYTEA
+            );
+        """)

migrations/pg/0002_catalog.py

@@ -0,0 +1,42 @@
+# 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/>.
+#
+
+from nucliadb.common.maindb.pg import PGTransaction
+
+
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        await cur.execute(r"""
+            CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
+            CREATE EXTENSION IF NOT EXISTS btree_gin;
+            CREATE TABLE catalog (
+                kbid UUID,
+                rid UUID,
+                title TEXT,
+                created_at TIMESTAMP,
+                modified_at TIMESTAMP,
+                labels TEXT[],
+                PRIMARY KEY(kbid, rid)
+            );
+            CREATE INDEX ON catalog USING GIN(kbid, labels);
+            CREATE INDEX ON catalog USING GIN(kbid, regexp_split_to_array(lower(title), '\W'::text));
+            CREATE INDEX ON catalog(kbid, created_at);
+            CREATE INDEX ON catalog(kbid, modified_at);
+        """)

nucliadb/ingest/tests/unit/test_settings.py → migrations/pg/0003_catalog_kbid_index.py

@@ -17,8 +17,10 @@
 # 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/>.
 #
-from nucliadb.ingest.settings import DriverConfig
 
+from nucliadb.common.maindb.pg import PGTransaction
 
-def test_case_insenstive_driver_config():
-    assert DriverConfig("PG") == DriverConfig.PG
+
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        await cur.execute("CREATE INDEX ON catalog(kbid);")

nucliadb/common/cluster/base.py

@@ -20,16 +20,16 @@
 from abc import ABCMeta, abstractmethod
 from typing import AsyncIterator, Optional
 
+from nucliadb_protos import nodereader_pb2, noderesources_pb2, utils_pb2
 from nucliadb_protos.nodereader_pb2_grpc import NodeReaderStub
 from nucliadb_protos.nodewriter_pb2 import (
     NewShardRequest,
     NewVectorSetRequest,
     OpStatus,
+    VectorIndexConfig,
 )
 from nucliadb_protos.nodewriter_pb2_grpc import NodeWriterStub
 
-from nucliadb_protos import nodereader_pb2, noderesources_pb2, utils_pb2
-
 
 class AbstractIndexNode(metaclass=ABCMeta):
     label: str = "index-node"
@@ -93,15 +93,29 @@ class AbstractIndexNode(metaclass=ABCMeta):
     async def new_shard(
         self,
         kbid: str,
-        similarity: utils_pb2.VectorSimilarity.ValueType,
-        release_channel: utils_pb2.ReleaseChannel.ValueType,
-        normalize_vectors: bool,
+        vector_index_config: VectorIndexConfig,
+    ) -> noderesources_pb2.ShardCreated:
+        req = NewShardRequest(
+            kbid=kbid,
+            release_channel=utils_pb2.ReleaseChannel.STABLE,
+            config=vector_index_config,
+            # Deprecated fields, only for backwards compatibility with older nodes
+            similarity=vector_index_config.similarity,
+            normalize_vectors=vector_index_config.normalize_vectors,
+        )
+
+        resp = await self.writer.NewShard(req)  # type: ignore
+        return resp
+
+    async def new_shard_with_vectorsets(
+        self,
+        kbid: str,
+        vectorsets_configs: dict[str, VectorIndexConfig],
     ) -> noderesources_pb2.ShardCreated:
         req = NewShardRequest(
             kbid=kbid,
-            similarity=similarity,
-            release_channel=release_channel,
-            normalize_vectors=normalize_vectors,
+            release_channel=utils_pb2.ReleaseChannel.STABLE,
+            vectorsets_configs=vectorsets_configs,
         )
 
         resp = await self.writer.NewShard(req)  # type: ignore
@@ -120,15 +134,15 @@ class AbstractIndexNode(metaclass=ABCMeta):
         self,
         shard_id: str,
         vectorset: str,
-        *,
-        similarity: utils_pb2.VectorSimilarity.ValueType = utils_pb2.VectorSimilarity.COSINE,
-        normalize_vectors: bool = False,
+        config: VectorIndexConfig,
     ) -> OpStatus:
-        req = NewVectorSetRequest()
-        req.id.shard.id = shard_id
-        req.id.vectorset = vectorset
-        req.similarity = similarity
-        req.normalize_vectors = normalize_vectors
+        req = NewVectorSetRequest(
+            id=noderesources_pb2.VectorSetID(
+                shard=noderesources_pb2.ShardId(id=shard_id), vectorset=vectorset
+            ),
+            config=config,
+        )
+
         resp = await self.writer.AddVectorSet(req)  # type: ignore
         return resp
 

nucliadb/common/cluster/discovery/base.py

@@ -113,7 +113,7 @@ async def _get_index_node_metadata(
     channel = get_traced_grpc_channel(grpc_address, "discovery", variant="_writer")
     if read_replica:
         # on a read replica, we need to use the replication service
-        stub = replication_pb2_grpc.ReplicationServiceStub(channel)  # type: ignore
+        stub = replication_pb2_grpc.ReplicationServiceStub(channel)
     else:
         stub = nodewriter_pb2_grpc.NodeWriterStub(channel)  # type: ignore
     try:
@@ -127,9 +127,7 @@ async def _get_index_node_metadata(
         or None
     )
     if read_replica and primary_id is None:
-        raise Exception(
-            "Primary node id not found when it is expected to be a read replica"
-        )
+        raise Exception("Primary node id not found when it is expected to be a read replica")
 
     return IndexNodeMetadata(
         node_id=metadata.node_id,
@@ -141,18 +139,14 @@ async def _get_index_node_metadata(
     )
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=4
-)
-async def _get_standalone_index_node_metadata(
-    settings: Settings, address: str
-) -> IndexNodeMetadata:
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=4)
+async def _get_standalone_index_node_metadata(settings: Settings, address: str) -> IndexNodeMetadata:
     if ":" not in address:
         grpc_address = f"{address}:{settings.standalone_node_port}"
     else:
         grpc_address = address
     channel = get_traced_grpc_channel(grpc_address, "standalone_proxy")
-    stub = standalone_pb2_grpc.StandaloneClusterServiceStub(channel)  # type: ignore
+    stub = standalone_pb2_grpc.StandaloneClusterServiceStub(channel)
     resp: standalone_pb2.NodeInfoResponse = await stub.NodeInfo(standalone_pb2.NodeInfoRequest())  # type: ignore
     return IndexNodeMetadata(
         node_id=resp.id,
@@ -177,9 +171,7 @@ class AbstractClusterDiscovery(abc.ABC):
     async def finalize(self) -> None:
         """ """
 
-    async def _query_node_metadata(
-        self, address: str, read_replica: bool = False
-    ) -> IndexNodeMetadata:
+    async def _query_node_metadata(self, address: str, read_replica: bool = False) -> IndexNodeMetadata:
         if self.settings.standalone_mode:
             return await _get_standalone_index_node_metadata(self.settings, address)
         else:

nucliadb/common/cluster/discovery/k8s.py

@@ -69,9 +69,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
     ) -> IndexNodeMetadata:
         async with self.update_lock:
             if pod_name not in self.node_id_cache:
-                self.node_id_cache[pod_name] = await self._query_node_metadata(
-                    node_ip, read_replica
-                )
+                self.node_id_cache[pod_name] = await self._query_node_metadata(node_ip, read_replica)
             else:
                 self.node_id_cache[pod_name].address = node_ip
             self.node_id_cache[pod_name].updated_at = time.time()
@@ -84,12 +82,10 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
         This method will update global node state by utilizing the cluster manager
         to add or remove nodes.
         """
-        status: kubernetes_asyncio.client.models.v1_pod_status.V1PodStatus = event[
+        status: kubernetes_asyncio.client.models.v1_pod_status.V1PodStatus = event["object"].status
+        event_metadata: kubernetes_asyncio.client.models.v1_object_meta.V1ObjectMeta = event[
             "object"
-        ].status
-        event_metadata: kubernetes_asyncio.client.models.v1_object_meta.V1ObjectMeta = (
-            event["object"].metadata
-        )
+        ].metadata
 
         ready = status.container_statuses is not None
         if event["type"] == "DELETED":
@@ -199,9 +195,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
                             except NodeConnectionError:  # pragma: no cover
                                 pass
                             except Exception:  # pragma: no cover
-                                logger.exception(
-                                    "Error while updating node", exc_info=True
-                                )
+                                logger.exception("Error while updating node", exc_info=True)
                     except (
                         asyncio.CancelledError,
                         KeyboardInterrupt,
@@ -259,11 +253,9 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
                             continue
                         existing = self.node_id_cache[pod_name]
                         try:
-                            self.node_id_cache[pod_name] = (
-                                await self._query_node_metadata(
-                                    existing.address,
-                                    read_replica=existing.primary_id is not None,
-                                )
+                            self.node_id_cache[pod_name] = await self._query_node_metadata(
+                                existing.address,
+                                read_replica=existing.primary_id is not None,
                             )
                         except NodeConnectionError:  # pragma: no cover
                             self._maybe_remove_stale_node(pod_name)
@@ -301,9 +293,7 @@ class KubernetesDiscovery(AbstractClusterDiscovery):
 
     async def initialize(self) -> None:
         self.cluster_task = asyncio.create_task(self.watch_k8s_for_updates())
-        self.update_node_data_cache_task = asyncio.create_task(
-            self.update_node_data_cache()
-        )
+        self.update_node_data_cache_task = asyncio.create_task(self.update_node_data_cache())
         await self._wait_ready()
 
     async def finalize(self) -> None:

nucliadb/common/cluster/discovery/manual.py

@@ -46,9 +46,7 @@ class ManualDiscovery(AbstractClusterDiscovery):
             except asyncio.CancelledError:
                 return
             except Exception:
-                logger.exception(
-                    "Error while watching cluster members. Will retry at started interval"
-                )
+                logger.exception("Error while watching cluster members. Will retry at started interval")
             finally:
                 await asyncio.sleep(15)
 

nucliadb/common/cluster/discovery/utils.py

@@ -40,9 +40,7 @@ async def setup_cluster_discovery() -> None:
             # already loaded
             return util
 
-        klass: Union[
-            Type[ManualDiscovery], Type[KubernetesDiscovery], Type[SingleNodeDiscovery]
-        ]
+        klass: Union[Type[ManualDiscovery], Type[KubernetesDiscovery], Type[SingleNodeDiscovery]]
         if settings.cluster_discovery_mode == ClusterDiscoveryMode.MANUAL:
             klass = ManualDiscovery
         elif settings.cluster_discovery_mode == ClusterDiscoveryMode.KUBERNETES:

nucliadb/common/cluster/grpc_node_dummy.py

@@ -22,16 +22,15 @@ from typing import Any
 from nucliadb_protos.nodereader_pb2 import (
     EdgeList,
     RelationEdge,
-    RelationSearchResponse,
 )
-from nucliadb_protos.noderesources_pb2 import EmptyResponse
-from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
 from nucliadb_protos.noderesources_pb2 import (
+    EmptyResponse,
     ShardCreated,
     ShardId,
     ShardIds,
     VectorSetList,
 )
+from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
 from nucliadb_protos.nodewriter_pb2 import OpStatus
 from nucliadb_protos.utils_pb2 import Relation
 
@@ -90,15 +89,8 @@ class DummyReaderStub:  # pragma: no cover
         self.calls.setdefault("GetShard", []).append(data)
         return NodeResourcesShard(shard_id="shard", fields=2, paragraphs=2, sentences=2)
 
-    async def RelationSearch(self, data):  # pragma: no cover
-        self.calls.setdefault("RelationSearch", []).append(data)
-        result = RelationSearchResponse()
-        return result
-
     async def RelationEdges(self, data):  # pragma: no cover
         self.calls.setdefault("RelationEdges", []).append(data)
         result = EdgeList()
-        result.list.append(
-            RelationEdge(edge_type=Relation.RelationType.ENTITY, property="dummy")
-        )
+        result.list.append(RelationEdge(edge_type=Relation.RelationType.ENTITY, property="dummy"))
         return result

nucliadb/common/cluster/index_node.py

@@ -19,16 +19,13 @@
 #
 from typing import Optional
 
-from lru import LRU  # type: ignore
-from nucliadb_protos.nodereader_pb2_grpc import NodeReaderStub
-from nucliadb_protos.nodewriter_pb2_grpc import NodeWriterStub
+from lru import LRU
 
-from nucliadb.common.cluster.base import AbstractIndexNode  # type: ignore
-from nucliadb.common.cluster.grpc_node_dummy import (  # type: ignore
-    DummyReaderStub,
-    DummyWriterStub,
-)
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.cluster.grpc_node_dummy import DummyReaderStub, DummyWriterStub
 from nucliadb.ingest import SERVICE_NAME
+from nucliadb_protos.nodereader_pb2_grpc import NodeReaderStub
+from nucliadb_protos.nodewriter_pb2_grpc import NodeWriterStub
 from nucliadb_utils.grpc import get_traced_grpc_channel
 
 from .settings import settings
@@ -41,9 +38,7 @@ class IndexNode(AbstractIndexNode):
     _writer: Optional[NodeWriterStub] = None
     _reader: Optional[NodeReaderStub] = None
 
-    def _get_service_address(
-        self, port_map: dict[str, int], port: Optional[int]
-    ) -> str:
+    def _get_service_address(self, port_map: dict[str, int], port: Optional[int]) -> str:
         hostname = self.address.split(":")[0]
         if port is None:
             # For testing purposes we need to be able to have a writing port
@@ -60,10 +55,8 @@ class IndexNode(AbstractIndexNode):
                 grpc_address = self._get_service_address(
                     settings.writer_port_map, settings.node_writer_port
                 )
-                channel = get_traced_grpc_channel(
-                    grpc_address, SERVICE_NAME, variant="_writer"
-                )
-                WRITE_CONNECTIONS[self.address] = NodeWriterStub(channel)  # type: ignore
+                channel = get_traced_grpc_channel(grpc_address, SERVICE_NAME, variant="_writer")
+                WRITE_CONNECTIONS[self.address] = NodeWriterStub(channel)
             else:
                 WRITE_CONNECTIONS[self.address] = DummyWriterStub()
             self._writer = WRITE_CONNECTIONS[self.address]
@@ -76,10 +69,8 @@ class IndexNode(AbstractIndexNode):
                 grpc_address = self._get_service_address(
                     settings.reader_port_map, settings.node_reader_port
                 )
-                channel = get_traced_grpc_channel(
-                    grpc_address, SERVICE_NAME, variant="_reader"
-                )
-                READ_CONNECTIONS[self.address] = NodeReaderStub(channel)  # type: ignore
+                channel = get_traced_grpc_channel(grpc_address, SERVICE_NAME, variant="_reader")
+                READ_CONNECTIONS[self.address] = NodeReaderStub(channel)
             else:
                 READ_CONNECTIONS[self.address] = DummyReaderStub()
             self._reader = READ_CONNECTIONS[self.address]