nucliadb 6.2.1.post2971__py3-none-any.whl → 6.2.1.post2977__py3-none-any.whl

nucliadb/common/cluster/standalone/utils.py

@@ -19,13 +19,10 @@
 
 import logging
 import os
-import shutil
 import uuid
-from socket import gethostname
 
 from nucliadb.common.cluster.settings import StandaloneNodeRole
 from nucliadb.common.cluster.settings import settings as cluster_settings
-from nucliadb.common.cluster.standalone.index_node import StandaloneIndexNode
 
 logger = logging.getLogger(__name__)
 
@@ -46,46 +43,6 @@ def get_standalone_node_id() -> str:
         return str(uuid.UUID(bytes=f.read()))
 
 
-_SELF_INDEX_NODE = None
-
-
-def get_self() -> StandaloneIndexNode:
-    """
-    This returns an instance of the standalone index node
-    so when API requests come into this mode, we don't
-    make another grpc request since this node can service it directly.
-    """
-    if not is_index_node():
-        raise Exception("This node is not an Index Node. You should not reach this code path.")
-    global _SELF_INDEX_NODE
-    node_id = get_standalone_node_id()
-    if _SELF_INDEX_NODE is None or node_id != _SELF_INDEX_NODE.id:
-        if "NUCLIADB_SERVICE_HOST" in os.environ:
-            hn = os.environ["HOSTNAME"]
-            ns = os.environ.get("NAMESPACE", "nucliadb")
-            host = f"{hn}.{ns}"
-        else:
-            host = gethostname()
-        _SELF_INDEX_NODE = StandaloneIndexNode(id=node_id, address=host, shard_count=0, available_disk=0)
-    try:
-        _, _, available_disk = shutil.disk_usage(cluster_settings.data_path)
-        _SELF_INDEX_NODE.available_disk = available_disk
-    except FileNotFoundError:  # pragma: no cover
-        ...
-    try:
-        _shards_dir = os.path.join(cluster_settings.data_path, "shards")
-        _SELF_INDEX_NODE.shard_count = len(
-            [
-                shard_dir
-                for shard_dir in os.listdir(_shards_dir)
-                if os.path.isdir(os.path.join(_shards_dir, shard_dir))
-            ]
-        )
-    except FileNotFoundError:  # pragma: no cover
-        ...
-    return _SELF_INDEX_NODE
-
-
 def is_index_node() -> bool:
     return cluster_settings.standalone_node_role in (
         StandaloneNodeRole.ALL,

nucliadb/common/cluster/utils.py

@@ -23,20 +23,11 @@ from typing import TYPE_CHECKING, Optional, Union
 import backoff
 
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.discovery.utils import (
-    setup_cluster_discovery,
-    teardown_cluster_discovery,
-)
 from nucliadb.common.cluster.manager import (
     KBShardManager,
     StandaloneKBShardManager,
-    clear_index_nodes,
 )
 from nucliadb.common.cluster.settings import settings
-from nucliadb.common.cluster.standalone.service import (
-    start_grpc as start_standalone_grpc,
-)
-from nucliadb.common.cluster.standalone.utils import is_index_node
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import nodereader_pb2, writer_pb2
 from nucliadb_utils import const
@@ -62,12 +53,8 @@ async def setup_cluster() -> Union[KBShardManager, StandaloneKBShardManager]:
             # already setup
             return get_utility(Utility.SHARD_MANAGER)
 
-        await setup_cluster_discovery()
         mng: Union[KBShardManager, StandaloneKBShardManager]
         if settings.standalone_mode:
-            if is_index_node():
-                server = await start_standalone_grpc()
-                set_utility(_STANDALONE_SERVER, server)
             mng = StandaloneKBShardManager()
         else:
             mng = KBShardManager()
@@ -76,7 +63,6 @@ async def setup_cluster() -> Union[KBShardManager, StandaloneKBShardManager]:
 
 
 async def teardown_cluster():
-    await teardown_cluster_discovery()
     if get_utility(Utility.SHARD_MANAGER):
         clean_utility(Utility.SHARD_MANAGER)
 
@@ -85,8 +71,6 @@ async def teardown_cluster():
         await std_server.stop(None)
         clean_utility(_STANDALONE_SERVER)
 
-    clear_index_nodes()
-
 
 def get_shard_manager() -> KBShardManager:
     return get_utility(Utility.SHARD_MANAGER)  # type: ignore

nucliadb/common/nidx.py

@@ -21,7 +21,7 @@
 import os
 from typing import Optional
 
-from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxSearcherStub
+from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxIndexerStub, NidxSearcherStub
 
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.settings import settings
@@ -37,12 +37,10 @@ from nucliadb_utils.settings import FileBackendConfig, indexing_settings, storag
 from nucliadb_utils.storages.settings import settings as extended_storage_settings
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
-NIDX_ENABLED = bool(os.environ.get("NIDX_ENABLED"))
-
 
 class NidxUtility:
-    api_client = None
-    searcher_client = None
+    api_client: NidxApiStub
+    searcher_client: NidxSearcherStub
 
     async def initialize(self):
         raise NotImplementedError()
@@ -98,6 +96,9 @@ class NidxBindingUtility(NidxUtility):
 
         self.config = {
             "METADATA__DATABASE_URL": ingest_settings.driver_pg_url,
+            "SEARCHER__METADATA_REFRESH_INTERVAL": str(
+                indexing_settings.index_searcher_refresh_interval
+            ),
             **_storage_config("INDEXER", None),
             **_storage_config("STORAGE", "nidx"),
         }
@@ -123,16 +124,8 @@ class NidxBindingUtility(NidxUtility):
         self.binding.wait_for_sync()
 
 
-class NidxServiceUtility(NidxUtility):
-    """Implements Nidx utility connecting to the network service"""
-
+class NidxNatsIndexer:
     def __init__(self):
-        if indexing_settings.index_nidx_subject is None:
-            raise ValueError("INDEX_NIDX_SUBJECT needed for nidx utility")
-
-        if not settings.nidx_api_address or not settings.nidx_searcher_address:
-            raise ValueError("NIDX_API_ADDRESS and NIDX_SEARCHER_ADDRESS are required")
-
         self.nats_connection_manager = NatsConnectionManager(
             service_name="NidxIndexer",
             nats_servers=indexing_settings.index_jetstream_servers,
@@ -142,10 +135,6 @@ class NidxServiceUtility(NidxUtility):
 
     async def initialize(self):
         await self.nats_connection_manager.initialize()
-        self.api_client = NidxApiStub(get_traced_grpc_channel(settings.nidx_api_address, "nidx_api"))
-        self.searcher_client = NidxSearcherStub(
-            get_traced_grpc_channel(settings.nidx_searcher_address, "nidx_searcher")
-        )
 
     async def finalize(self):
         await self.nats_connection_manager.finalize()
@@ -158,18 +147,68 @@ class NidxServiceUtility(NidxUtility):
         return res.seq
 
 
-async def start_nidx_utility() -> Optional[NidxUtility]:
-    if not NIDX_ENABLED:
-        return None
+class NidxGrpcIndexer:
+    def __init__(self, address):
+        self.address = address
 
-    nidx = get_nidx()
+    async def initialize(self):
+        self.client = NidxIndexerStub(get_traced_grpc_channel(self.address, "nidx_indexer"))
+
+    async def finalize(self):
+        pass
+
+    async def index(self, writer: IndexMessage) -> int:
+        await self.client.Index(writer)
+        return 0
+
+
+class NidxServiceUtility(NidxUtility):
+    """Implements Nidx utility connecting to the network service"""
+
+    def __init__(self):
+        if not settings.nidx_api_address or not settings.nidx_searcher_address:
+            raise ValueError("NIDX_API_ADDRESS and NIDX_SEARCHER_ADDRESS are required")
+
+        if indexing_settings.index_nidx_subject:
+            self.indexer = NidxNatsIndexer()
+        elif settings.nidx_indexer_address is not None:
+            self.indexer = NidxGrpcIndexer(settings.nidx_indexer_address)
+        else:
+            raise ValueError("NIDX_INDEXER_ADDRESS or INDEX_NIDX_SUBJECT are required")
+
+    async def initialize(self):
+        await self.indexer.initialize()
+        self.api_client = NidxApiStub(get_traced_grpc_channel(settings.nidx_api_address, "nidx_api"))
+        self.searcher_client = NidxSearcherStub(
+            get_traced_grpc_channel(settings.nidx_searcher_address, "nidx_searcher")
+        )
+
+    async def finalize(self):
+        await self.indexer.finalize()
+
+    async def index(self, writer: IndexMessage) -> int:
+        return await self.indexer.index(writer)
+
+
+async def start_nidx_utility() -> Optional[NidxUtility]:
+    nidx = get_utility(Utility.NIDX)
     if nidx:
         return nidx
 
     nidx_utility: NidxUtility
     if settings.standalone_mode:
-        nidx_utility = NidxBindingUtility()
+        if (
+            settings.nidx_api_address is not None
+            and settings.nidx_searcher_address is not None
+            and settings.nidx_indexer_address is not None
+        ):
+            # Standalone with nidx service (via grpc). This is used in clustered standalone mode
+            nidx_utility = NidxServiceUtility()
+        else:
+            # Normal standalone mode with binding
+            nidx_utility = NidxBindingUtility()
     else:
+        # Component deploy with nidx service via grpc & nats (cloud)
         nidx_utility = NidxServiceUtility()
 
     await nidx_utility.initialize()
@@ -178,30 +217,33 @@ async def start_nidx_utility() -> Optional[NidxUtility]:
 
 
 async def stop_nidx_utility():
-    nidx_utility = get_nidx()
+    nidx_utility = get_utility(Utility.NIDX)
     if nidx_utility:
         clean_utility(Utility.NIDX)
         await nidx_utility.finalize()
 
 
-def get_nidx() -> Optional[NidxUtility]:
-    return get_utility(Utility.NIDX)
+def get_nidx() -> NidxUtility:
+    nidx = get_utility(Utility.NIDX)
+    if nidx is None:
+        raise Exception("nidx not initialized")
+    return nidx
 
 
-def get_nidx_api_client() -> Optional["NidxApiStub"]:
+def get_nidx_api_client() -> "NidxApiStub":
     nidx = get_nidx()
-    if nidx:
+    if nidx.api_client:
         return nidx.api_client
     else:
-        return None
+        raise Exception("nidx not initialized")
 
 
-def get_nidx_searcher_client() -> Optional["NidxSearcherStub"]:
+def get_nidx_searcher_client() -> "NidxSearcherStub":
     nidx = get_nidx()
-    if nidx:
+    if nidx.searcher_client:
         return nidx.searcher_client
     else:
-        return None
+        raise Exception("nidx not initialized")
 
 
 # TODO: Remove the index node abstraction
@@ -252,9 +294,6 @@ class FakeNode(AbstractIndexNode):
         return "nidx"
 
 
-def get_nidx_fake_node() -> Optional[FakeNode]:
+def get_nidx_fake_node() -> FakeNode:
     nidx = get_nidx()
-    if nidx:
-        return FakeNode(nidx.api_client, nidx.searcher_client)
-    else:
-        return None
+    return FakeNode(nidx.api_client, nidx.searcher_client)

nucliadb/health.py

@@ -40,13 +40,6 @@ def nats_manager_healthy() -> bool:
     return nats_manager.healthy()
 
 
-def nodes_health_check() -> bool:
-    from nucliadb.common.cluster import manager
-    from nucliadb.ingest.settings import DriverConfig, settings
-
-    return len(manager.INDEX_NODES) > 0 or settings.driver == DriverConfig.LOCAL
-
-
 def pubsub_check() -> bool:
     driver: Optional[PubSubDriver] = get_utility(Utility.PUBSUB)
     if driver is None:

nucliadb/ingest/app.py

@@ -22,10 +22,6 @@ import importlib.metadata
 from typing import Awaitable, Callable
 
 from nucliadb import health
-from nucliadb.common.cluster.discovery.utils import (
-    setup_cluster_discovery,
-    teardown_cluster_discovery,
-)
 from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.context import ApplicationContext
@@ -89,13 +85,9 @@ async def initialize() -> list[Callable[[], Awaitable[None]]]:
         )
         finalizers.append(stop_nats_manager)
 
-        await setup_cluster_discovery()
-        finalizers.append(teardown_cluster_discovery)
-
     health.register_health_checks(
         [
             health.nats_manager_healthy,
-            health.nodes_health_check,
             health.pubsub_check,
         ]
     )

nucliadb/ingest/consumer/auditing.py

@@ -113,7 +113,7 @@ class IndexAuditHandler:
 
         for shard_obj in shard_groups:
             # TODO: Uses node for auditing, don't want to suddenly change metrics
-            node, shard_id = choose_node(shard_obj, use_nidx=False)
+            node, shard_id = choose_node(shard_obj)
             shard: nodereader_pb2.Shard = await node.reader.GetShard(
                 nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
             )

nucliadb/ingest/consumer/shard_creator.py

@@ -103,7 +103,7 @@ class ShardCreatorHandler:
         async with locking.distributed_lock(locking.NEW_SHARD_LOCK.format(kbid=kbid)):
             # remember, a lock will do at least 1+ reads and 1 write.
             # with heavy writes, this adds some simple k/v pressure
-            node, shard_id = choose_node(current_shard, use_nidx=True)
+            node, shard_id = choose_node(current_shard)
             shard: nodereader_pb2.Shard = await node.reader.GetShard(
                 nodereader_pb2.GetShardRequest(shard_id=noderesources_pb2.ShardId(id=shard_id))  # type: ignore
             )

nucliadb/ingest/orm/entities.py

@@ -37,6 +37,7 @@ from nucliadb.common.datamanagers.entities import (
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.ingest.settings import settings
+from nucliadb.search.search.shards import query_shard
 from nucliadb_protos.knowledgebox_pb2 import (
     DeletedEntitiesGroups,
     EntitiesGroup,
@@ -54,8 +55,6 @@ from nucliadb_protos.nodereader_pb2 import (
 from nucliadb_protos.utils_pb2 import RelationNode
 from nucliadb_protos.writer_pb2 import GetEntitiesResponse
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
-from nucliadb_utils.utilities import has_feature
 
 from .exceptions import EntityManagementException
 
@@ -218,14 +217,13 @@ class EntitiesManager:
                     ],
                 ),
             )
-            response = await node.reader.Search(request)  # type: ignore
+            response = await query_shard(node, shard_id, request)
             return response.relation
 
         results = await shard_manager.apply_for_all_shards(
             self.kbid,
             do_entities_search,
             settings.relation_search_timeout,
-            use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": self.kbid}),
             use_read_replica_nodes=self.use_read_replica_nodes,
         )
         for result in results:
@@ -315,7 +313,7 @@ class EntitiesManager:
                 paragraph=False,
                 faceted=Faceted(labels=["/e"]),
             )
