nucliadb 6.2.1.post2954__py3-none-any.whl → 6.2.1.post2972__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

@@ -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"),
         }
@@ -158,11 +159,8 @@ class NidxServiceUtility(NidxUtility):
         return res.seq
 
 
-async def start_nidx_utility() -> Optional[NidxUtility]:
-    if not NIDX_ENABLED:
-        return None
-
-    nidx = get_nidx()
+async def start_nidx_utility() -> NidxUtility:
+    nidx = get_utility(Utility.NIDX)
     if nidx:
         return nidx
 
@@ -178,30 +176,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 +253,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/predict.py

@@ -21,7 +21,7 @@ import json
 import os
 import random
 from enum import Enum
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncGenerator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
@@ -268,7 +268,7 @@ class PredictEngine:
     @predict_observer.wrap({"type": "chat_ndjson"})
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
+    ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         """
         Chat query using the new stream format
         Format specs: https://github.com/ndjson/ndjson-spec
@@ -444,7 +444,7 @@ class DummyPredictEngine(PredictEngine):
 
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
+    ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         self.calls.append(("chat_query_ndjson", item))
 
         async def generate():
@@ -555,7 +555,7 @@ def get_answer_generator(response: aiohttp.ClientResponse):
 
 def get_chat_ndjson_generator(
     response: aiohttp.ClientResponse,
-) -> AsyncIterator[GenerativeChunk]:
+) -> AsyncGenerator[GenerativeChunk, None]:
     async def _parse_generative_chunks(gen):
         async for chunk in gen:
             try:

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/chat/ask.py

@@ -129,7 +129,7 @@ class AskResult:
         main_results: KnowledgeboxFindResults,
         prequeries_results: Optional[list[PreQueryResult]],
         nuclia_learning_id: Optional[str],
-        predict_answer_stream: AsyncGenerator[GenerativeChunk, None],
+        predict_answer_stream: Optional[AsyncGenerator[GenerativeChunk, None]],
         prompt_context: PromptContext,
         prompt_context_order: PromptContextOrder,
         auditor: ChatAuditor,
@@ -396,6 +396,9 @@ class AskResult:
         This method does not assume any order in the stream of items, but it assumes that at least
         the answer text is streamed in order.
         """
+        if self.predict_answer_stream is None:
+            # In some cases, clients may want to skip the answer generation step
+            return
         async for generative_chunk in self.predict_answer_stream:
             item = generative_chunk.chunk
             if isinstance(item, TextGenerativeResponse):
@@ -562,14 +565,18 @@ async def ask(
         rerank_context=False,
         top_k=ask_request.top_k,
     )
-    with metrics.time("stream_start"):
-        predict = get_predict()
-        (
-            nuclia_learning_id,
-            nuclia_learning_model,
-            predict_answer_stream,
-        ) = await predict.chat_query_ndjson(kbid, chat_model)
-        debug_chat_model = chat_model
+
+    nuclia_learning_id = None
+    nuclia_learning_model = None
+    predict_answer_stream = None
+    if ask_request.generate_answer:
+        with metrics.time("stream_start"):
+            predict = get_predict()
+            (
+                nuclia_learning_id,
+                nuclia_learning_model,
+                predict_answer_stream,
+            ) = await predict.chat_query_ndjson(kbid, chat_model)
 
     auditor = ChatAuditor(
         kbid=kbid,
@@ -590,13 +597,13 @@ async def ask(
         main_results=retrieval_results.main_query,
         prequeries_results=retrieval_results.prequeries,
         nuclia_learning_id=nuclia_learning_id,
-        predict_answer_stream=predict_answer_stream,  # type: ignore
+        predict_answer_stream=predict_answer_stream,
         prompt_context=prompt_context,
         prompt_context_order=prompt_context_order,
         auditor=auditor,
         metrics=metrics,
         best_matches=retrieval_results.best_matches,
-        debug_chat_model=debug_chat_model,
+        debug_chat_model=chat_model,
     )
 
 

nucliadb/search/search/chat/query.py

@@ -349,7 +349,7 @@ class ChatAuditor:
         learning_id: Optional[str],
         query_context: PromptContext,
         query_context_order: PromptContextOrder,
-        model: str,
+        model: Optional[str],
     ):
         self.kbid = kbid
         self.user_id = user_id

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,6 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
-from nucliadb.common.cluster.discovery.utils import (
-    setup_cluster_discovery,
-    teardown_cluster_discovery,
-)
 from nucliadb.train import SERVICE_NAME
 from nucliadb.train.utils import (
     start_shard_manager,
@@ -40,7 +36,6 @@ 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_shard_manager()
     await start_train_grpc(SERVICE_NAME)
     await start_audit_utility(SERVICE_NAME)
@@ -50,5 +45,4 @@ async def lifespan(app: FastAPI):
     await stop_audit_utility()
     await stop_train_grpc()
     await stop_shard_manager()
-    await teardown_cluster_discovery()
     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]: