nucliadb 6.4.0.post4192__py3-none-any.whl → 6.4.0.post4196__py3-none-any.whl

nucliadb/common/cache.py

@@ -18,19 +18,28 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-import asyncio
 import contextlib
+import logging
 from abc import ABC, abstractmethod
 from contextvars import ContextVar
 from dataclasses import dataclass
 from functools import cached_property
 from typing import Generic, Optional, TypeVar
 
-from lru import LRU
+import backoff
+from async_lru import _LRUCacheWrapper, alru_cache
+from typing_extensions import ParamSpec
 
+from nucliadb.common.ids import FieldId
+from nucliadb.common.maindb.utils import get_driver
+from nucliadb.ingest.fields.base import FieldTypes
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
 from nucliadb_protos.utils_pb2 import ExtractedText
 from nucliadb_telemetry.metrics import Counter, Gauge
+from nucliadb_utils.utilities import get_storage
+
+logger = logging.getLogger(__name__)
 
 # specific metrics per cache type
 cached_resources = Gauge("nucliadb_cached_resources")
@@ -39,6 +48,7 @@ resource_cache_ops = Counter("nucliadb_resource_cache_ops", labels={"type": ""})
 extracted_text_cache_ops = Counter("nucliadb_extracted_text_cache_ops", labels={"type": ""})
 
 
+K = ParamSpec("K")
 T = TypeVar("T")
 
 
@@ -48,7 +58,7 @@ class CacheMetrics:
     ops: Counter
 
 
-class Cache(Generic[T], ABC):
+class Cache(Generic[K, T], ABC):
     """Low-level bounded cache implementation with access to per-key async locks
     in case cache users want to lock concurrent access.
 
@@ -56,52 +66,43 @@ class Cache(Generic[T], ABC):
 
     """
 
-    def __init__(self, cache_size: int) -> None:
-        self.cache: LRU[str, T] = LRU(cache_size, callback=self._evicted_callback)
-        self.locks: dict[str, asyncio.Lock] = {}
-
-    def _evicted_callback(self, key: str, value: T):
-        self.locks.pop(key, None)
-        self.metrics.ops.inc({"type": "evict"})
-
-    def get(self, key: str) -> Optional[T]:
-        return self.cache.get(key)
-
-    # Get a lock for a specific key. Locks will be evicted at the same time as
-    # key-value pairs
-    def get_lock(self, key: str) -> asyncio.Lock:
-        return self.locks.setdefault(key, asyncio.Lock())
-
-    def set(self, key: str, value: T):
-        len_before = len(self.cache)
-
-        self.cache[key] = value
+    cache: _LRUCacheWrapper[Optional[T]]
 
-        len_after = len(self.cache)
-        if len_after - len_before > 0:
-            self.metrics._cache_size.inc(len_after - len_before)
+    async def get(self, *args: K.args, **kwargs: K.kwargs) -> Optional[T]:
+        result = await self.cache(*args)
+        # Do not cache None
+        if result is None:
+            self.cache.cache_invalidate(*args)
+        return result
 
-    def __contains__(self, key: str) -> bool:
-        return self.cache.__contains__(key)
-
-    def clear(self):
-        self.metrics._cache_size.dec(len(self.cache))
-        self.cache.clear()
-        self.locks.clear()
+    def finalize(self):
+        info = self.cache.cache_info()
+        self.metrics.ops.inc({"type": "miss"}, value=info.misses)
+        self.metrics.ops.inc({"type": "hit"}, value=info.hits)
 
     @abstractmethod
     @cached_property
     def metrics(self) -> CacheMetrics: ...
 
 
-class ResourceCache(Cache[ResourceORM]):
+class ResourceCache(Cache[[str, str], ResourceORM]):
+    def __init__(self, cache_size: int) -> None:
+        @alru_cache(maxsize=cache_size)
+        async def _get_resource(kbid: str, rid: str) -> Optional[ResourceORM]:
+            storage = await get_storage()
+            async with get_driver().transaction(read_only=True) as txn:
+                kb = KnowledgeBoxORM(txn, storage, kbid)
+                return await kb.get(rid)
+
+        self.cache = _get_resource
+
     metrics = CacheMetrics(
         _cache_size=cached_resources,
         ops=resource_cache_ops,
     )
 
 
