nucliadb 6.7.2.post4862__py3-none-any.whl → 6.9.2.post5282__py3-none-any.whl

nucliadb/export_import/datamanager.py

@@ -54,7 +54,7 @@ class ExportImportDataManager:
 
     async def get_metadata(self, type: str, kbid: str, id: str) -> Metadata:
         key = self._get_maindb_metadata_key(type, kbid, id)
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             data = await txn.get(key)
         if data is None or data == b"":
             raise MetadataNotFound()
@@ -89,13 +89,13 @@ class ExportImportDataManager:
         metadata.modified = datetime.now(timezone.utc)
         key = self._get_maindb_metadata_key(type, metadata.kbid, metadata.id)
         data = metadata.model_dump_json().encode("utf-8")
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             await txn.set(key, data)
             await txn.commit()
 
     async def delete_metadata(self, type: str, metadata: Metadata):
         key = self._get_maindb_metadata_key(type, metadata.kbid, metadata.id)
-        async with self.driver.transaction() as txn:
+        async with self.driver.rw_transaction() as txn:
             await txn.delete(key)
             await txn.commit()
 

nucliadb/ingest/consumer/pull.py

@@ -31,6 +31,7 @@ from opentelemetry.trace import (
     Link,
 )
 
+from nucliadb.common.http_clients.exceptions import ServiceUnavailableException
 from nucliadb.common.http_clients.processing import (
     ProcessingHTTPClient,
     ProcessingPullMessageProgressUpdater,
@@ -209,6 +210,12 @@ class PullV2Worker:
                             payload_length = len(base64.b64decode(data.payload))
                         logger.error(f"Message too big for transaction: {payload_length}")
                     raise e
+
+                except ServiceUnavailableException as ex:
+                    logger.warning(f"Processing api is unavailable, will retry shortly: {ex}")
+                    await processing_http_client.reset_session()
+                    await asyncio.sleep(self.pull_time_error_backoff)
+
                 except Exception:
                     logger.exception("Unhandled error pulling messages from processing")
                     await asyncio.sleep(self.pull_time_error_backoff)

nucliadb/ingest/consumer/service.py

@@ -22,21 +22,18 @@ import sys
 from functools import partial
 from typing import Awaitable, Callable, Optional
 
-from nucliadb.common.back_pressure.materializer import BackPressureMaterializer
-from nucliadb.common.back_pressure.settings import settings as back_pressure_settings
 from nucliadb.common.maindb.utils import setup_driver
 from nucliadb.ingest import SERVICE_NAME, logger
 from nucliadb.ingest.consumer.consumer import IngestConsumer
 from nucliadb.ingest.consumer.pull import PullV2Worker
 from nucliadb.ingest.settings import settings
 from nucliadb_utils.exceptions import ConfigurationError
-from nucliadb_utils.settings import indexing_settings, transaction_settings
+from nucliadb_utils.settings import transaction_settings
 from nucliadb_utils.utilities import (
     get_audit,
     get_nats_manager,
     get_pubsub,
     get_storage,
-    start_nats_manager,
 )
 
 from .auditing import IndexAuditHandler, ResourceWritesAuditHandler
@@ -57,27 +54,6 @@ async def _exit_tasks(tasks: list[asyncio.Task]) -> None:
     await asyncio.gather(*tasks, return_exceptions=True)
 
 
-async def start_back_pressure() -> BackPressureMaterializer:
-    logger.info("Starting back pressure materializer")
-    nats_manager = await start_nats_manager(
-        SERVICE_NAME,
-        indexing_settings.index_jetstream_servers,
-        indexing_settings.index_jetstream_auth,
-    )
-    back_pressure = BackPressureMaterializer(
-        nats_manager,
-        indexing_check_interval=back_pressure_settings.indexing_check_interval,
-        ingest_check_interval=back_pressure_settings.ingest_check_interval,
-    )
-    await back_pressure.start()
-    return back_pressure
-
-
-async def stop_back_pressure(materializer: BackPressureMaterializer) -> None:
-    await materializer.stop()
-    await materializer.nats_manager.finalize()
-
-
 async def start_ingest_consumers(
     service_name: Optional[str] = None,
 ) -> Callable[[], Awaitable[None]]:
@@ -164,9 +140,8 @@ async def start_shard_creator() -> Callable[[], Awaitable[None]]:
     driver = await setup_driver()
     pubsub = await get_pubsub()
     assert pubsub is not None, "Pubsub is not configured"
-    storage = await get_storage(service_name=SERVICE_NAME)
 
-    shard_creator = ShardCreatorHandler(driver=driver, storage=storage, pubsub=pubsub)
+    shard_creator = ShardCreatorHandler(driver=driver, pubsub=pubsub)
     await shard_creator.initialize()
 
     return shard_creator.finalize

nucliadb/ingest/consumer/shard_creator.py

@@ -25,14 +25,14 @@ from typing import Any
 
 from nidx_protos import nodereader_pb2, noderesources_pb2
 
-from nucliadb.common import locking
+from nucliadb.common import datamanagers, locking
+from nucliadb.common.cluster.settings import settings
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.nidx import get_nidx_api_client
 from nucliadb_protos import writer_pb2
 from nucliadb_utils import const
 from nucliadb_utils.cache.pubsub import PubSubDriver
-from nucliadb_utils.storages.storage import Storage
 
 from . import metrics
 from .utils import DelayedTaskHandler
@@ -52,12 +52,10 @@ class ShardCreatorHandler:
         self,
         *,
         driver: Driver,
-        storage: Storage,
         pubsub: PubSubDriver,
         check_delay: float = 10.0,
     ):
         self.driver = driver
-        self.storage = storage
         self.pubsub = pubsub
         self.shard_manager = get_shard_manager()
         self.task_handler = DelayedTaskHandler(check_delay)
@@ -91,7 +89,7 @@ class ShardCreatorHandler:
     @metrics.handler_histo.wrap({"type": "shard_creator"})
     async def process_kb(self, kbid: str) -> None:
         logger.info({"message": "Processing notification for kbid", "kbid": kbid})
-        async with self.driver.transaction(read_only=True) as txn:
+        async with self.driver.ro_transaction() as txn:
             current_shard = await self.shard_manager.get_current_active_shard(txn, kbid)
 
         if current_shard is None:
@@ -111,4 +109,17 @@ class ShardCreatorHandler:
                     shard_id=noderesources_pb2.ShardId(id=current_shard.nidx_shard_id)
                 )  # type: ignore
             )
-            await self.shard_manager.maybe_create_new_shard(kbid, shard.paragraphs)
+
+            if not should_create_new_shard(shard.paragraphs):
+                return
+
+            logger.info({"message": "Adding shard", "kbid": kbid})
+            async with datamanagers.with_rw_transaction() as txn:
+                kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
+                prewarm = kb_config is not None and kb_config.prewarm_enabled
+                await self.shard_manager.create_shard_by_kbid(txn, kbid, prewarm_enabled=prewarm)
+                await txn.commit()
+
+
+def should_create_new_shard(num_paragraphs: int) -> bool:
+    return num_paragraphs > settings.max_shard_paragraphs

nucliadb/ingest/fields/base.py

@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any, Generic, Optional, Type, TypeVar
 from google.protobuf.message import DecodeError, Message
 
 from nucliadb.common import datamanagers
+from nucliadb.common.ids import FieldId
 from nucliadb.ingest.fields.exceptions import InvalidFieldClass, InvalidPBClass
 from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import (
@@ -46,10 +47,8 @@ from nucliadb_protos.resources_pb2 import (
 )
 from nucliadb_protos.utils_pb2 import ExtractedText, VectorObject
 from nucliadb_protos.writer_pb2 import Error, FieldStatus
-from nucliadb_utils import const
 from nucliadb_utils.storages.exceptions import CouldNotCopyNotFound
 from nucliadb_utils.storages.storage import Storage, StorageField
-from nucliadb_utils.utilities import has_feature
 
 logger = logging.getLogger(__name__)
 
@@ -125,6 +124,14 @@ class Field(Generic[PbType]):
     def uuid(self) -> str:
         return self.resource.uuid
 
+    @property
+    def field_id(self) -> FieldId:
+        return FieldId(
+            rid=self.resource.uuid,
+            type=self.type,
+            key=self.id,
+        )
+
     @property
     def storage(self) -> Storage:
         return self.resource.storage
@@ -215,21 +222,6 @@ class Field(Generic[PbType]):
     ) -> None:
         # Try delete vectors
         sf = self._get_extracted_vectors_storage_field(vectorset, storage_key_kind)
-
-        if has_feature(const.Features.DEBUG_MISSING_VECTORS):
-            # This is a very chatty log. It is just a temporary hint while debugging an issue.
-            logger.info(
-                "Deleting vectors from storage",
-                extra={
-                    "kbid": self.kbid,
-                    "rid": self.resource.uuid,
-                    "field": f"{self.type}/{self.id}",
-                    "vectorset": vectorset,
-                    "storage_key_kind": storage_key_kind,
-                    "key": sf.key,
-                    "bucket": sf.bucket,
-                },
-            )
         try:
             await self.storage.delete_upload(sf.key, sf.bucket)
         except KeyError:

nucliadb/ingest/fields/conversation.py

@@ -21,13 +21,16 @@ import uuid
 from typing import Any, Optional
 
 from nucliadb.ingest.fields.base import Field
-from nucliadb_protos.resources_pb2 import CloudFile, FieldConversation
+from nucliadb_protos.resources_pb2 import CloudFile, FieldConversation, SplitMetadata, SplitsMetadata
 from nucliadb_protos.resources_pb2 import Conversation as PBConversation
 from nucliadb_utils.storages.storage import StorageField
 
+MAX_CONVERSATION_MESSAGES = 50 * 1024
+
 PAGE_SIZE = 200
 
 CONVERSATION_PAGE_VALUE = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/{page}"
+CONVERSATION_SPLITS_METADATA = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/splits_metadata"
 CONVERSATION_METADATA = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}"
 
 
@@ -52,9 +55,22 @@ class Conversation(Field[PBConversation]):
     ):
         super(Conversation, self).__init__(id, resource, pb, value)
         self.value = {}
+        self._splits_metadata: Optional[SplitsMetadata] = None
+        self.metadata = None
+
+    async def delete_value(self):
+        await self.resource.txn.delete_by_prefix(
+            CONVERSATION_METADATA.format(kbid=self.kbid, uuid=self.uuid, type=self.type, field=self.id)
+        )
+        self._split_metadata = None
         self.metadata = None
+        self.value.clear()
 
     async def set_value(self, payload: PBConversation):
+        if payload.replace_field:
+            # As we need to overwrite the value of the conversation, first delete any previous data.
+            await self.delete_value()
+
         metadata = await self.get_metadata()
         metadata.extract_strategy = payload.extract_strategy
         metadata.split_strategy = payload.split_strategy
@@ -70,10 +86,13 @@ class Conversation(Field[PBConversation]):
             last_page = PBConversation()
             metadata.pages += 1
 
+        self._splits_metadata = await self.get_splits_metadata()
+
         # Make sure message attachment files are on our region. This is needed
         # to support the hybrid-onprem deployment as the attachments must be stored
         # at the storage services of the client's premises.
         for message in payload.messages:
+            self._splits_metadata.metadata.setdefault(message.ident, SplitMetadata())
             new_message_files = []
             for idx, file in enumerate(message.content.attachments):
                 if self.storage.needs_move(file, self.kbid):
@@ -117,6 +136,7 @@ class Conversation(Field[PBConversation]):
 
         # Finally, set the metadata
         await self.db_set_metadata(metadata)
+        await self.set_splits_metadata(self._splits_metadata)
 
     async def get_value(self, page: Optional[int] = None) -> Optional[PBConversation]:
         # If no page was requested, force fetch of metadata
@@ -203,3 +223,29 @@ class Conversation(Field[PBConversation]):
         self.metadata = payload
         self.resource.modified = True
         self._created = False
+
+    async def get_splits_metadata(self) -> SplitsMetadata:
+        if self._splits_metadata is None:
+            field_key = CONVERSATION_SPLITS_METADATA.format(
+                kbid=self.kbid,
+                uuid=self.uuid,
+                type=self.type,
+                field=self.id,
+            )
+            payload = await self.resource.txn.get(field_key)
+            if payload is None:
+                return SplitsMetadata()
+            self._splits_metadata = SplitsMetadata()
+            self._splits_metadata.ParseFromString(payload)
+        return self._splits_metadata
+
+    async def set_splits_metadata(self, payload: SplitsMetadata) -> None:
+        key = CONVERSATION_SPLITS_METADATA.format(
+            kbid=self.kbid,
+            uuid=self.uuid,
+            type=self.type,
+            field=self.id,
+        )
+        await self.resource.txn.set(key, payload.SerializeToString())
+        self._split_metadata = payload
+        self.resource.modified = True

nucliadb/ingest/orm/brain_v2.py

@@ -193,7 +193,7 @@ class ResourceBrain:
         if field_author is not None and field_author.WhichOneof("author") == "data_augmentation":
             field_type, field_id = field_key.split("/")
             da_task_id = ids.extract_data_augmentation_id(field_id)
-            if da_task_id is None:  # pragma: nocover
+            if da_task_id is None:  # pragma: no cover
                 logger.warning(
                     "Data augmentation field id has an unexpected format! Skipping label",
                     extra={
@@ -217,6 +217,7 @@ class ResourceBrain:
         replace_field: bool,
         skip_paragraphs_index: Optional[bool],
         skip_texts_index: Optional[bool],
+        append_splits: Optional[set[str]] = None,
     ) -> None:
         # We need to add the extracted text to the texts section of the Resource so that
         # the paragraphs can be indexed
@@ -234,6 +235,7 @@ class ResourceBrain:
             user_field_metadata,
             replace_field=replace_field,
             skip_paragraphs=skip_paragraphs_index,
+            append_splits=append_splits,
         )
 
     @observer.wrap({"type": "apply_field_paragraphs"})
@@ -246,6 +248,7 @@ class ResourceBrain:
         user_field_metadata: Optional[UserFieldMetadata],
         replace_field: bool,
         skip_paragraphs: Optional[bool],
+        append_splits: Optional[set[str]] = None,
     ) -> None:
         if skip_paragraphs is not None:
             self.brain.skip_paragraphs = skip_paragraphs
@@ -254,7 +257,12 @@ class ResourceBrain:
         paragraph_pages = ParagraphPages(page_positions) if page_positions else None
         # Splits of the field
         for subfield, field_metadata in field_computed_metadata.split_metadata.items():
-            extracted_text_str = extracted_text.split_text[subfield] if extracted_text else None
+            if should_skip_split_indexing(subfield, replace_field, append_splits):
+                continue
+            if subfield not in extracted_text.split_text:
+                # No extracted text for this split
+                continue
+            extracted_text_str = extracted_text.split_text[subfield]
             for idx, paragraph in enumerate(field_metadata.paragraphs):
                 key = f"{self.rid}/{field_key}/{subfield}/{paragraph.start}-{paragraph.end}"
                 denied_classifications = set(user_paragraph_classifications.denied.get(key, []))
@@ -308,7 +316,7 @@ class ResourceBrain:
                 self.brain.paragraphs[field_key].paragraphs[key].CopyFrom(p)
 
         # Main field
-        extracted_text_str = extracted_text.text if extracted_text else None
+        extracted_text_str = extracted_text.text
         for idx, paragraph in enumerate(field_computed_metadata.metadata.paragraphs):
             key = f"{self.rid}/{field_key}/{paragraph.start}-{paragraph.end}"
             denied_classifications = set(user_paragraph_classifications.denied.get(key, []))
@@ -496,9 +504,12 @@ class ResourceBrain:
         replace_field: bool = False,
         # cut to specific dimension if specified
         vector_dimension: Optional[int] = None,
+        append_splits: Optional[set[str]] = None,
     ):
         fid = ids.FieldId.from_string(f"{self.rid}/{field_id}")
         for subfield, vectors in vo.split_vectors.items():
+            if should_skip_split_indexing(subfield, replace_field, append_splits):
+                continue
             _field_id = ids.FieldId(
                 rid=fid.rid,
                 type=fid.type,
@@ -792,3 +803,10 @@ class ParagraphPages:
             if len(self._materialized) > 0:
                 return self._materialized[-1]
             return 0
+
+
+def should_skip_split_indexing(
+    split: str, replace_field: bool, append_splits: Optional[set[str]]
+) -> bool:
+    # When replacing the whole field, reindex all splits. Otherwise, we're only indexing the splits that are appended
+    return not replace_field and append_splits is not None and split not in append_splits