nucliadb 6.4.0.post4210__py3-none-any.whl → 6.4.0.post4213__py3-none-any.whl

nucliadb/ingest/app.py

@@ -103,7 +103,7 @@ async def initialize_pull_workers() -> list[Callable[[], Awaitable[None]]]:
     finalizers = await initialize_grpc()
     pull_workers = await consumer_service.start_pull_workers(SERVICE_NAME)
 
-    return [pull_workers] + finalizers
+    return pull_workers + finalizers
 
 
 async def main_consumer():  # pragma: no cover
@@ -117,7 +117,7 @@ async def main_consumer():  # pragma: no cover
     ingest_consumers = await consumer_service.start_ingest_consumers(SERVICE_NAME)
 
     await run_until_exit(
-        [grpc_health_finalizer, pull_workers, ingest_consumers, metrics_server.shutdown] + finalizers
+        [grpc_health_finalizer, ingest_consumers, metrics_server.shutdown] + pull_workers + finalizers
     )
 
 

nucliadb/ingest/consumer/pull.py

@@ -19,11 +19,14 @@
 #
 import asyncio
 import base64
+from datetime import datetime, timezone
 from typing import Optional
 
 from aiohttp.client_exceptions import ClientConnectorError
 
 from nucliadb.common import datamanagers
+from nucliadb.common.back_pressure.materializer import BackPressureMaterializer
+from nucliadb.common.back_pressure.utils import BackPressureException
 from nucliadb.common.http_clients.processing import ProcessingHTTPClient, get_nua_api_id
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.ingest import logger, logger_activity
@@ -57,6 +60,7 @@ class PullWorker:
         local_subscriber: bool = False,
         pull_time_empty_backoff: float = 5.0,
         pull_api_timeout: int = 60,
+        back_pressure: Optional[BackPressureMaterializer] = None,
     ):
         self.partition = partition
         self.pull_time_error_backoff = pull_time_error_backoff
@@ -65,6 +69,7 @@ class PullWorker:
         self.local_subscriber = local_subscriber
 
         self.processor = Processor(driver, storage, pubsub, partition)
+        self.back_pressure = back_pressure
 
     def __str__(self) -> str:
         return f"PullWorker(partition={self.partition})"
@@ -112,11 +117,29 @@ class PullWorker:
                 transaction_check=False,
             )
 
+    async def back_pressure_check(self) -> None:
+        if self.back_pressure is None:
+            return
+        while True:
+            try:
+                self.back_pressure.check_indexing()
+                self.back_pressure.check_ingest()
+                break
+            except BackPressureException as exc:
+                sleep_time = (datetime.now(timezone.utc) - exc.data.try_after).total_seconds()
+                logger.warning(f"Back pressure active! Sleeping for {sleep_time} seconds", exc_info=True)
+                await asyncio.sleep(sleep_time)
+            except Exception as e:
+                errors.capture_exception(e)
+                logger.exception("Error while checking back pressure. Moving on")
+                break
+
     async def loop(self):
         """
         Run this forever
         """
         while True:
+            await self.back_pressure_check()
             try:
                 await self._loop()
             except ReallyStopPulling:

nucliadb/ingest/consumer/service.py

@@ -22,18 +22,22 @@ 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.back_pressure.utils import is_back_pressure_enabled
 from nucliadb.common.maindb.utils import setup_driver
 from nucliadb.ingest import SERVICE_NAME, logger
 from nucliadb.ingest.consumer.consumer import IngestConsumer, IngestProcessedConsumer
 from nucliadb.ingest.consumer.pull import PullWorker
 from nucliadb.ingest.settings import settings
 from nucliadb_utils.exceptions import ConfigurationError
-from nucliadb_utils.settings import transaction_settings
+from nucliadb_utils.settings import indexing_settings, transaction_settings
 from nucliadb_utils.utilities import (
     get_audit,
     get_nats_manager,
     get_pubsub,
     get_storage,
+    start_nats_manager,
 )
 
 from .auditing import IndexAuditHandler, ResourceWritesAuditHandler
@@ -54,12 +58,38 @@ async def _exit_tasks(tasks: list[asyncio.Task]) -> None:
     await asyncio.gather(*tasks, return_exceptions=True)
 
 
+async def start_back_pressure() -> BackPressureMaterializer:
+    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_pull_workers(
     service_name: Optional[str] = None,
-) -> Callable[[], Awaitable[None]]:
+) -> list[Callable[[], Awaitable[None]]]:
+    finalizers: list[Callable[[], Awaitable[None]]] = []
+
     driver = await setup_driver()
     pubsub = await get_pubsub()
     storage = await get_storage(service_name=service_name or SERVICE_NAME)