-class ExtractedTextCache(Cache[ExtractedText]):
+class ExtractedTextCache(Cache[[str, FieldId], ExtractedText]):
     """
     Used to cache extracted text from a resource in memory during the process
     of search results hydration.
@@ -111,6 +112,30 @@ class ExtractedTextCache(Cache[ExtractedText]):
     fetched for each field where the text block is found.
     """
 
+    def __init__(self, cache_size: int) -> None:
+        @alru_cache(maxsize=cache_size)
+        @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
+        async def _get_extracted_text(kbid: str, field_id: FieldId) -> Optional[ExtractedText]:
+            storage = await get_storage()
+            try:
+                sf = storage.file_extracted(
+                    kbid, field_id.rid, field_id.type, field_id.key, FieldTypes.FIELD_TEXT.value
+                )
+                return await storage.download_pb(sf, ExtractedText)
+            except Exception:
+                logger.warning(
+                    "Error getting extracted text for field. Retrying",
+                    exc_info=True,
+                    extra={
+                        "kbid": kbid,
+                        "resource_id": field_id.rid,
+                        "field": f"{field_id.type}/{field_id.key}",
+                    },
+                )
+                raise
+
+        self.cache = _get_extracted_text
+
     metrics = CacheMetrics(
         _cache_size=cached_extracted_texts,
         ops=extracted_text_cache_ops,
@@ -154,7 +179,7 @@ def _use_cache(klass: type[Cache], context_var: ContextVar, /, **kwargs):
         yield cache
     finally:
         context_var.reset(token)
-        cache.clear()
+        cache.finalize()
 
 
 @contextlib.contextmanager

nucliadb/learning_proxy.py

@@ -30,7 +30,6 @@ import backoff
 import httpx
 from fastapi import Request, Response
 from fastapi.responses import StreamingResponse
-from lru import LRU
 from pydantic import BaseModel, Field, model_validator
 from typing_extensions import Self
 
@@ -480,8 +479,7 @@ class ProxiedLearningConfig(LearningConfigService):
             yield client
 
 
-_IN_MEMORY_CONFIGS: dict[str, LearningConfiguration]
-_IN_MEMORY_CONFIGS = LRU(50)  # type: ignore
+_IN_MEMORY_CONFIGS: dict[str, LearningConfiguration] = {}
 
 
 class InMemoryLearningConfig(LearningConfigService):

nucliadb/search/search/cache.py

@@ -45,26 +45,12 @@ async def get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
     """
     Will try to get the resource from the cache, if it's not there it will fetch it from the ORM and cache it.
     """
-    orm_resource: Optional[ResourceORM] = None
-
     resource_cache = get_resource_cache()
     if resource_cache is None:
         logger.warning("Resource cache not set")
         return await _orm_get_resource(kbid, uuid)
 
-    async with resource_cache.get_lock(uuid):
-        if uuid not in resource_cache:
-            resource_cache.metrics.ops.inc({"type": "miss"})
-            orm_resource = await _orm_get_resource(kbid, uuid)
-        else:
-            resource_cache.metrics.ops.inc({"type": "hit"})
-
-        if orm_resource is not None:
-            resource_cache.set(uuid, orm_resource)
-        else:
-            orm_resource = resource_cache.get(uuid)
-
-    return orm_resource
+    return await resource_cache.get(kbid, uuid)
 
 
 async def _orm_get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
@@ -75,30 +61,17 @@ async def _orm_get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
 
 
 async def get_field_extracted_text(field: Field) -> Optional[ExtractedText]:
+    if field.extracted_text is not None:
+        return field.extracted_text
+
     cache = get_extracted_text_cache()
     if cache is None:
         logger.warning("Extracted text cache not set")
         return await field.get_extracted_text()
 
-    key = f"{field.kbid}/{field.uuid}/{field.id}"
-    extracted_text = cache.get(key)
-    if extracted_text is not None:
-        cache.metrics.ops.inc({"type": "hit"})
-        return extracted_text
-
-    async with cache.get_lock(key):
-        # Check again in case another task already fetched it
-        extracted_text = cache.get(key)
-        if extracted_text is not None:
-            cache.metrics.ops.inc({"type": "hit"})
-            return extracted_text
-
-        cache.metrics.ops.inc({"type": "miss"})
-        extracted_text = await field_get_extracted_text(field)
-        if extracted_text is not None:
-            # Only cache if we actually have extracted text
-            cache.set(key, extracted_text)
-        return extracted_text
+    extracted_text = await cache.get(field.kbid, FieldId(field.uuid, field.type, field.id))
+    field.extracted_text = extracted_text
+    return extracted_text
 
 
 @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)

nucliadb/train/generators/utils.py

@@ -32,20 +32,11 @@ from nucliadb_utils.utilities import get_storage
 
 async def get_resource_from_cache_or_db(kbid: str, uuid: str) -> Optional[ResourceORM]:
     resource_cache = get_resource_cache()
-    orm_resource: Optional[ResourceORM] = None
     if resource_cache is None:
         return await _get_resource_from_db(kbid, uuid)
         logger.warning("Resource cache is not set")
 
-    if uuid not in resource_cache:
-        resource_cache.metrics.ops.inc({"type": "miss"})
-        orm_resource = await _get_resource_from_db(kbid, uuid)
-        if orm_resource is not None:
-            resource_cache.set(uuid, orm_resource)
-    else:
-        resource_cache.metrics.ops.inc({"type": "hit"})
-        orm_resource = resource_cache.get(uuid)
-    return orm_resource
+    return await resource_cache.get(kbid, uuid)
 
 
 async def _get_resource_from_db(kbid: str, uuid: str) -> Optional[ResourceORM]:

{nucliadb-6.4.0.post4192.dist-info → nucliadb-6.4.0.post4196.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4192
+Version: 6.4.0.post4196
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,17 +20,17 @@ 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.post4192
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4192
-Requires-Dist: nucliadb-protos>=6.4.0.post4192
-Requires-Dist: nucliadb-models>=6.4.0.post4192
-Requires-Dist: nidx-protos>=6.4.0.post4192
+Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4196
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4196
+Requires-Dist: nucliadb-protos>=6.4.0.post4196
+Requires-Dist: nucliadb-models>=6.4.0.post4196
+Requires-Dist: nidx-protos>=6.4.0.post4196
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]
 Requires-Dist: argdantic
 Requires-Dist: aiohttp>=3.11.11
-Requires-Dist: lru-dict>=1.1.7
+Requires-Dist: async-lru>=2
 Requires-Dist: backoff
 Requires-Dist: aiofiles>=0.8.0
 Requires-Dist: psutil>=5.9.7

{nucliadb-6.4.0.post4192.dist-info → nucliadb-6.4.0.post4196.dist-info}/RECORD

@@ -39,7 +39,7 @@ migrations/pg/0004_catalog_facets.py,sha256=FJFASHjfEHG3sNve9BP2HnnLO4xr7dnR6Qpc
 migrations/pg/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/__init__.py,sha256=_abCmDJ_0ku483Os4UAjPX7Nywm39cQgAV_DiyjsKeQ,891
 nucliadb/health.py,sha256=UIxxA4oms4HIsCRZM_SZsdkIZIlgzmOxw-qSHLlWuak,3465
-nucliadb/learning_proxy.py,sha256=Gf76qXxjl1lrHEFaCpOUfjjf0ab6eGLNxLMJz3-M_mo,19354
+nucliadb/learning_proxy.py,sha256=ZFCb4sTWGmNE1SNNAF2eJEcFksn1VMW7AD8-vu7o_U4,19294
 nucliadb/metrics_exporter.py,sha256=N9ncgPdmkP5_hqlxcBrH7ZFrgM7BkpaZrndmUpdfDKc,4981
 nucliadb/openapi.py,sha256=wDiw0dVEvTpJvbatkJ0JZLkKm9RItZT5PWRHjqRfqTA,2272
 nucliadb/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -53,7 +53,7 @@ nucliadb/backups/settings.py,sha256=SyzsInj1BRbBI0atg5IXWbMbOZ_eVg4eSQ3IcnUhCxQ,
 nucliadb/backups/tasks.py,sha256=WkL1LgdYBHbV_A5ilyYv5p3zmXwxH68TDudytN5f7zk,4225
 nucliadb/backups/utils.py,sha256=_Vogjqcru5oqNZM-bZ0q7Ju79Bv1PD-LVFEa7Z-Q13I,1261
 nucliadb/common/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/common/cache.py,sha256=r3zFCivPq9BWY1S4UCraY_DnSiXh7Bd0C7wmwGDX0D8,5255
