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

nucliadb/common/back_pressure/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+from .materializer import maybe_back_pressure, start_materializer, stop_materializer  # noqa

nucliadb/common/back_pressure/cache.py

@@ -0,0 +1,86 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 contextlib
+import logging
+import threading
+from datetime import datetime, timezone
+from typing import Optional
+
+from cachetools import TTLCache
+
+from nucliadb.common.back_pressure.utils import BackPressureData, BackPressureException
+from nucliadb_telemetry import metrics
+
+logger = logging.getLogger(__name__)
+
+
+RATE_LIMITED_REQUESTS_COUNTER = metrics.Counter(
+    "nucliadb_rate_limited_requests", labels={"type": "", "cached": ""}
+)
+
+
+class BackPressureCache:
+    """
+    Global cache for storing already computed try again in times.
+    It allows us to avoid making the same calculations multiple
+    times if back pressure has been applied.
+    """
+
+    def __init__(self):
+        self._cache = TTLCache(maxsize=1024, ttl=5 * 60)
+        self._lock = threading.Lock()
+
+    def get(self, key: str) -> Optional[BackPressureData]:
+        with self._lock:
+            data = self._cache.get(key, None)
+            if data is None:
+                return None
+            if datetime.now(timezone.utc) >= data.try_after:
+                # The key has expired, so remove it from the cache
+                self._cache.pop(key, None)
+                return None
+            return data
+
+    def set(self, key: str, data: BackPressureData):
+        with self._lock:
+            self._cache[key] = data
+
+
+_cache = BackPressureCache()
+
+
+@contextlib.contextmanager
+def cached_back_pressure(cache_key: str):
+    """
+    Context manager that handles the caching of the try again in time so that
+    we don't recompute try again times if we have already applied back pressure.
+    """
+    data: Optional[BackPressureData] = _cache.get(cache_key)
+    if data is not None:
+        back_pressure_type = data.type
+        RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "true"})
+        raise BackPressureException(data)
+    try:
+        yield
+    except BackPressureException as exc:
+        back_pressure_type = exc.data.type
+        RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "false"})
+        _cache.set(cache_key, exc.data)
+        raise exc

nucliadb/common/back_pressure/materializer.py

@@ -0,0 +1,315 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 logging
+import threading
+from typing import Optional
+
+from cachetools import TTLCache
+from fastapi import HTTPException
+
+from nucliadb.common.back_pressure.cache import cached_back_pressure
+from nucliadb.common.back_pressure.settings import settings
+from nucliadb.common.back_pressure.utils import (
+    BackPressureData,
+    BackPressureException,
+    estimate_try_after,
+    get_nats_consumer_pending_messages,
+    is_back_pressure_enabled,
+)
+from nucliadb.common.context import ApplicationContext
+from nucliadb.common.http_clients.processing import ProcessingHTTPClient
+from nucliadb_telemetry import metrics
+from nucliadb_utils import const
+from nucliadb_utils.nats import NatsConnectionManager
+from nucliadb_utils.settings import is_onprem_nucliadb
+
+logger = logging.getLogger(__name__)
+
+
+back_pressure_observer = metrics.Observer("nucliadb_back_pressure", labels={"type": ""})
+
+
+class BackPressureMaterializer:
+    """
+    Singleton class that will run in the background gathering the different
+    stats to apply back pressure and materializing it in memory. This allows us
+    to do stale-reads when checking if back pressure is needed for a particular
+    request - thus not slowing it down.
+    """
+
+    def __init__(
+        self,
+        nats_manager: NatsConnectionManager,
+        indexing_check_interval: int = 30,
+        ingest_check_interval: int = 30,
+    ):
+        self.nats_manager = nats_manager
+        self.processing_http_client = ProcessingHTTPClient()
+
+        self.indexing_check_interval = indexing_check_interval
+        self.ingest_check_interval = ingest_check_interval
+
+        self.ingest_pending: int = 0
+        self.indexing_pending: int = 0
+
+        self._tasks: list[asyncio.Task] = []
+        self._running = False
+
+        self.processing_pending_cache = TTLCache(maxsize=1024, ttl=60)  # type: ignore
+        self.processing_pending_locks: dict[str, asyncio.Lock] = {}
+
+    async def start(self):
+        self._tasks.append(asyncio.create_task(self._get_indexing_pending_task()))
+        self._tasks.append(asyncio.create_task(self._get_ingest_pending_task()))
+        self._running = True
+
+    async def stop(self):
+        for task in self._tasks:
+            task.cancel()
+        self._tasks.clear()
+        await self.processing_http_client.close()
+        self._running = False
+
+    @property
+    def running(self) -> bool:
+        return self._running
+
+    async def get_processing_pending(self, kbid: str) -> int:
+        """
+        We don't materialize the pending messages for every kbid, but values are cached for some time.
+        """
+        cached = self.processing_pending_cache.get(kbid)
+        if cached is not None:
+            return cached
+
+        lock = self.processing_pending_locks.setdefault(kbid, asyncio.Lock())
+        async with lock:
+            # Check again if the value has been cached while we were waiting for the lock
+            cached = self.processing_pending_cache.get(kbid)
+            if cached is not None:
+                return cached
+
+            # Get the pending messages and cache the result
+            try:
+                with back_pressure_observer({"type": "get_processing_pending"}):
+                    pending = await self._get_processing_pending(kbid)
+            except Exception:  # pragma: no cover
+                # Do not cache if there was an error
+                logger.exception(
+                    "Error getting pending messages to process. Back pressure on proccessing for KB can't be applied.",
+                    exc_info=True,
+                    extra={"kbid": kbid},
+                )
+                return 0
+
+            if pending > 0:
+                logger.info(
+                    f"Processing returned {pending} pending messages for KB",
+                    extra={"kbid": kbid},
+                )
+            self.processing_pending_cache[kbid] = pending
+            return pending
+
+    async def _get_processing_pending(self, kbid: str) -> int:
+        response = await self.processing_http_client.stats(kbid=kbid, timeout=0.5)
+        return response.incomplete
+
+    def get_indexing_pending(self) -> int:
+        return self.indexing_pending
+
+    def get_ingest_pending(self) -> int:
+        return self.ingest_pending
+
+    async def _get_indexing_pending_task(self):
+        try:
+            while True:
+                try:
+                    with back_pressure_observer({"type": "get_indexing_pending"}):
+                        self.indexing_pending = await get_nats_consumer_pending_messages(
+                            self.nats_manager,
+                            stream="nidx",
+                            consumer="nidx",
+                        )
+                except Exception:  # pragma: no cover
+                    logger.exception(
+                        "Error getting pending messages to index",
+                        exc_info=True,
+                    )
+                await asyncio.sleep(self.indexing_check_interval)
+        except asyncio.CancelledError:
+            pass
+
+    async def _get_ingest_pending_task(self):
+        try:
+            while True:
+                try:
+                    with back_pressure_observer({"type": "get_ingest_pending"}):
+                        self.ingest_pending = await get_nats_consumer_pending_messages(
+                            self.nats_manager,
+                            stream=const.Streams.INGEST_PROCESSED.name,
+                            consumer=const.Streams.INGEST_PROCESSED.group,
+                        )
+                except Exception:  # pragma: no cover
+                    logger.exception(
+                        "Error getting pending messages to ingest",
+                        exc_info=True,
+                    )
+                await asyncio.sleep(self.ingest_check_interval)
+        except asyncio.CancelledError:
+            pass
+
+    def check_indexing(self):
+        max_pending = settings.max_indexing_pending
+        if max_pending <= 0:
+            # Indexing back pressure is disabled
+            return
+        pending = self.get_indexing_pending()
+        if pending > max_pending:
+            try_after = estimate_try_after(
+                rate=settings.indexing_rate,
+                pending=pending,
+                max_wait=settings.max_wait_time,
+            )
+            data = BackPressureData(type="indexing", try_after=try_after)
+            raise BackPressureException(data)
+
+    def check_ingest(self):
+        max_pending = settings.max_ingest_pending
+        if max_pending <= 0:
+            # Ingest back pressure is disabled
+            return
+        ingest_pending = self.get_ingest_pending()
+        if ingest_pending > max_pending:
+            try_after = estimate_try_after(
+                rate=settings.ingest_rate,
+                pending=ingest_pending,
+                max_wait=settings.max_wait_time,
+            )
+            data = BackPressureData(type="ingest", try_after=try_after)
+            raise BackPressureException(data)
+
+    async def check_processing(self, kbid: str):
+        max_pending = settings.max_processing_pending
+        if max_pending <= 0:
+            # Processing back pressure is disabled
+            return
+
+        kb_pending = await self.get_processing_pending(kbid)
+        if kb_pending > max_pending:
+            try_after = estimate_try_after(
+                rate=settings.processing_rate,
+                pending=kb_pending,
+                max_wait=settings.max_wait_time,
+            )
+            data = BackPressureData(type="processing", try_after=try_after)
+            raise BackPressureException(data)
+
+
+MATERIALIZER: Optional[BackPressureMaterializer] = None
+materializer_lock = threading.Lock()
+
+
+async def start_materializer(context: ApplicationContext):
+    global MATERIALIZER
+    if MATERIALIZER is not None:
+        logger.warning("BackPressureMaterializer already started")
+        return
+    with materializer_lock:
+        if MATERIALIZER is not None:
+            return
+        logger.info("Initializing materializer")
+        try:
+            nats_manager = context.nats_manager
+        except AttributeError:
+            logger.warning(
+                "Could not initialize materializer. Nats manager not found or not initialized yet"
+            )
+            return
+        materializer = BackPressureMaterializer(
+            nats_manager,
+            indexing_check_interval=settings.indexing_check_interval,
+            ingest_check_interval=settings.ingest_check_interval,
+        )
+        await materializer.start()
+        MATERIALIZER = materializer
+
+
+async def stop_materializer():
+    global MATERIALIZER
+    if MATERIALIZER is None or not MATERIALIZER.running:
+        logger.warning("BackPressureMaterializer already stopped")
+        return
+    with materializer_lock:
+        if MATERIALIZER is None:
+            return
+        logger.info("Stopping materializer")
+        await MATERIALIZER.stop()
+        MATERIALIZER = None
+
+
+def get_materializer() -> BackPressureMaterializer:
+    global MATERIALIZER
+    if MATERIALIZER is None:
+        raise RuntimeError("BackPressureMaterializer not initialized")
+    return MATERIALIZER
+
+
+async def maybe_back_pressure(kbid: str, resource_uuid: Optional[str] = None) -> None:
+    """
+    This function does system checks to see if we need to put back pressure on writes.
+    In that case, a HTTP 429 will be raised with the estimated time to try again.
+    """
+    if not is_back_pressure_enabled() or is_onprem_nucliadb():
+        return
+    await back_pressure_checks(kbid, resource_uuid)
+
+
+async def back_pressure_checks(kbid: str, resource_uuid: Optional[str] = None):
+    """
+    Will raise a 429 if back pressure is needed:
+    - If the processing engine is behind.
+    - If ingest processed consumer is behind.
+    - If the indexing on nodes affected by the request (kbid, and resource_uuid) is behind.
+    """
+    materializer = get_materializer()
+    try:
+        with cached_back_pressure(f"{kbid}-{resource_uuid}"):
+            materializer.check_indexing()
+            materializer.check_ingest()
+            await materializer.check_processing(kbid)
+    except BackPressureException as exc:
+        logger.info(
+            "Back pressure applied",
+            extra={
+                "kbid": kbid,
+                "resource_uuid": resource_uuid,
+                "try_after": exc.data.try_after,
+                "back_pressure_type": exc.data.type,
+            },
+        )
+        raise HTTPException(
+            status_code=429,
+            detail={
+                "message": f"Too many messages pending to ingest. Retry after {exc.data.try_after}",
+                "try_after": exc.data.try_after.timestamp(),
+                "back_pressure_type": exc.data.type,
+            },
+        ) from exc