-            response: SearchResponse = await node.reader.Search(request)  # type: ignore
+            response: SearchResponse = await query_shard(node, shard_id, request)
             try:
                 facetresults = response.document.facets["/e"].facetresults
                 return {facet.tag.split("/")[-1] for facet in facetresults}
@@ -327,7 +325,6 @@ class EntitiesManager:
             self.kbid,
             query_indexed_entities_group_names,
             settings.relation_types_timeout,
-            use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": self.kbid}),
             use_read_replica_nodes=self.use_read_replica_nodes,
         )
         for result in results:

nucliadb/purge/orphan_shards.py

@@ -33,6 +33,7 @@ from nucliadb.common.cluster.manager import KBShardManager
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
+from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb.ingest import logger
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.logs import setup_logging
@@ -135,10 +136,9 @@ async def _get_stored_shards(driver: Driver) -> dict[str, ShardLocation]:
                 continue
             else:
                 for shard_object_pb in kb_shards:
-                    for shard_replica_pb in shard_object_pb.replicas:
-                        shard_replica_id = shard_replica_pb.shard.id
-                        node_id = shard_replica_pb.node
-                        stored_shards[shard_replica_id] = ShardLocation(kbid=kbid, node_id=node_id)
+                    stored_shards[shard_object_pb.nidx_shard_id] = ShardLocation(
+                        kbid=kbid, node_id="nidx"
+                    )
     return stored_shards
 
 
