nucliadb 6.4.0.post4127__py3-none-any.whl → 6.4.0.post4132__py3-none-any.whl

nucliadb/search/requesters/utils.py

@@ -35,8 +35,6 @@ from nidx_protos.nodereader_pb2 import (
     SuggestResponse,
 )
 
-from nucliadb.common.cluster import manager as cluster_manager
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.exceptions import ShardsNotFound
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.search import logger
@@ -78,7 +76,7 @@ async def node_query(
     method: Method,
     pb_query: SuggestRequest,
     timeout: Optional[float] = None,
-) -> tuple[list[SuggestResponse], bool, list[tuple[AbstractIndexNode, str]]]: ...
+) -> tuple[list[SuggestResponse], bool, list[str]]: ...
 
 
 @overload
@@ -87,7 +85,7 @@ async def node_query(
     method: Method,
     pb_query: SearchRequest,
     timeout: Optional[float] = None,
-) -> tuple[list[SearchResponse], bool, list[tuple[AbstractIndexNode, str]]]: ...
+) -> tuple[list[SearchResponse], bool, list[str]]: ...
 
 
 @overload
@@ -96,7 +94,7 @@ async def node_query(
     method: Method,
     pb_query: GraphSearchRequest,
     timeout: Optional[float] = None,
-) -> tuple[list[GraphSearchResponse], bool, list[tuple[AbstractIndexNode, str]]]: ...
+) -> tuple[list[GraphSearchResponse], bool, list[str]]: ...
 
 
 async def node_query(
@@ -104,7 +102,7 @@ async def node_query(
     method: Method,
     pb_query: REQUEST_TYPE,
     timeout: Optional[float] = None,
-) -> tuple[Sequence[Union[T, BaseException]], bool, list[tuple[AbstractIndexNode, str]]]:
+) -> tuple[Sequence[Union[T, BaseException]], bool, list[str]]:
     timeout = timeout or settings.search_timeout
     shard_manager = get_shard_manager()
     try:
@@ -116,21 +114,17 @@ async def node_query(
         )
 
     ops = []
-    queried_nodes = []
+    queried_shards = []
     incomplete_results = False
 
     for shard_obj in shard_groups:
-        try:
-            node, shard_id = cluster_manager.choose_node(shard_obj)
-        except KeyError:
-            incomplete_results = True
-        else:
-            if shard_id is not None:
-                # At least one node is alive for this shard group
-                # let's add it ot the query list if has a valid value
-                func = METHODS[method]
-                ops.append(func(node, shard_id, pb_query))  # type: ignore
-                queried_nodes.append((node, shard_id))
+        shard_id = shard_obj.nidx_shard_id
+        if shard_id is not None:
+            # At least one node is alive for this shard group
+            # let's add it ot the query list if has a valid value
+            func = METHODS[method]
+            ops.append(func(shard_id, pb_query))  # type: ignore
+            queried_shards.append(shard_id)
 
     if not ops:
         logger.warning(f"No node found for any of this resources shards {kbid}")
@@ -146,8 +140,7 @@ async def node_query(
         )
     except asyncio.TimeoutError as exc:  # pragma: no cover
         logger.warning(
-            "Timeout while querying nodes",
-            extra={"nodes": debug_nodes_info(queried_nodes)},
+            "Timeout while querying nidx",
         )
         results = [exc]
 
@@ -164,7 +157,7 @@ async def node_query(
         )
         raise error
 
-    return results, incomplete_results, queried_nodes
+    return results, incomplete_results, queried_shards
 
 
 def validate_node_query_results(results: list[Any]) -> Optional[HTTPException]:
@@ -201,15 +194,3 @@ def validate_node_query_results(results: list[Any]) -> Optional[HTTPException]:
             return HTTPException(status_code=status_code, detail=reason)
 
     return None
-
-
-def debug_nodes_info(nodes: list[tuple[AbstractIndexNode, str]]) -> list[dict[str, str]]:
-    details: list[dict[str, str]] = []
-    for node, shard_id in nodes:
-        info = {
-            "id": node.id,
-            "shard_id": shard_id,
-            "address": "nidx",
-        }
-        details.append(info)
-    return details

nucliadb/search/search/find.py

@@ -23,7 +23,7 @@ from time import time
 from nucliadb.common.external_index_providers.base import ExternalIndexManager
 from nucliadb.common.external_index_providers.manager import get_external_index_manager
 from nucliadb.common.models_utils import to_proto
-from nucliadb.search.requesters.utils import Method, debug_nodes_info, node_query
+from nucliadb.search.requesters.utils import Method, node_query
 from nucliadb.search.search.find_merge import (
     build_find_response,
     compose_find_resources,
@@ -105,7 +105,7 @@ async def _index_node_retrieval(
         ) = await legacy_convert_retrieval_to_proto(parsed)
 
     with metrics.time("node_query"):
-        results, query_incomplete_results, queried_nodes = await node_query(
+        results, query_incomplete_results, queried_shards = await node_query(
             kbid, Method.SEARCH, pb_query
         )
     incomplete_results = incomplete_results or query_incomplete_results
@@ -139,10 +139,6 @@ async def _index_node_retrieval(
             retrieval_rephrased_question=rephrased_query,
         )
 
-    if item.debug:
-        search_results.nodes = debug_nodes_info(queried_nodes)
-
-    queried_shards = [shard_id for _, shard_id in queried_nodes]
     search_results.shards = queried_shards
     search_results.autofilters = autofilters
 
@@ -156,7 +152,6 @@ async def _index_node_retrieval(
                 "client": x_ndb_client,
                 "query": item.model_dump_json(),
                 "time": search_time,
-                "nodes": debug_nodes_info(queried_nodes),
                 "durations": metrics.steps(),
             },
         )
@@ -169,7 +164,6 @@ async def _index_node_retrieval(
                 "client": x_ndb_client,
                 "query": item.model_dump_json(),
                 "time": search_time,
-                "nodes": debug_nodes_info(queried_nodes),
                 "durations": metrics.steps(),
             },
         )

nucliadb/search/search/shards.py

@@ -17,7 +17,6 @@
 # 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/>.
 #
-import asyncio
 
 import backoff
 from grpc import StatusCode
@@ -33,16 +32,7 @@ from nidx_protos.nodereader_pb2 import (
 )
 from nidx_protos.noderesources_pb2 import Shard
 
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb_telemetry import metrics
-
-node_observer = metrics.Observer(
-    "node_client",
-    labels={"type": "", "node_id": ""},
-    error_mappings={
-        "timeout": asyncio.CancelledError,
-    },
-)
+from nucliadb.common.nidx import get_nidx_api_client, get_nidx_searcher_client
 
 
 def should_giveup(e: Exception):
@@ -54,43 +44,37 @@ def should_giveup(e: Exception):
 @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:
+async def query_shard(shard: str, query: SearchRequest) -> SearchResponse:
     req = SearchRequest()
     req.CopyFrom(query)
     req.shard = shard
-    with node_observer({"type": "search", "node_id": node.id}):
-        return await node.reader.Search(req)  # type: ignore
+    return await get_nidx_searcher_client().Search(req)
 
 
 @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:
+async def get_shard(shard_id: str) -> Shard:
     req = GetShardRequest()
     req.shard_id.id = shard_id
-    with node_observer({"type": "get_shard", "node_id": node.id}):
-        return await node.reader.GetShard(req)  # type: ignore
+    return await get_nidx_api_client().GetShard(req)
 
 
 @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:
+async def suggest_shard(shard: str, query: SuggestRequest) -> SuggestResponse:
     req = SuggestRequest()
     req.CopyFrom(query)
     req.shard = shard
-    with node_observer({"type": "suggest", "node_id": node.id}):
-        return await node.reader.Suggest(req)  # type: ignore
+    return await get_nidx_searcher_client().Suggest(req)
 
 
 @backoff.on_exception(
     backoff.expo, Exception, jitter=None, factor=0.1, max_tries=3, giveup=should_giveup
 )
-async def graph_search_shard(
-    node: AbstractIndexNode, shard: str, query: GraphSearchRequest
-) -> GraphSearchResponse:
+async def graph_search_shard(shard: str, query: GraphSearchRequest) -> GraphSearchResponse:
     req = GraphSearchRequest()
     req.CopyFrom(query)
     req.shard = shard
-    with node_observer({"type": "graph_search", "node_id": node.id}):
-        return await node.reader.GraphSearch(req)  # type: ignore
+    return await get_nidx_searcher_client().GraphSearch(req)

nucliadb/train/generator.py

@@ -54,7 +54,7 @@ from nucliadb_protos.dataset_pb2 import TaskType, TrainSet
 async def generate_train_data(kbid: str, shard: str, trainset: TrainSet):
     # Get the data structure to generate data
     shard_manager = get_shard_manager()
-    node, shard_replica_id = await shard_manager.get_reader(kbid, shard)
+    shard_replica_id = await shard_manager.get_shard_id(kbid, shard)
 
     if trainset.batch_size == 0:
         trainset.batch_size = 50
@@ -62,24 +62,22 @@ async def generate_train_data(kbid: str, shard: str, trainset: TrainSet):
     batch_generator: Optional[AsyncIterator[TrainBatch]] = None
 
     if trainset.type == TaskType.FIELD_CLASSIFICATION:
-        batch_generator = field_classification_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = field_classification_batch_generator(kbid, trainset, shard_replica_id)
     elif trainset.type == TaskType.IMAGE_CLASSIFICATION:
-        batch_generator = image_classification_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = image_classification_batch_generator(kbid, trainset, shard_replica_id)
     elif trainset.type == TaskType.PARAGRAPH_CLASSIFICATION:
-        batch_generator = paragraph_classification_batch_generator(
-            kbid, trainset, node, shard_replica_id
-        )
+        batch_generator = paragraph_classification_batch_generator(kbid, trainset, shard_replica_id)
     elif trainset.type == TaskType.TOKEN_CLASSIFICATION:
-        batch_generator = token_classification_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = token_classification_batch_generator(kbid, trainset, shard_replica_id)
     elif trainset.type == TaskType.SENTENCE_CLASSIFICATION:
-        batch_generator = sentence_classification_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = sentence_classification_batch_generator(kbid, trainset, shard_replica_id)
     elif trainset.type == TaskType.PARAGRAPH_STREAMING:
-        batch_generator = paragraph_streaming_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = paragraph_streaming_batch_generator(kbid, trainset, shard_replica_id)
 
     elif trainset.type == TaskType.QUESTION_ANSWER_STREAMING:
-        batch_generator = question_answer_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = question_answer_batch_generator(kbid, trainset, shard_replica_id)
     elif trainset.type == TaskType.FIELD_STREAMING:
-        batch_generator = field_streaming_batch_generator(kbid, trainset, node, shard_replica_id)
+        batch_generator = field_streaming_batch_generator(kbid, trainset, shard_replica_id)
 
     if batch_generator is None:
         raise HTTPException(

nucliadb/train/generators/field_classifier.py

@@ -22,8 +22,8 @@ from typing import AsyncGenerator
 
 from nidx_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
@@ -37,10 +37,9 @@ from nucliadb_protos.dataset_pb2 import (
 def field_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[FieldClassificationBatch, None]:
-    generator = generate_field_classification_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_field_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, FieldClassificationBatch)
     return batch_generator
 
@@ -48,7 +47,6 @@ def field_classification_batch_generator(
 async def generate_field_classification_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[TextLabel, None]:
     labelset = f"/l/{trainset.filter.labels[0]}"
@@ -59,7 +57,7 @@ async def generate_field_classification_payloads(
     request.filter.labels.append(labelset)
     total = 0
 
-    async for document_item in node.stream_get_fields(request):
+    async for document_item in get_nidx_searcher_client().Documents(request):
         text_labels = []
         for label in document_item.labels:
             if label.startswith(labelset):

nucliadb/train/generators/field_streaming.py

@@ -22,8 +22,8 @@ from typing import AsyncGenerator, Optional
 
 from nidx_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
@@ -38,10 +38,9 @@ from nucliadb_protos.utils_pb2 import ExtractedText
 def field_streaming_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[FieldStreamingBatch, None]:
-    generator = generate_field_streaming_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_field_streaming_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, FieldStreamingBatch)
     return batch_generator
 
@@ -49,7 +48,6 @@ def field_streaming_batch_generator(
 async def generate_field_streaming_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[FieldSplitData, None]:
     # Query how many resources has each label
@@ -77,7 +75,7 @@ async def generate_field_streaming_payloads(
     total = 0
     resources = set()
 
-    async for document_item in node.stream_get_fields(request):
+    async for document_item in get_nidx_searcher_client().Documents(request):
         text_labels = []
         for label in document_item.labels:
             text_labels.append(label)

nucliadb/train/generators/image_classifier.py

@@ -20,7 +20,6 @@
 
 from typing import AsyncGenerator
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.train.generators.utils import batchify
 from nucliadb_protos.dataset_pb2 import (
     ImageClassification,
@@ -32,10 +31,9 @@ from nucliadb_protos.dataset_pb2 import (
 def image_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ImageClassificationBatch, None]:
-    generator = generate_image_classification_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_image_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ImageClassificationBatch)
     return batch_generator
 
@@ -43,7 +41,6 @@ def image_classification_batch_generator(
 async def generate_image_classification_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ImageClassification, None]:
     # NOTE: image classifications are no longer supported, as the page selection annotations were removed

nucliadb/train/generators/paragraph_classifier.py

@@ -23,7 +23,7 @@ from typing import AsyncGenerator
 from fastapi import HTTPException
 from nidx_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train.generators.utils import batchify, get_paragraph
 from nucliadb_protos.dataset_pb2 import (
     Label,
@@ -36,7 +36,6 @@ from nucliadb_protos.dataset_pb2 import (
 def paragraph_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ParagraphClassificationBatch, None]:
     if len(trainset.filter.labels) != 1:
@@ -45,7 +44,7 @@ def paragraph_classification_batch_generator(
             detail="Paragraph Classification should be of 1 labelset",
         )
 
-    generator = generate_paragraph_classification_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_paragraph_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ParagraphClassificationBatch)
     return batch_generator
 
@@ -53,7 +52,6 @@ def paragraph_classification_batch_generator(
 async def generate_paragraph_classification_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[TextLabel, None]:
     labelset = f"/l/{trainset.filter.labels[0]}"
@@ -63,7 +61,7 @@ async def generate_paragraph_classification_payloads(
     request.shard_id.id = shard_replica_id
     request.filter.labels.append(labelset)
 
-    async for paragraph_item in node.stream_get_paragraphs(request):
+    async for paragraph_item in get_nidx_searcher_client().Paragraphs(request):
         text_labels = []
         for label in paragraph_item.labels:
             if label.startswith(labelset):

nucliadb/train/generators/paragraph_streaming.py

@@ -22,8 +22,8 @@ from typing import AsyncGenerator
 
 from nidx_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
@@ -36,10 +36,9 @@ from nucliadb_protos.dataset_pb2 import (
 def paragraph_streaming_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ParagraphStreamingBatch, None]:
-    generator = generate_paragraph_streaming_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_paragraph_streaming_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, ParagraphStreamingBatch)
     return batch_generator
 
@@ -47,7 +46,6 @@ def paragraph_streaming_batch_generator(
 async def generate_paragraph_streaming_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[ParagraphStreamItem, None]:
     """Streams paragraphs ordered as if they were read sequentially from each
@@ -57,7 +55,7 @@ async def generate_paragraph_streaming_payloads(
     request = StreamRequest()
     request.shard_id.id = shard_replica_id
 
-    async for document_item in node.stream_get_fields(request):
+    async for document_item in get_nidx_searcher_client().Documents(request):
         field_id = f"{document_item.uuid}{document_item.field}"
         rid, field_type, field = field_id.split("/")
 

nucliadb/train/generators/question_answer_streaming.py

@@ -22,8 +22,8 @@ from typing import AsyncGenerator
 
 from nidx_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR, FIELD_TYPE_STR_TO_PB
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import (
     batchify,
@@ -45,10 +45,9 @@ from nucliadb_protos.resources_pb2 import (
 def question_answer_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[QuestionAnswerStreamingBatch, None]:
-    generator = generate_question_answer_streaming_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_question_answer_streaming_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, QuestionAnswerStreamingBatch)
     return batch_generator
 
@@ -56,13 +55,12 @@ def question_answer_batch_generator(
 async def generate_question_answer_streaming_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ):
     request = StreamRequest()
     request.shard_id.id = shard_replica_id
 
-    async for document_item in node.stream_get_fields(request):
+    async for document_item in get_nidx_searcher_client().Documents(request):
         field_id = f"{document_item.uuid}{document_item.field}"
         rid, field_type, field = field_id.split("/")
 

nucliadb/train/generators/sentence_classifier.py

@@ -23,8 +23,8 @@ from typing import AsyncGenerator
 from fastapi import HTTPException
 from nidx_protos.nodereader_pb2 import StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
@@ -38,7 +38,6 @@ from nucliadb_protos.dataset_pb2 import (
 def sentence_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[SentenceClassificationBatch, None]:
     if len(trainset.filter.labels) == 0:
@@ -47,7 +46,7 @@ def sentence_classification_batch_generator(
             detail="Sentence Classification should be at least of 1 labelset",
         )
 
-    generator = generate_sentence_classification_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_sentence_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, SentenceClassificationBatch)
     return batch_generator
 
@@ -55,7 +54,6 @@ def sentence_classification_batch_generator(
 async def generate_sentence_classification_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[MultipleTextSameLabels, None]:
     labelsets = []
@@ -67,7 +65,7 @@ async def generate_sentence_classification_payloads(
         labelsets.append(labelset)
         request.filter.labels.append(labelset)
 
-    async for paragraph_item in node.stream_get_paragraphs(request):
+    async for paragraph_item in get_nidx_searcher_client().Paragraphs(request):
         text_labels: list[str] = []
         for label in paragraph_item.labels:
             for labelset in labelsets:

nucliadb/train/generators/token_classifier.py

@@ -23,8 +23,8 @@ from typing import AsyncGenerator, cast
 
 from nidx_protos.nodereader_pb2 import StreamFilter, StreamRequest
 
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.nidx import get_nidx_searcher_client
 from nucliadb.train import logger
 from nucliadb.train.generators.utils import batchify, get_resource_from_cache_or_db
 from nucliadb_protos.dataset_pb2 import (
@@ -41,10 +41,9 @@ MAIN = "__main__"
 def token_classification_batch_generator(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[TokenClassificationBatch, None]:
-    generator = generate_token_classification_payloads(kbid, trainset, node, shard_replica_id)
+    generator = generate_token_classification_payloads(kbid, trainset, shard_replica_id)
     batch_generator = batchify(generator, trainset.batch_size, TokenClassificationBatch)
     return batch_generator
 
@@ -52,7 +51,6 @@ def token_classification_batch_generator(
 async def generate_token_classification_payloads(
     kbid: str,
     trainset: TrainSet,
-    node: AbstractIndexNode,
     shard_replica_id: str,
 ) -> AsyncGenerator[TokensClassification, None]:
     request = StreamRequest()
@@ -60,7 +58,7 @@ async def generate_token_classification_payloads(
     for entitygroup in trainset.filter.labels:
         request.filter.labels.append(f"/e/{entitygroup}")
         request.filter.conjunction = StreamFilter.Conjunction.OR
-    async for field_item in node.stream_get_fields(request):
+    async for field_item in get_nidx_searcher_client().Documents(request):
         _, field_type, field = field_item.field.split("/")
         (
             split_text,

nucliadb/train/nodes.py

@@ -21,7 +21,6 @@ from typing import AsyncIterator, Optional
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster import manager
-from nucliadb.common.cluster.base import AbstractIndexNode
 
 # XXX: this keys shouldn't be exposed outside datamanagers
 from nucliadb.common.datamanagers.resources import KB_RESOURCE_SLUG_BASE
@@ -54,15 +53,14 @@ class TrainShardManager(manager.KBShardManager):
         self.driver = driver
         self.storage = storage
 
-    async def get_reader(self, kbid: str, shard: str) -> tuple[AbstractIndexNode, str]:
+    async def get_shard_id(self, kbid: str, shard: str) -> str:
         shards = await self.get_shards_by_kbid_inner(kbid)
         try:
             shard_object: ShardObject = next(filter(lambda x: x.shard == shard, shards.shards))
         except StopIteration:
             raise KeyError("Shard not found")
 
-        node_obj, shard_id = manager.choose_node(shard_object)
-        return node_obj, shard_id
+        return shard_object.nidx_shard_id
 
     async def get_kb_obj(self, txn: Transaction, kbid: str) -> Optional[KnowledgeBox]:
         if kbid is None:

{nucliadb-6.4.0.post4127.dist-info → nucliadb-6.4.0.post4132.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4127
+Version: 6.4.0.post4132
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4127
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4127
-Requires-Dist: nucliadb-protos>=6.4.0.post4127
-Requires-Dist: nucliadb-models>=6.4.0.post4127
-Requires-Dist: nidx-protos>=6.4.0.post4127
+Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4132
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4132
+Requires-Dist: nucliadb-protos>=6.4.0.post4132
+Requires-Dist: nucliadb-models>=6.4.0.post4132
+Requires-Dist: nidx-protos>=6.4.0.post4132
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]