nucliadb 6.3.7.post4079__py3-none-any.whl → 6.3.7.post4091__py3-none-any.whl

nucliadb/common/context/__init__.py

@@ -18,17 +18,19 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+from typing import Optional
 
 from nucliadb.common.cluster.manager import KBShardManager
 from nucliadb.common.cluster.settings import in_standalone_mode
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
-from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
+from nucliadb.common.nidx import NidxUtility, start_nidx_utility, stop_nidx_utility
 from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.partition import PartitionUtility
 from nucliadb_utils.settings import indexing_settings
 from nucliadb_utils.storages.storage import Storage
+from nucliadb_utils.transaction import TransactionUtility
 from nucliadb_utils.utilities import (
     get_storage,
     start_nats_manager,
@@ -42,16 +44,34 @@ from nucliadb_utils.utilities import (
 
 
 class ApplicationContext:
-    kv_driver: Driver
-    shard_manager: KBShardManager
-    blob_storage: Storage
-    partitioning: PartitionUtility
-    nats_manager: NatsConnectionManager
-
-    def __init__(self, service_name: str = "service") -> None:
+    def __init__(
+        self,
+        service_name: str = "service",
+        kv_driver: bool = True,
+        blob_storage: bool = True,
+        shard_manager: bool = True,
+        partitioning: bool = True,
+        nats_manager: bool = True,
+        transaction: bool = True,
+        nidx: bool = True,
+    ) -> None:
         self.service_name = service_name
         self._initialized: bool = False
         self._lock = asyncio.Lock()
+        self._kv_driver: Optional[Driver] = None
+        self._blob_storage: Optional[Storage] = None
+        self._shard_manager: Optional[KBShardManager] = None
+        self._partitioning: Optional[PartitionUtility] = None
+        self._nats_manager: Optional[NatsConnectionManager] = None
+        self._transaction: Optional[TransactionUtility] = None
+        self._nidx: Optional[NidxUtility] = None
+        self.enabled_kv_driver = kv_driver
+        self.enabled_blob_storage = blob_storage
+        self.enabled_shard_manager = shard_manager
+        self.enabled_partitioning = partitioning
+        self.enabled_nats_manager = nats_manager
+        self.enabled_transaction = transaction
+        self.enabled_nidx = nidx
 
     async def initialize(self) -> None:
         if self._initialized:
@@ -63,30 +83,75 @@ class ApplicationContext:
             self._initialized = True
 
     async def _initialize(self):
-        self.kv_driver = await setup_driver()
-        self.blob_storage = await get_storage()
-        self.shard_manager = await setup_cluster()
-        self.partitioning = start_partitioning_utility()
-        if not in_standalone_mode():
-            self.nats_manager = await start_nats_manager(
+        if self.enabled_kv_driver:
+            self._kv_driver = await setup_driver()
+        if self.enabled_blob_storage:
+            self._blob_storage = await get_storage()
+        if self.enabled_shard_manager:
+            self._shard_manager = await setup_cluster()
+        if self.enabled_partitioning:
+            self._partitioning = start_partitioning_utility()
+        if not in_standalone_mode() and self.enabled_nats_manager:
+            self._nats_manager = await start_nats_manager(
                 self.service_name,
                 indexing_settings.index_jetstream_servers,
                 indexing_settings.index_jetstream_auth,
             )
-        self.transaction = await start_transaction_utility(self.service_name)
-        self.nidx = await start_nidx_utility()
+        if self.enabled_transaction:
+            self._transaction = await start_transaction_utility(self.service_name)
+        if self.enabled_nidx:
+            self._nidx = await start_nidx_utility()
+
+    @property
+    def kv_driver(self) -> Driver:
+        assert self._kv_driver is not None, "Driver not initialized"
+        return self._kv_driver
+
+    @property
+    def shard_manager(self) -> KBShardManager:
+        assert self._shard_manager is not None, "Shard manager not initialized"
+        return self._shard_manager
+
+    @property
+    def blob_storage(self) -> Storage:
+        assert self._blob_storage is not None, "Blob storage not initialized"
+        return self._blob_storage
+
+    @property
+    def partitioning(self) -> PartitionUtility:
+        assert self._partitioning is not None, "Partitioning not initialized"
+        return self._partitioning
+
+    @property
+    def nats_manager(self) -> NatsConnectionManager:
+        assert self._nats_manager is not None, "NATS manager not initialized"
+        return self._nats_manager
+
+    @property
+    def transaction(self) -> TransactionUtility:
+        assert self._transaction is not None, "Transaction utility not initialized"
+        return self._transaction
+
+    @property
+    def nidx(self) -> NidxUtility:
+        assert self._nidx is not None, "Nidx utility not initialized"
+        return self._nidx
 
     async def finalize(self) -> None:
         if not self._initialized:
             return
-
-        await stop_nidx_utility()
-        await stop_transaction_utility()
-        if not in_standalone_mode():
+        if self.enabled_nidx:
+            await stop_nidx_utility()
+        if self.enabled_transaction:
+            await stop_transaction_utility()
+        if not in_standalone_mode() and self.enabled_nats_manager:
             await stop_nats_manager()
-
-        stop_partitioning_utility()
-        await teardown_cluster()
-        await teardown_driver()
-        await teardown_storage()
+        if self.enabled_partitioning:
+            stop_partitioning_utility()
+        if self.enabled_shard_manager:
+            await teardown_cluster()
+        if self.enabled_blob_storage:
+            await teardown_storage()
+        if self.enabled_kv_driver:
+            await teardown_driver()
         self._initialized = False

nucliadb/common/context/fastapi.py

@@ -19,6 +19,7 @@
 #
 
 from contextlib import asynccontextmanager
+from typing import Optional
 
 from fastapi import FastAPI
 from starlette.routing import Mount
@@ -27,8 +28,9 @@ from nucliadb.common.context import ApplicationContext
 
 
 @asynccontextmanager
-async def inject_app_context(app: FastAPI):
-    context = ApplicationContext()
+async def inject_app_context(app: FastAPI, context: Optional[ApplicationContext] = None):
+    if context is None:
+        context = ApplicationContext()
 
     app.state.context = context
 

nucliadb/train/app.py

@@ -33,18 +33,15 @@ from nucliadb_telemetry.fastapi.utils import (
     client_disconnect_handler,
     global_exception_handler,
 )
-from nucliadb_utils.audit.stream import AuditMiddleware
 from nucliadb_utils.authentication import NucliaCloudAuthenticationBackend
 from nucliadb_utils.fastapi.openapi import extend_openapi
 from nucliadb_utils.fastapi.versioning import VersionedFastAPI
 from nucliadb_utils.settings import running_settings
-from nucliadb_utils.utilities import get_audit
 
 middleware = []
 middleware.extend(
     [
         Middleware(AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()),
-        Middleware(AuditMiddleware, audit_utility_getter=get_audit),
     ]
 )
 

nucliadb/train/generator.py

@@ -45,6 +45,7 @@ from nucliadb.train.generators.sentence_classifier import (
 from nucliadb.train.generators.token_classifier import (
     token_classification_batch_generator,
 )
+from nucliadb.train.settings import settings
 from nucliadb.train.types import TrainBatch
 from nucliadb.train.utils import get_shard_manager
 from nucliadb_protos.dataset_pb2 import TaskType, TrainSet
@@ -88,7 +89,7 @@ async def generate_train_data(kbid: str, shard: str, trainset: TrainSet):
 
     # 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):
+    with resource_cache(size=settings.resource_cache_size):
         async for item in batch_generator:
             payload = item.SerializeToString()
             yield len(payload).to_bytes(4, byteorder="big", signed=False)

nucliadb/train/lifecycle.py

@@ -22,32 +22,37 @@ from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
+from nucliadb.common.context import ApplicationContext
 from nucliadb.common.context.fastapi import inject_app_context
-from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb.train import SERVICE_NAME
 from nucliadb.train.utils import (
-    start_shard_manager,
+    start_shard_manager as start_train_shard_manager,
+)
+from nucliadb.train.utils import (
     start_train_grpc,
-    stop_shard_manager,
     stop_train_grpc,
 )
+from nucliadb.train.utils import (
+    stop_shard_manager as stop_train_shard_manager,
+)
 from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
-from nucliadb_utils.utilities import start_audit_utility, stop_audit_utility
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     await setup_telemetry(SERVICE_NAME)
-    await start_nidx_utility()
-    await start_shard_manager()
+    await start_train_shard_manager()
     await start_train_grpc(SERVICE_NAME)
-    await start_audit_utility(SERVICE_NAME)
-
-    async with inject_app_context(app):
-        yield
-
-    await stop_audit_utility()
-    await stop_train_grpc()
-    await stop_shard_manager()
-    await stop_nidx_utility()
-    await clean_telemetry(SERVICE_NAME)
+    try:
+        context = ApplicationContext(
+            service_name="train",
+            partitioning=False,
+            nats_manager=False,
+            transaction=False,
+        )
+        async with inject_app_context(app, context):
+            yield
+    finally:
+        await stop_train_grpc()
+        await stop_train_shard_manager()
+        await clean_telemetry(SERVICE_NAME)

nucliadb/train/settings.py

@@ -33,5 +33,7 @@ class Settings(DriverSettings):
 
     internal_search_api: str = "http://search.nuclia.svc.cluster.local:8030/api/v1/kb/{kbid}/search"
 
+    resource_cache_size: int = 2
+
 
 settings = Settings()

{nucliadb-6.3.7.post4079.dist-info → nucliadb-6.3.7.post4091.dist-info}/METADATA

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

{nucliadb-6.3.7.post4079.dist-info → nucliadb-6.3.7.post4091.dist-info}/RECORD

@@ -70,8 +70,8 @@ nucliadb/common/cluster/settings.py,sha256=JPwV_0U_i618Tn66GWUq6qCKNjy4TWkGEGld9
 nucliadb/common/cluster/utils.py,sha256=K0BZ75m12Tbi7KS1Tr3LqXJSPhJzUHkBMo4j8W8nQpA,4639
 nucliadb/common/cluster/standalone/__init__.py,sha256=itSI7dtTwFP55YMX4iK7JzdMHS5CQVUiB1XzQu4UBh8,833
 nucliadb/common/cluster/standalone/utils.py,sha256=af3r-x_GF7A6dwIAhZLR-r-SZQEVxsFrDKeMfUTA6G0,1908
-nucliadb/common/context/__init__.py,sha256=ZLUvKuIPaolKeA3aeZa2JcHwCIaEauNu8WpdKsiINXo,3354
-nucliadb/common/context/fastapi.py,sha256=j3HZ3lne6mIfw1eEar2het8RWzv6UruUZpXaKieSLOs,1527
+nucliadb/common/context/__init__.py,sha256=C3SsVpwyqHj2PPAUgtf0GVmq9KVQi3q5dJCsMIs9fpY,5999
+nucliadb/common/context/fastapi.py,sha256=mH_8n5t7quNSPivNM2JS5EQf2sTVJsdzXW6LaY7EHAA,1629
 nucliadb/common/datamanagers/__init__.py,sha256=jksw4pXyXb05SG3EN-BPBrhc1u1Ge_m21PYqD7NYQEs,2118
 nucliadb/common/datamanagers/atomic.py,sha256=WihdtBWQIAuElZQjh1xQ--q5dJowwlkovqsW-OB_t2k,3230
 nucliadb/common/datamanagers/cluster.py,sha256=iU0b7AESm1Yi8Wp3pIKgqixZGNMjeBrxSpvEKsaZKgY,1831
@@ -296,16 +296,16 @@ nucliadb/tests/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,83
 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=w2Jv605Zv4o1F9SQc1NlpuVHiJ3vyaCufv2DKTDlXb4,4158
-nucliadb/train/lifecycle.py,sha256=a96KuAVZ0sf9TVVW6v6szVXn2eGBaboszwnv_HdWiZk,1820
+nucliadb/train/app.py,sha256=z6xlGVVVaJmZZmLPIVTgkjD-wIz5b0NYlXAQp7hBHYw,2652
+nucliadb/train/generator.py,sha256=B3bdkH-IWiqib-ALpU_g8wFe0xOmdE0kEPxvki28yEU,4229
+nucliadb/train/lifecycle.py,sha256=3HadM4GRsYb2m-v4jtdr9C-KBEBx8GlrJDArPYi3SWQ,1960
 nucliadb/train/models.py,sha256=BmgmMjDsu_1Ih5JDAqo6whhume90q0ASJcDP9dkMQm8,1198
 nucliadb/train/nodes.py,sha256=HROQMRw2g5sJTnuBagh3B0id3iWonRJ68tg3skOme9k,5748
 nucliadb/train/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/train/resource.py,sha256=3qQ_9Zdt5JAbtD-wpmt7OeDGRNKS-fQdKAuIQfznZm0,16219
 nucliadb/train/run.py,sha256=evz6CKVfJOzkbHMoaYz2mTMlKjJnNOb1O8zBBWMpeBw,1400
 nucliadb/train/servicer.py,sha256=scbmq8FriKsJGkOcoZB2Fg_IyIExn9Ux4W30mGDlkJQ,5728
-nucliadb/train/settings.py,sha256=rrLtgdBmuthtIObLuZUaeuo4VBGU2PJRazquQbtPBeI,1383
+nucliadb/train/settings.py,sha256=8_-XCO_nrE98cMJpe0fYkXeF2bkWKff1VX-2qdwcsjs,1417
 nucliadb/train/types.py,sha256=xyVYy8kHipAWoDb7Pn7dCYQ_efHPzDW_3AXg5M-aV28,1519
 nucliadb/train/upload.py,sha256=fTjH1KEL-0ogf3LV0T6ODO0QdPGwdZShSUtFUCAcUlA,3256
 nucliadb/train/uploader.py,sha256=xdLGz1ToDue9Q_M8A-_KYkO-V6fWKYOZQ6IGM4FuwWA,6424
@@ -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.3.7.post4079.dist-info/METADATA,sha256=WP6AMoJHWMGwWeqYxHCKJbhd7mTgzJDQD6f4fCrdp6w,4226
-nucliadb-6.3.7.post4079.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-nucliadb-6.3.7.post4079.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.3.7.post4079.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.3.7.post4079.dist-info/RECORD,,
+nucliadb-6.3.7.post4091.dist-info/METADATA,sha256=en368pfXpQ5ec1oDjrmsWTHGZbhM4PYtqKEbImJziR0,4226
+nucliadb-6.3.7.post4091.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+nucliadb-6.3.7.post4091.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.3.7.post4091.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.3.7.post4091.dist-info/RECORD,,