nucliadb 6.3.5.post4032__py3-none-any.whl → 6.3.5.post4033__py3-none-any.whl

nucliadb/common/cache.py

@@ -19,6 +19,7 @@
 #
 
 import asyncio
+import contextlib
 from abc import ABC, abstractmethod
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -80,20 +81,14 @@ class Cache(Generic[T], ABC):
         if len_after - len_before > 0:
             self.metrics._cache_size.inc(len_after - len_before)
 
-    def contains(self, key: str) -> bool:
-        return key in self.cache
+    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 __del__(self):
-        # we want to clear the cache before deleting the object and set the
-        # metric appropriately
-        # XXX: apparently, this doesn't work properly. Don't rely on it
-        self.clear()
-
     @abstractmethod
     @cached_property
     def metrics(self) -> CacheMetrics: ...
@@ -105,11 +100,6 @@ class ResourceCache(Cache[ResourceORM]):
         ops=resource_cache_ops,
     )
 
-    # This cache size is an arbitrary number, once we have a metric in place and
-    # we analyze memory consumption, we can adjust it with more knoweldge
-    def __init__(self, cache_size: int = 128) -> None:
-        super().__init__(cache_size)
-
 
 class ExtractedTextCache(Cache[ExtractedText]):
     """
@@ -126,9 +116,6 @@ class ExtractedTextCache(Cache[ExtractedText]):
         ops=extracted_text_cache_ops,
     )
 
-    def __init__(self, cache_size: int = 128):
-        super().__init__(cache_size)
-
 
 # Global caches (per asyncio task)
 
@@ -139,46 +126,44 @@ etcache: ContextVar[Optional[ExtractedTextCache]] = ContextVar("etcache", defaul
 # Cache management
 
 
-# Get or create a resource cache specific to the current asyncio task (and all
-# its subtasks). If you spawn subtasks that use this cache, make sure to create
-# it in the parent task, otherwise each subtask will have its own independent
-# cache instance
-def get_or_create_resource_cache(clear: bool = False) -> ResourceCache:
-    cache: Optional[ResourceCache] = rcache.get()
-    if cache is None or clear:
-        cache = ResourceCache()
-        rcache.set(cache)
-    return cache
-
-
 def get_resource_cache() -> Optional[ResourceCache]:
     return rcache.get()
 
 
-def set_resource_cache() -> None:
-    cache = ResourceCache()
-    rcache.set(cache)
+def get_extracted_text_cache() -> Optional[ExtractedTextCache]:
+    return etcache.get()
+
 
+@contextlib.contextmanager
+def _use_cache(klass: type[Cache], context_var: ContextVar, /, **kwargs):
+    """Context manager that manages a context var cache. It's responsible of
+    cache creation and cleanup.
 
-# Delete resource cache and all its content
-def delete_resource_cache() -> None:
-    cache = rcache.get()
-    if cache is not None:
-        rcache.set(None)
-        cache.clear()
+    Note the configured cache is specific to the current asyncio task (and all
+    its subtasks). If you spawn subtasks that should share a cache, make sure
+    the parent task is the one using this decorator, otherwise, each subtask
+    will use its own independent cache instance
 
+    Do not use the cache object outside the scope of this context manager!
+    Otherwise, metrics and cleanup could get wrong.
 
-def get_extracted_text_cache() -> Optional[ExtractedTextCache]:
-    return etcache.get()
+    """
+    cache = klass(**kwargs)
+    token = context_var.set(cache)
+    try:
+        yield cache
+    finally:
+        context_var.reset(token)
+        cache.clear()
 
 
-def set_extracted_text_cache() -> None:
-    value = ExtractedTextCache()
-    etcache.set(value)
+@contextlib.contextmanager
+def resource_cache(size: int):
+    with _use_cache(ResourceCache, rcache, cache_size=size) as cache:
+        yield cache
 
 
-def delete_extracted_text_cache() -> None:
-    cache = etcache.get()
-    if cache is not None:
-        etcache.set(None)
-        cache.clear()
+@contextlib.contextmanager
+def extracted_text_cache(size: int):
+    with _use_cache(ExtractedTextCache, etcache, cache_size=size) as cache:
+        yield cache

nucliadb/search/search/cache.py

@@ -24,12 +24,10 @@ from typing import Optional
 import backoff
 
 from nucliadb.common.cache import (
-    delete_extracted_text_cache,
-    delete_resource_cache,
+    extracted_text_cache,
     get_extracted_text_cache,
     get_resource_cache,
-    set_extracted_text_cache,
-    set_resource_cache,
+    resource_cache,
 )
 from nucliadb.common.ids import FieldId
 from nucliadb.common.maindb.utils import get_driver
@@ -55,7 +53,7 @@ async def get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
         return await _orm_get_resource(kbid, uuid)
 
     async with resource_cache.get_lock(uuid):
-        if not resource_cache.contains(uuid):
+        if uuid not in resource_cache:
             resource_cache.metrics.ops.inc({"type": "miss"})
             orm_resource = await _orm_get_resource(kbid, uuid)
         else:
@@ -146,10 +144,9 @@ def request_caches():
     ...     resource = await get_resource(kbid, uuid)
     ...     extracted_text = await get_extracted_text_from_field_id(kbid, rid, field_id)
     """
-    set_resource_cache()
-    set_extracted_text_cache()
-    try:
+
+    # This cache size is an arbitrary number, once we have a metric in place and
+    # we analyze memory consumption, we can adjust it with more knoweldge
+    cache_size = 50
+    with resource_cache(cache_size), extracted_text_cache(cache_size):
         yield
-    finally:
-        delete_resource_cache()
-        delete_extracted_text_cache()

nucliadb/train/generator.py

@@ -22,6 +22,7 @@ from typing import AsyncIterator, Optional
 
 from fastapi import HTTPException
 
+from nucliadb.common.cache import resource_cache
 from nucliadb.train.generators.field_classifier import (
     field_classification_batch_generator,
 )
@@ -85,7 +86,10 @@ async def generate_train_data(kbid: str, shard: str, trainset: TrainSet):
             detail=f"Invalid train type '{TaskType.Name(trainset.type)}'",
         )
 
-    async for item in batch_generator:
-        payload = item.SerializeToString()
-        yield len(payload).to_bytes(4, byteorder="big", signed=False)
-        yield payload
+    # This cache size is an arbitrary number, once we have a metric in place and
+    # we analyze memory consumption, we can adjust it with more knoweldge
+    with resource_cache(size=20):
+        async for item in batch_generator:
+            payload = item.SerializeToString()
+            yield len(payload).to_bytes(4, byteorder="big", signed=False)
+            yield payload

nucliadb/train/generators/utils.py

@@ -20,7 +20,7 @@
 
 from typing import Any, AsyncGenerator, AsyncIterator, Optional, Type
 
-from nucliadb.common.cache import get_or_create_resource_cache
+from nucliadb.common.cache import get_resource_cache
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
@@ -31,22 +31,30 @@ from nucliadb_utils.utilities import get_storage
 
 
 async def get_resource_from_cache_or_db(kbid: str, uuid: str) -> Optional[ResourceORM]:
-    resource_cache = get_or_create_resource_cache()
+    resource_cache = get_resource_cache()
     orm_resource: Optional[ResourceORM] = None
-    if not resource_cache.contains(uuid):
+    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"})
-        storage = await get_storage(service_name=SERVICE_NAME)
-        async with get_driver().transaction(read_only=True) as transaction:
-            kb = KnowledgeBoxORM(transaction, storage, kbid)
-            orm_resource = await kb.get(uuid)
-            if orm_resource is not None:
-                resource_cache.set(uuid, orm_resource)
+        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
 
 
+async def _get_resource_from_db(kbid: str, uuid: str) -> Optional[ResourceORM]:
+    storage = await get_storage(service_name=SERVICE_NAME)
+    async with get_driver().transaction(read_only=True) as transaction:
+        kb = KnowledgeBoxORM(transaction, storage, kbid)
+        return await kb.get(uuid)
+
+
 async def get_paragraph(kbid: str, paragraph_id: str) -> str:
     if paragraph_id.count("/") == 5:
         rid, field_type, field, split_str, start_end = paragraph_id.split("/")

{nucliadb-6.3.5.post4032.dist-info → nucliadb-6.3.5.post4033.dist-info}/METADATA

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

{nucliadb-6.3.5.post4032.dist-info → nucliadb-6.3.5.post4033.dist-info}/RECORD

@@ -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=bwkKJ3tgsDE794SIklkOrAwY8vPSvyambC0y8ziMthg,5669
+nucliadb/common/cache.py,sha256=r3zFCivPq9BWY1S4UCraY_DnSiXh7Bd0C7wmwGDX0D8,5255
 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=cQZ4-NftiMljoWQ7-Zl7nWfr6u_FY8u_wc9kTvKQcAg,6999
 nucliadb/search/search/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/cache.py,sha256=F7rT0X-bp02Qh-heNo1tRhXG3iqTkL41F1wA_HIZdoM,5404
+nucliadb/search/search/cache.py,sha256=WUQe5tbAbJSjbcqqMWve9w_rziZBOCi_HCqaw31ghJU,5428
 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=XJHIFnZmXM_8Kb37lb4lg1GYG7cZ1plT-qAIb_QziX4,6184
@@ -293,7 +293,7 @@ nucliadb/tests/config.py,sha256=JN_Jhgj-fwM9_8IeO9pwxr6C1PiwRDrXxm67Y38rU30,2080
 nucliadb/tests/vectors.py,sha256=CcNKx-E8LPpyvRyljbmb-Tn_wST9Juw2CBoogWrKiTk,62843
 nucliadb/train/__init__.py,sha256=NVwe5yULoHXb80itIJT8YJYEz2xbiOPQ7_OMys6XJw8,1301
 nucliadb/train/app.py,sha256=TiRttTvekLuZdIvi46E4HyuumDTkR4G4Luqq3fEdjes,2824
-nucliadb/train/generator.py,sha256=0_zqWsLUHmJZl0lXhGorO5CWSkl42-k78dqb1slZ5h0,3904
+nucliadb/train/generator.py,sha256=w2Jv605Zv4o1F9SQc1NlpuVHiJ3vyaCufv2DKTDlXb4,4158
 nucliadb/train/lifecycle.py,sha256=a96KuAVZ0sf9TVVW6v6szVXn2eGBaboszwnv_HdWiZk,1820
 nucliadb/train/models.py,sha256=BmgmMjDsu_1Ih5JDAqo6whhume90q0ASJcDP9dkMQm8,1198
 nucliadb/train/nodes.py,sha256=HROQMRw2g5sJTnuBagh3B0id3iWonRJ68tg3skOme9k,5748
@@ -321,7 +321,7 @@ nucliadb/train/generators/paragraph_streaming.py,sha256=dsM7a5hBd2iokvFuxnZhQeko
 nucliadb/train/generators/question_answer_streaming.py,sha256=P7-de4W4yW2mgEQ82fF2OZVyx6QJHXezY52qDciDcmw,5680
 nucliadb/train/generators/sentence_classifier.py,sha256=DuvXfnWvLhklYR_qFGk2LqUyl2JE7CMVFwuHaPyC9Ys,5121
 nucliadb/train/generators/token_classifier.py,sha256=0848GqoXh8ywU82cPUrkzOM53-lZ1MVCw--8yDABigY,9557
-nucliadb/train/generators/utils.py,sha256=bArGm3MyFvHAIImhnQ3GZ6u2mZAhmRd3t3AbfDK-Aeg,3756
+nucliadb/train/generators/utils.py,sha256=0jbCfD50RORC2T-pmvRHX4iXbioZQ7ZoRKwImpd7U_4,3981
 nucliadb/writer/__init__.py,sha256=S298mrZL3vr62OrBqi97mdLxgR5cReMlRJgnaQHZV7s,1304
 nucliadb/writer/app.py,sha256=ABBO8-u4pDAa61b3mCdD0TFhuHAYcxMkgpZSGgWARuE,2736
 nucliadb/writer/back_pressure.py,sha256=JaiC2JAugVA92gDHzABZFiuQexiOKZC9C-3Jn9VF-M0,17898
@@ -361,8 +361,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.3.5.post4032.dist-info/METADATA,sha256=hbwiwOO_3khU7Rm8IyIFwpmCR3KwiILVV-agkTKmNgA,4301
-nucliadb-6.3.5.post4032.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-nucliadb-6.3.5.post4032.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.5.post4032.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.5.post4032.dist-info/RECORD,,
+nucliadb-6.3.5.post4033.dist-info/METADATA,sha256=3k6HCLT1QjmqSlXDmUdknc4RqNGQQQ6Up3klneIqzwU,4301
+nucliadb-6.3.5.post4033.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+nucliadb-6.3.5.post4033.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.3.5.post4033.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.3.5.post4033.dist-info/RECORD,,