+nucliadb/common/cache.py,sha256=NM69CVvNjlh58jiVUF1JeYPmBO7_L4rB3tffxK0k_vI,6549
 nucliadb/common/constants.py,sha256=QpigxJh_CtD85Evy0PtV5cVq6x0U_f9xfIcXz1ymkUg,869
 nucliadb/common/counters.py,sha256=8lOi3A2HeLDDlcNaS2QT1SfD3350VPBjiY3FkmHH1V8,977
 nucliadb/common/ids.py,sha256=4QjoIofes_vtKj2HsFWZf8VVIVWXxdkYtLpx1n618Us,8239
@@ -228,7 +228,7 @@ nucliadb/search/api/v1/resource/utils.py,sha256=-NjZqAQtFEXKpIh8ui5S26ItnJ5rzmmG
 nucliadb/search/requesters/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb/search/requesters/utils.py,sha256=o5JtXX5KrqMtUJo3u6rw9EOOKXPiw-GaF0oGuZu7PPc,6225
 nucliadb/search/search/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/cache.py,sha256=WUQe5tbAbJSjbcqqMWve9w_rziZBOCi_HCqaw31ghJU,5428
+nucliadb/search/search/cache.py,sha256=-6l3i2Qi8ig2SM_FCgOLIaQ48XVj7L5ctd5PdQRY5mY,4458
 nucliadb/search/search/cut.py,sha256=ytY0_GY7ocNjfxTb4aosxEp4ZfhQNDP--JkhEMGD298,1153
 nucliadb/search/search/exceptions.py,sha256=klGLgAGGrXcSGix_W6418ZBMqDchAIGjN77ofkOScEI,1039
 nucliadb/search/search/fetch.py,sha256=eiljOKim-4OOEZn-3fyVZSYxztCH156BXYdqlIwVdN4,6181
@@ -325,7 +325,7 @@ nucliadb/train/generators/paragraph_streaming.py,sha256=1xsc_IqP-1M0TzYTqu5qCvWB
 nucliadb/train/generators/question_answer_streaming.py,sha256=yZZD0GpuHdV-BT4O8CV1sYVDk8ri8yGPVhfjSp3FTBQ,5626
 nucliadb/train/generators/sentence_classifier.py,sha256=bp-UeIbZ0vm4ujbTbZnqdI0JWijuMi53cmeLbP7dofI,5063
 nucliadb/train/generators/token_classifier.py,sha256=DdyMbrpxIVGWdTcz3SEN_3HwxKffUV3JGyTZzlCET8c,9503
-nucliadb/train/generators/utils.py,sha256=0jbCfD50RORC2T-pmvRHX4iXbioZQ7ZoRKwImpd7U_4,3981
+nucliadb/train/generators/utils.py,sha256=ZNwvEVPZr-eP0MW3ABN7a11hPQKaa0NdVaRcgBcTp5w,3601
 nucliadb/writer/__init__.py,sha256=S298mrZL3vr62OrBqi97mdLxgR5cReMlRJgnaQHZV7s,1304
 nucliadb/writer/app.py,sha256=ABBO8-u4pDAa61b3mCdD0TFhuHAYcxMkgpZSGgWARuE,2736
 nucliadb/writer/back_pressure.py,sha256=4OwFGq9pvAbChB3WBZAY36lclfD-gD2ouC6YsKA4bIo,16892
@@ -365,8 +365,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.post4192.dist-info/METADATA,sha256=U5VsHHLt_nINiBGZmirG1AEaniLWqSXM9mIM2I_Bz44,4226
-nucliadb-6.4.0.post4192.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-nucliadb-6.4.0.post4192.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.0.post4192.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.0.post4192.dist-info/RECORD,,
+nucliadb-6.4.0.post4196.dist-info/METADATA,sha256=jtJRm9idqV9MmQAmslJDXhOsEzTsiLBHKsdTvCJ6pqM,4223
+nucliadb-6.4.0.post4196.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+nucliadb-6.4.0.post4196.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4196.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4196.dist-info/RECORD,,