+    back_pressure = None
+    if is_back_pressure_enabled():
+        back_pressure = await start_back_pressure()
+        finalizers.append(partial(stop_back_pressure, back_pressure))
     tasks = []
     for partition in settings.partitions:
         worker = PullWorker(
@@ -70,12 +100,14 @@ async def start_pull_workers(
             pubsub=pubsub,
             local_subscriber=transaction_settings.transaction_local,
             pull_api_timeout=settings.pull_api_timeout,
+            back_pressure=back_pressure,
         )
         task = asyncio.create_task(worker.loop())
         task.add_done_callback(_handle_task_result)
         tasks.append(task)
-
-    return partial(_exit_tasks, tasks)
+    if len(tasks):
+        finalizers.append(partial(_exit_tasks, tasks))
+    return finalizers
 
 
 async def start_ingest_consumers(

{nucliadb-6.4.0.post4210.dist-info → nucliadb-6.4.0.post4213.dist-info}/METADATA

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

{nucliadb-6.4.0.post4210.dist-info → nucliadb-6.4.0.post4213.dist-info}/RECORD

@@ -122,7 +122,7 @@ nucliadb/export_import/models.py,sha256=dbjScNkiMRv4X3Ktudy1JRliD25bfoDTy3JmEZgQ
 nucliadb/export_import/tasks.py,sha256=DWbdqY97ffoyfipelGXz3Jqz1iam6JCjQSh367Fc3NA,2947
 nucliadb/export_import/utils.py,sha256=8XOVMYXXw8b4ikojG7RjQ4tKN3Xu7nfu2yCUOqD50sk,23216
 nucliadb/ingest/__init__.py,sha256=fsw3C38VP50km3R-nHL775LNGPpJ4JxqXJ2Ib1f5SqE,1011
-nucliadb/ingest/app.py,sha256=KCptzFq1Msq4eHFxvEol4TFwSLdmkG2v1EfQ3C8PhyY,7547
+nucliadb/ingest/app.py,sha256=BKmjpdBEskHcRIHwOnI_jG4gFGs6dV0KKVH9MLJeA48,7546
 nucliadb/ingest/partitions.py,sha256=2NIhMYbNT0TNBL6bX1UMSi7vxFGICstCKEqsB0TXHOE,2410
 nucliadb/ingest/processing.py,sha256=QmkHq-BU4vub7JRWe9VHvQ2DcAmT6-CzgFXuZxXhcBU,20953
 nucliadb/ingest/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -134,8 +134,8 @@ nucliadb/ingest/consumer/auditing.py,sha256=xK21DIa_ZAiOJVVbnkmT4jgCRGshNGyPyxsq
 nucliadb/ingest/consumer/consumer.py,sha256=OgS1fr5Yo55u-XbC6zypTH1aJ562Y1vZHnPDlJJpCXQ,13703
 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=EYT0ImngMQgatStG68p2GSrPQBbJxeuq8nFm8DdAbwk,9280
-nucliadb/ingest/consumer/service.py,sha256=BLM_dmKZkFBsYl3sj4MZZp5M3kkxHLuO7sE18PqIatw,6538
+nucliadb/ingest/consumer/pull.py,sha256=vv1AyN0EhVgbgnZyT0D_1_IB4hWy7jPd4lAWPAOHGNc,10374
+nucliadb/ingest/consumer/service.py,sha256=mWzMQS1QkWmJNrkIahEZsn7jb8NbY9FRvPz89NeTT-4,7842
 nucliadb/ingest/consumer/shard_creator.py,sha256=w0smEu01FU_2cjZnsfBRNqT_Ntho11X17zTMST-vKbc,4359
 nucliadb/ingest/consumer/utils.py,sha256=jpX8D4lKzuPCpArQLZeX_Zczq3pfen_zAf8sPJfOEZU,2642
 nucliadb/ingest/fields/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -369,8 +369,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.post4210.dist-info/METADATA,sha256=SB9gIMgWxoWNtUEexRLH85E0PL-MnroGhJ6aOambTT4,4223
-nucliadb-6.4.0.post4210.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-nucliadb-6.4.0.post4210.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.0.post4210.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.0.post4210.dist-info/RECORD,,
+nucliadb-6.4.0.post4213.dist-info/METADATA,sha256=PAI_c9PMh-wJWIS4SmAAltmQcXStRUi6tKINdrNKJRM,4223
+nucliadb-6.4.0.post4213.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+nucliadb-6.4.0.post4213.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4213.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4213.dist-info/RECORD,,