nucliadb/common/back_pressure/settings.py

@@ -0,0 +1,72 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+
+class BackPressureSettings(BaseSettings):
+    enabled: bool = Field(
+        default=False,
+        description="Enable or disable back pressure.",
+        alias="back_pressure_enabled",
+    )
+    indexing_rate: float = Field(
+        default=10,
+        description="Estimation of the indexing rate in messages per second. This is used to calculate the try again in time",  # noqa
+    )
+    ingest_rate: float = Field(
+        default=4,
+        description="Estimation of the ingest processed consumer rate in messages per second. This is used to calculate the try again in time",  # noqa
+    )
+    processing_rate: float = Field(
+        default=1,
+        description="Estimation of the processing rate in messages per second. This is used to calculate the try again in time",  # noqa
+    )
+    max_indexing_pending: int = Field(
+        default=1000,
+        description="Max number of messages pending to index in a node queue before rate limiting writes. Set to 0 to disable indexing back pressure checks",  # noqa
+        alias="back_pressure_max_indexing_pending",
+    )
+    max_ingest_pending: int = Field(
+        # Disabled by default
+        default=0,
+        description="Max number of messages pending to be ingested by processed consumers before rate limiting writes. Set to 0 to disable ingest back pressure checks",  # noqa
+        alias="back_pressure_max_ingest_pending",
+    )
+    max_processing_pending: int = Field(
+        default=1000,
+        description="Max number of messages pending to process per Knowledge Box before rate limiting writes. Set to 0 to disable processing back pressure checks",  # noqa
+        alias="back_pressure_max_processing_pending",
+    )
+    indexing_check_interval: int = Field(
+        default=30,
+        description="Interval in seconds to check the indexing pending messages",
+    )
+    ingest_check_interval: int = Field(
+        default=30,
+        description="Interval in seconds to check the ingest pending messages",
+    )
+    max_wait_time: int = Field(
+        default=60,
+        description="Max time in seconds to wait before trying again after back pressure",
+    )
+
+
+settings = BackPressureSettings()

nucliadb/common/back_pressure/utils.py

@@ -0,0 +1,59 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+
+from nucliadb.common.back_pressure.settings import settings
+from nucliadb_utils.nats import NatsConnectionManager
+
+
+@dataclass
+class BackPressureData:
+    type: str
+    try_after: datetime
+
+
+class BackPressureException(Exception):
+    def __init__(self, data: BackPressureData):
+        self.data = data
+
+    def __str__(self):
+        return f"Back pressure applied for {self.data.type}. Try again after {self.data.try_after}"
+
+
+def is_back_pressure_enabled() -> bool:
+    return settings.enabled
+
+
+def estimate_try_after(rate: float, pending: int, max_wait: int) -> datetime:
+    """
+    This function estimates the time to try again based on the rate and the number of pending messages.
+    """
+    delta_seconds = min(pending / rate, max_wait)
+    return datetime.now(timezone.utc) + timedelta(seconds=delta_seconds)
+
+
+async def get_nats_consumer_pending_messages(
+    nats_manager: NatsConnectionManager, *, stream: str, consumer: str
+) -> int:
+    # get raw js client
+    js = nats_manager.js
+    consumer_info = await js.consumer_info(stream, consumer)
+    return consumer_info.num_pending

nucliadb/common/external_index_providers/base.py

@@ -55,16 +55,19 @@ class VectorsetExternalIndex:
     similarity: VectorSimilarity.ValueType
 
 
-class TextBlockMatch(BaseModel):
+class ScoredTextBlock(BaseModel):
+    paragraph_id: ParagraphId
+    score: float
+    score_type: SCORE_TYPE
+
+
+class TextBlockMatch(ScoredTextBlock):
     """
     Model a text block/paragraph retrieved from an external index with all the information
     needed in order to later hydrate retrieval results.
     """
 
-    paragraph_id: ParagraphId
     position: TextPosition
-    score: float
-    score_type: SCORE_TYPE
     order: int
     page_with_visual: bool = False
     fuzzy_search: bool

nucliadb/search/search/chat/ask.py