@@ -241,6 +241,7 @@ async def main():
     """
     args = parse_arguments()
 
+    await start_nidx_utility()
     await setup_cluster()
     driver = await setup_driver()
 
@@ -253,6 +254,7 @@ async def main():
     finally:
         await teardown_driver()
         await teardown_cluster()
+        await stop_nidx_utility()
 
 
 def run() -> int:  # pragma: no cover

nucliadb/search/api/v1/knowledgebox.py

@@ -48,9 +48,7 @@ from nucliadb_protos.noderesources_pb2 import Shard
 from nucliadb_protos.writer_pb2 import ShardObject as PBShardObject
 from nucliadb_protos.writer_pb2 import Shards
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
 from nucliadb_utils.authentication import requires, requires_one
-from nucliadb_utils.utilities import has_feature
 
 MAX_PARAGRAPHS_FOR_SMALL_KB = 250_000
 
@@ -166,9 +164,7 @@ async def get_node_index_counts(kbid: str) -> tuple[IndexCounts, list[str]]:
     queried_shards = []
     for shard_object in shard_groups:
         try:
-            node, shard_id = choose_node(
-                shard_object, use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": kbid})
-            )
+            node, shard_id = choose_node(shard_object)
         except KeyError:
             raise HTTPException(
                 status_code=500,

nucliadb/search/requesters/utils.py

@@ -123,7 +123,6 @@ async def node_query(
         try:
             node, shard_id = cluster_manager.choose_node(
                 shard_obj,
-                use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": kbid}),
                 use_read_replica_nodes=use_read_replica_nodes,
                 target_shard_replicas=target_shard_replicas,
             )
@@ -224,7 +223,7 @@ def validate_node_query_results(results: list[Any]) -> Optional[HTTPException]:
                     )
             else:
                 errors.capture_exception(result)
-                logger.exception("Error while querying shard data", exc_info=result)
+                logger.exception(f"Error while querying shard data {result}", exc_info=result)
 
             return HTTPException(status_code=status_code, detail=reason)
 

nucliadb/search/search/shards.py

@@ -19,6 +19,10 @@
 #
 import asyncio
 
+import backoff
+from grpc import StatusCode
+from grpc.aio import AioRpcError
+
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb_protos.nodereader_pb2 import (
     GetShardRequest,
@@ -39,6 +43,15 @@ node_observer = metrics.Observer(
 )
 
 
+def should_giveup(e: Exception):
+    if isinstance(e, AioRpcError) and e.code() != StatusCode.NOT_FOUND:
+        return True
+    return False
+
+
+@backoff.on_exception(
+    backoff.expo, Exception, jitter=None, factor=0.1, max_tries=3, giveup=should_giveup
+)
 async def query_shard(node: AbstractIndexNode, shard: str, query: SearchRequest) -> SearchResponse:
     req = SearchRequest()
     req.CopyFrom(query)
@@ -47,6 +60,9 @@ async def query_shard(node: AbstractIndexNode, shard: str, query: SearchRequest)
         return await node.reader.Search(req)  # type: ignore
 
 
+@backoff.on_exception(
+    backoff.expo, Exception, jitter=None, factor=0.1, max_tries=3, giveup=should_giveup
+)
 async def get_shard(node: AbstractIndexNode, shard_id: str) -> Shard:
     req = GetShardRequest()
     req.shard_id.id = shard_id
@@ -54,6 +70,9 @@ async def get_shard(node: AbstractIndexNode, shard_id: str) -> Shard:
         return await node.reader.GetShard(req)  # type: ignore
 
 
+@backoff.on_exception(
+    backoff.expo, Exception, jitter=None, factor=0.1, max_tries=3, giveup=should_giveup
+)
 async def suggest_shard(node: AbstractIndexNode, shard: str, query: SuggestRequest) -> SuggestResponse:
     req = SuggestRequest()
     req.CopyFrom(query)

nucliadb/standalone/introspect.py

@@ -32,7 +32,6 @@ import psutil
 from fastapi import FastAPI
 from pydantic import BaseModel
 
-from nucliadb.common.cluster import manager as cluster_manager
 from nucliadb.standalone.settings import Settings
 from nucliadb_telemetry.settings import LogOutputType, LogSettings
 
@@ -83,7 +82,6 @@ async def stream_tar(app: FastAPI) -> AsyncGenerator[bytes, None]:
         with tarfile.open(tar_file, mode="w:gz") as tar:
             await add_system_info(temp_dir, tar)
             await add_dependencies(temp_dir, tar)
-            await add_cluster_info(temp_dir, tar)
             settings: Settings = app.settings.copy()  # type: ignore
             await add_settings(temp_dir, tar, settings)
             if settings.log_output_type == LogOutputType.FILE:
@@ -145,29 +143,6 @@ def _add_dependencies_to_tar(temp_dir: str, tar: tarfile.TarFile):
     tar.add(dependendies_file, arcname="dependencies.txt")
 
 
-async def add_cluster_info(temp_dir: str, tar: tarfile.TarFile):
-    loop = asyncio.get_event_loop()
-    await loop.run_in_executor(None, _add_cluster_info_to_tar, temp_dir, tar)
-
-
-def _add_cluster_info_to_tar(temp_dir: str, tar: tarfile.TarFile):
-    cluster_info = ClusterInfo(
-        nodes=[
-            NodeInfo(
-                id=node.id,
-                address=node.address,
-                shard_count=node.shard_count,
-                primary_id=node.primary_id,
-            )
-            for node in cluster_manager.get_index_nodes()
-        ]
-    )
-    cluster_info_file = os.path.join(temp_dir, "cluster_info.txt")
-    with open(cluster_info_file, "w") as f:
-        f.write(cluster_info.model_dump_json(indent=4))
-    tar.add(cluster_info_file, arcname="cluster_info.txt")
-
-
 async def add_settings(temp_dir: str, tar: tarfile.TarFile, settings: Settings):
     loop = asyncio.get_event_loop()
     await loop.run_in_executor(None, _add_settings_to_tar, temp_dir, tar, settings)

nucliadb/train/lifecycle.py

@@ -22,10 +22,7 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
-from nucliadb.common.cluster.discovery.utils import (
-    setup_cluster_discovery,
-    teardown_cluster_discovery,
-)
+from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb.train import SERVICE_NAME
 from nucliadb.train.utils import (
     start_shard_manager,
@@ -40,7 +37,7 @@ from nucliadb_utils.utilities import start_audit_utility, stop_audit_utility
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     await setup_telemetry(SERVICE_NAME)
-    await setup_cluster_discovery()
+    await start_nidx_utility()
     await start_shard_manager()
     await start_train_grpc(SERVICE_NAME)
     await start_audit_utility(SERVICE_NAME)
@@ -50,5 +47,5 @@ async def lifespan(app: FastAPI):
     await stop_audit_utility()
     await stop_train_grpc()
     await stop_shard_manager()
-    await teardown_cluster_discovery()
+    await stop_nidx_utility()
     await clean_telemetry(SERVICE_NAME)

nucliadb/train/nodes.py

@@ -45,9 +45,7 @@ from nucliadb_protos.train_pb2 import (
     TrainSentence,
 )
 from nucliadb_protos.writer_pb2 import ShardObject
-from nucliadb_utils import const
 from nucliadb_utils.storages.storage import Storage
-from nucliadb_utils.utilities import has_feature
 
 
 class TrainShardManager(manager.KBShardManager):
@@ -63,9 +61,7 @@ class TrainShardManager(manager.KBShardManager):
         except StopIteration:
             raise KeyError("Shard not found")
 
-        node_obj, shard_id = manager.choose_node(
-            shard_object, use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": kbid})
-        )
+        node_obj, shard_id = manager.choose_node(shard_object)
         return node_obj, shard_id
 
     async def get_kb_obj(self, txn: Transaction, kbid: str) -> Optional[KnowledgeBox]: