nucliadb 6.4.0.post4279__py3-none-any.whl → 6.4.0.post4287__py3-none-any.whl

nucliadb/common/http_clients/processing.py

@@ -188,7 +188,7 @@ JSON_HEADERS = {"Content-Type": "application/json"}
 
 
 class ProcessingHTTPClient:
-    def __init__(self):
+    def __init__(self) -> None:
         self.session = aiohttp.ClientSession()
         self.base_url = get_processing_api_url()
         self.base_url_v2 = get_processing_api_v2_url()

nucliadb/export_import/importer.py

@@ -31,7 +31,7 @@ from nucliadb.export_import.utils import (
     ExportStreamReader,
     TaskRetryHandler,
     import_binary,
-    import_broker_message,
+    restore_broker_message,
     set_entities_groups,
     set_labels,
 )
@@ -64,7 +64,7 @@ async def import_kb(
         items_count += 1
         if item_type == ExportedItemType.RESOURCE:
             bm = cast(writer_pb2.BrokerMessage, data)
-            await import_broker_message(context, kbid, bm)
+            await restore_broker_message(context, kbid, bm)
 
         elif item_type == ExportedItemType.BINARY:
             cf = cast(resources_pb2.CloudFile, data[0])

nucliadb/export_import/utils.py

@@ -104,39 +104,32 @@ BM_FIELDS = {
 }
 
 
-async def import_broker_message(
-    context: ApplicationContext, kbid: str, bm: writer_pb2.BrokerMessage
-) -> None:
-    bm.kbid = kbid
-    partition = context.partitioning.generate_partition(kbid, bm.uuid)
-    for pb in [get_writer_bm(bm), get_processor_bm(bm)]:
-        await transaction_commit(context, pb, partition)
-
-
 async def restore_broker_message(
     context: ApplicationContext, kbid: str, bm: writer_pb2.BrokerMessage
 ) -> None:
+    bm.kbid = kbid
     await send_writer_bm(context, bm)
     await send_processor_bm(context, bm)
 
 
-@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
 async def send_writer_bm(context: ApplicationContext, bm: writer_pb2.BrokerMessage) -> None:
+    await process_bm_grpc(context, get_writer_bm(bm))
+
+
+async def send_processor_bm(context: ApplicationContext, bm: writer_pb2.BrokerMessage) -> None:
+    await process_bm_grpc(context, get_processor_bm(bm))
+
+
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
+async def process_bm_grpc(context: ApplicationContext, bm: writer_pb2.BrokerMessage) -> None:
     async def _iterator() -> AsyncIterator[writer_pb2.BrokerMessage]:
-        yield get_writer_bm(bm)
+        yield bm
 
     ingest_grpc: WriterStub = get_ingest()
     response: writer_pb2.OpStatusWriter = await ingest_grpc.ProcessMessage(_iterator())  # type: ignore
     assert response.status == writer_pb2.OpStatusWriter.Status.OK, "Failed to process broker message"
 
 
-async def send_processor_bm(context: ApplicationContext, bm: writer_pb2.BrokerMessage) -> None:
-    # Then enqueue the processor part asynchronously
-    processor_bm = get_processor_bm(bm)
-    partition = context.partitioning.generate_partition(bm.kbid, bm.uuid)
-    await transaction_commit(context, processor_bm, partition)
-
-
 async def transaction_commit(
     context: ApplicationContext, bm: writer_pb2.BrokerMessage, partition: int
 ) -> None:

nucliadb/ingest/consumer/pull.py

@@ -19,6 +19,7 @@
 #
 import asyncio
 import base64
+import time
 from contextlib import contextmanager
 from datetime import datetime, timezone
 from typing import Optional
@@ -41,6 +42,7 @@ from nucliadb.common.http_clients.processing import (
 )
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.ingest import SERVICE_NAME, logger, logger_activity
+from nucliadb.ingest.consumer.consumer import consumer_observer
 from nucliadb.ingest.orm.exceptions import ReallyStopPulling
 from nucliadb.ingest.orm.processor import Processor
 from nucliadb_protos.writer_pb2 import BrokerMessage, BrokerMessageBlobReference
@@ -52,7 +54,7 @@ from nucliadb_utils.cache.pubsub import PubSubDriver
 from nucliadb_utils.settings import nuclia_settings
 from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.transaction import MaxTransactionSizeExceededError
-from nucliadb_utils.utilities import get_storage, get_transaction_utility
+from nucliadb_utils.utilities import get_storage, get_transaction_utility, pull_subscriber_utilization
 
 processing_pending_messages = Gauge("nucliadb_processing_pending_messages")
 
@@ -309,11 +311,13 @@ class PullV2Worker:
 
         logger.debug(f"Resource: {pb.uuid} KB: {pb.kbid} ProcessingID: {pb.processing_id}")
 
-        await self.processor.process(
-            pb,
-            seq,
-            transaction_check=False,
-        )
+        source = "writer" if pb.source == pb.MessageSource.WRITER else "processor"
+        with consumer_observer({"source": source, "partition": "-1"}):
+            await self.processor.process(
+                pb,
+                seq,
+                transaction_check=False,
+            )
 
     async def loop(self):
         """
@@ -331,6 +335,7 @@ class PullV2Worker:
                 await asyncio.sleep(10)
 
     async def _loop(self):
+        usage_metric = pull_subscriber_utilization
         headers = {}
         data = None
         if nuclia_settings.nuclia_service_account is not None:
@@ -346,6 +351,8 @@ class PullV2Worker:
         async with ProcessingHTTPClient() as processing_http_client:
             while True:
                 try:
+                    start_time = time.monotonic()
+
                     # The code is only really prepared to pull 1 message at a time. If changing this, review MessageProgressUpdate usage
                     pull = await processing_http_client.pull_v2(ack_tokens=ack_tokens, limit=1)
                     ack_tokens.clear()
@@ -353,10 +360,14 @@ class PullV2Worker:
                         processing_pending_messages.set(0)
                         logger_activity.debug(f"No messages waiting in processing pull")
                         await asyncio.sleep(self.pull_time_empty_backoff)
+                        usage_metric.inc({"status": "waiting"}, time.monotonic() - start_time)
                         continue
 
-                    logger.info("Message received from proxy", extra={"seq": [pull.messages[0].seq]})
+                    received_time = time.monotonic()
+                    usage_metric.inc({"status": "waiting"}, received_time - start_time)
                     processing_pending_messages.set(pull.pending)
+
+                    logger.info("Message received from proxy", extra={"seq": [pull.messages[0].seq]})
                     try:
                         for message in pull.messages:
                             async with ProcessingPullMessageProgressUpdater(
@@ -365,6 +376,8 @@ class PullV2Worker:
                                 with run_in_span(message.headers):
                                     await self.handle_message(message.seq, message.payload)
                                     ack_tokens.append(message.ack_token)
+
+                        usage_metric.inc({"status": "processing"}, time.monotonic() - received_time)
                     except Exception as e:
                         errors.capture_exception(e)
                         logger.exception("Error while pulling and processing message/s")

{nucliadb-6.4.0.post4279.dist-info → nucliadb-6.4.0.post4287.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4279
+Version: 6.4.0.post4287
 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.post4279
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4279
-Requires-Dist: nucliadb-protos>=6.4.0.post4279
-Requires-Dist: nucliadb-models>=6.4.0.post4279
-Requires-Dist: nidx-protos>=6.4.0.post4279
+Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4287
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4287
+Requires-Dist: nucliadb-protos>=6.4.0.post4287
+Requires-Dist: nucliadb-models>=6.4.0.post4287
+Requires-Dist: nidx-protos>=6.4.0.post4287
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.4.0.post4279.dist-info → nucliadb-6.4.0.post4287.dist-info}/RECORD

@@ -101,7 +101,7 @@ nucliadb/common/external_index_providers/settings.py,sha256=EGHnIkwxqe6aypwKegXT
 nucliadb/common/http_clients/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/common/http_clients/auth.py,sha256=srfpgAbs2wmqA9u_l-HxsV4YoO77Tse4y3gm3q2YvYM,2112
 nucliadb/common/http_clients/exceptions.py,sha256=47Y8OjkaGV_F18G07FpJhOzgWKUIexhlILyuVtICz8s,1100
-nucliadb/common/http_clients/processing.py,sha256=VzxzFArNsHWGmFoX0c5OrQB3vFW841aeyuP5NgzPQGo,9581
+nucliadb/common/http_clients/processing.py,sha256=crLfKo_2RJr9Uo2vuq11MWFa9tV2njA_v7ZBd95tjNU,9589
 nucliadb/common/http_clients/pypi.py,sha256=VHIUjwJEJVntVUo_FRoXIo8sLmluy7sa9-iXSITcrMY,1540
 nucliadb/common/http_clients/utils.py,sha256=yGUkHNS41abHiBoHqo_Mg3QSqGsS7rUtbfGftbEC57U,1529
 nucliadb/common/maindb/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -117,10 +117,10 @@ nucliadb/export_import/__init__.py,sha256=y-Is0Bxa8TMV6UiOW0deC_D3U465P65CQ5RjBj
 nucliadb/export_import/datamanager.py,sha256=xL8b0xvk45q6wx1l7J32JgPmpyjnF6fKiJi2F2B_UQY,6968
 nucliadb/export_import/exceptions.py,sha256=Dw8WqfG4r6MPJc5TPfbjMvCgXXWTcTOecGHRVU1h3kM,1949
 nucliadb/export_import/exporter.py,sha256=k2QVx1EjqFlDYiggriWiEJzwtMXzHbldsqWdpGQM3_U,7074
-nucliadb/export_import/importer.py,sha256=v5cq9Nn8c2zrY_K_00mydR52f8mdFxR7tLdtNLQ0qvk,4229
+nucliadb/export_import/importer.py,sha256=GNDMt4hdjbcLWdydVq8XFQKefzNJkQ1eTzhshUX64rk,4231
 nucliadb/export_import/models.py,sha256=dbjScNkiMRv4X3Ktudy1JRliD25bfoDTy3JmEZgQSCc,2121
 nucliadb/export_import/tasks.py,sha256=DWbdqY97ffoyfipelGXz3Jqz1iam6JCjQSh367Fc3NA,2947
-nucliadb/export_import/utils.py,sha256=8XOVMYXXw8b4ikojG7RjQ4tKN3Xu7nfu2yCUOqD50sk,23216
+nucliadb/export_import/utils.py,sha256=iutS86YblS8aLQ9PCZUyTJMN6lDV4DjcjaptQVBfBNA,22874
 nucliadb/ingest/__init__.py,sha256=fsw3C38VP50km3R-nHL775LNGPpJ4JxqXJ2Ib1f5SqE,1011
 nucliadb/ingest/app.py,sha256=Eympy8nbz09VDNPF28MuIeKMb7wgB9cTSOObS8uvL0o,8372
 nucliadb/ingest/partitions.py,sha256=2NIhMYbNT0TNBL6bX1UMSi7vxFGICstCKEqsB0TXHOE,2410
@@ -134,7 +134,7 @@ nucliadb/ingest/consumer/auditing.py,sha256=xK21DIa_ZAiOJVVbnkmT4jgCRGshNGyPyxsq
 nucliadb/ingest/consumer/consumer.py,sha256=GfdlrNlnt7PWYyk75xtyzn2SHZse7475U4U9q_9jKr0,13711
 nucliadb/ingest/consumer/materializer.py,sha256=tgD_rDI2twQzcz8kKNiW_L4YIth16IGh9mUfD5wiSD4,3858
 nucliadb/ingest/consumer/metrics.py,sha256=ji1l_4cKiHJthQd8YNem1ft4iMbw9KThmVvJmLcv3Xg,1075
-nucliadb/ingest/consumer/pull.py,sha256=gfdyQ8IMFA_bpGnEpmRB9qmOJywBEwxn7pGYaueZizU,16874
+nucliadb/ingest/consumer/pull.py,sha256=x39G6AcNXSnw_GRPxJfafmD5pehZzMBd6v_f_yrNbUI,17594
 nucliadb/ingest/consumer/service.py,sha256=WXBN8dY7MlmYWxqQHIbIO7w_SdVJRY1RuHAWlQUXf8o,8852
 nucliadb/ingest/consumer/shard_creator.py,sha256=w0smEu01FU_2cjZnsfBRNqT_Ntho11X17zTMST-vKbc,4359
 nucliadb/ingest/consumer/utils.py,sha256=jpX8D4lKzuPCpArQLZeX_Zczq3pfen_zAf8sPJfOEZU,2642
@@ -368,8 +368,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.4.0.post4279.dist-info/METADATA,sha256=ISm2mlidMfyHGlEXRXBJcbnPe52rP58sBjj5NLRFf68,4223
-nucliadb-6.4.0.post4279.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-nucliadb-6.4.0.post4279.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.0.post4279.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.0.post4279.dist-info/RECORD,,
+nucliadb-6.4.0.post4287.dist-info/METADATA,sha256=5QNnOEzt2bbjxt2hVf853EWspV23p_WQLeNr7FL9IQk,4223
+nucliadb-6.4.0.post4287.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+nucliadb-6.4.0.post4287.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4287.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4287.dist-info/RECORD,,