@@ -33,6 +33,8 @@ from nuclia_models.predict.generative_responses import (
 from pydantic_core import ValidationError
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
+from nucliadb.common.external_index_providers.base import ScoredTextBlock
+from nucliadb.common.ids import ParagraphId
 from nucliadb.models.responses import HTTPClientError
 from nucliadb.search import logger, predict
 from nucliadb.search.predict import (
@@ -63,6 +65,7 @@ from nucliadb.search.search.graph_strategy import get_graph_results
 from nucliadb.search.search.metrics import RAGMetrics
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.parsers.ask import fetcher_for_ask, parse_ask
+from nucliadb.search.search.rank_fusion import WeightedCombSum
 from nucliadb.search.search.rerankers import (
     get_reranker,
 )
@@ -865,6 +868,10 @@ async def retrieval_in_resource(
     )
 
 
+class _FindParagraph(ScoredTextBlock):
+    original: FindParagraph
+
+
 def compute_best_matches(
     main_results: KnowledgeboxFindResults,
     prequeries_results: Optional[list[PreQueryResult]] = None,
@@ -882,42 +889,46 @@ def compute_best_matches(
     `main_query_weight` is the weight given to the paragraphs matching the main query when calculating the final score.
     """
 
-    def iter_paragraphs(results: KnowledgeboxFindResults):
+    def extract_paragraphs(results: KnowledgeboxFindResults) -> list[_FindParagraph]:
+        paragraphs = []
         for resource in results.resources.values():
             for field in resource.fields.values():
                 for paragraph in field.paragraphs.values():
-                    yield paragraph
-
-    total_weights = main_query_weight + sum(prequery.weight for prequery, _ in prequeries_results or [])
-    paragraph_id_to_match: dict[str, RetrievalMatch] = {}
-    for paragraph in iter_paragraphs(main_results):
-        normalized_weight = main_query_weight / total_weights
-        rmatch = RetrievalMatch(
-            paragraph=paragraph,
-            weighted_score=paragraph.score * normalized_weight,
-        )
-        paragraph_id_to_match[paragraph.id] = rmatch
-
-    for prequery, prequery_results in prequeries_results or []:
-        for paragraph in iter_paragraphs(prequery_results):
-            normalized_weight = prequery.weight / total_weights
-            weighted_score = paragraph.score * normalized_weight
-            if paragraph.id in paragraph_id_to_match:
-                rmatch = paragraph_id_to_match[paragraph.id]
-                # If a paragraph is matched in various prequeries, the final score is the
-                # sum of the weighted scores
-                rmatch.weighted_score += weighted_score
-            else:
-                paragraph_id_to_match[paragraph.id] = RetrievalMatch(
-                    paragraph=paragraph,
-                    weighted_score=weighted_score,
-                )
+                    paragraphs.append(
+                        _FindParagraph(
+                            paragraph_id=ParagraphId.from_string(paragraph.id),
+                            score=paragraph.score,
+                            score_type=paragraph.score_type,
+                            original=paragraph,
+                        )
+                    )
+        return paragraphs
 
-    return sorted(
-        paragraph_id_to_match.values(),
-        key=lambda match: match.weighted_score,
-        reverse=True,
-    )
+    weights = {
+        "main": main_query_weight,
+    }
+    total_weight = main_query_weight
+    find_results = {
+        "main": extract_paragraphs(main_results),
+    }
+    for i, (prequery, prequery_results) in enumerate(prequeries_results or []):
+        weights[f"prequery-{i}"] = prequery.weight
+        total_weight += prequery.weight
+        find_results[f"prequery-{i}"] = extract_paragraphs(prequery_results)
+
+    normalized_weights = {key: value / total_weight for key, value in weights.items()}
+
+    # window does nothing here
+    rank_fusion = WeightedCombSum(window=0, weights=normalized_weights)
+
+    merged = []
+    for item in rank_fusion.fuse(find_results):
+        match = RetrievalMatch(
+            paragraph=item.original,
+            weighted_score=item.score,
+        )
+        merged.append(match)
+    return merged
 
 
 def calculate_prequeries_for_json_schema(