nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/ingest/settings.py

@@ -20,14 +20,13 @@
 from enum import Enum
 from typing import Optional
 
-from pydantic import BaseSettings, Field
+from pydantic import Field
+from pydantic_settings import BaseSettings
 
 
-class DriverConfig(str, Enum):
-    REDIS = "redis"
-    TIKV = "tikv"
+class DriverConfig(Enum):
     PG = "pg"
-    LOCAL = "local"
+    LOCAL = "local"  # Not recommended for production
     NOT_SET = "notset"  # setting not provided
 
     @classmethod
@@ -41,19 +40,7 @@ class DriverConfig(str, Enum):
 
 
 class DriverSettings(BaseSettings):
-    driver: DriverConfig = Field(
-        default=DriverConfig.NOT_SET, description="K/V storage driver"
-    )
-    driver_redis_url: Optional[str] = Field(
-        default=None, description="Redis URL. Example: redis://localhost:6379"
-    )
-    driver_tikv_url: Optional[list[str]] = Field(
-        default=None,
-        description=(
-            "TiKV PD (Placement Driver) URLs. The URL to the cluster manager of"
-            "TiKV. Example: '[\"tikv-pd.svc:2379\"]'"
-        ),
-    )
+    driver: DriverConfig = Field(default=DriverConfig.PG, description="K/V storage driver")
     driver_local_url: Optional[str] = Field(
         default=None,
         description="Local path to store data on file system. Example: /nucliadb/data/main",
@@ -62,13 +49,17 @@ class DriverSettings(BaseSettings):
         default=None,
         description="PostgreSQL DSN. The connection string to the PG server. Example: postgres://username:password@postgres:5432/nucliadb.",  # noqa
     )
+    driver_pg_connection_pool_min_size: int = Field(
+        default=10,
+        description="PostgreSQL min pool size. The minimum number of connections to the PostgreSQL server.",
+    )
     driver_pg_connection_pool_max_size: int = Field(
         default=20,
         description="PostgreSQL max pool size. The maximum number of connections to the PostgreSQL server.",
     )
-    driver_tikv_connection_pool_size: int = Field(
-        default=3,
-        description="TiKV max pool size. The maximum number of connections to the TiKV server.",
+    driver_pg_connection_pool_acquire_timeout_ms: int = Field(
+        default=1000,
+        description="PostgreSQL pool acquire timeout in ms. The maximum time to wait until a connection becomes available.",
     )
 
 
@@ -86,7 +77,7 @@ class Settings(DriverSettings):
     total_replicas: int = 1  # number of ingest processor replicas in the cluster
     nuclia_partitions: int = 50
 
-    max_receive_message_length: int = 4
+    max_receive_message_length: int = 500  # In MB
 
     # Search query timeouts
     relation_search_timeout: float = 10.0

nucliadb/ingest/utils.py

@@ -19,9 +19,8 @@
 #
 from typing import Optional
 
-from nucliadb_protos.writer_pb2_grpc import WriterStub
-
 from nucliadb.common.maindb.utils import setup_driver
+from nucliadb_protos.writer_pb2_grpc import WriterStub
 from nucliadb_utils.grpc import get_traced_grpc_channel
 from nucliadb_utils.settings import nucliadb_settings
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
@@ -37,11 +36,9 @@ async def start_ingest(service_name: Optional[str] = None):
     if nucliadb_settings.nucliadb_ingest is not None:
         # Its distributed lets create a GRPC client
         # We want Jaeger telemetry enabled
-        channel = get_traced_grpc_channel(
-            nucliadb_settings.nucliadb_ingest, service_name or "ingest"
-        )
+        channel = get_traced_grpc_channel(nucliadb_settings.nucliadb_ingest, service_name or "ingest")
         set_utility(Utility.CHANNEL, channel)
-        ingest = WriterStub(channel)  # type: ignore
+        ingest = WriterStub(channel)
         set_utility(Utility.INGEST, ingest)
     else:
         # Its not distributed create a ingest

nucliadb/learning_proxy.py

@@ -20,16 +20,21 @@
 import contextlib
 import json
 import logging
+import os
+from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
-from enum import Enum
+from enum import Enum, IntEnum
 from typing import Any, Optional, Union
 
 import backoff
 import httpx
 from fastapi import Request, Response
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
+from lru import LRU
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
 
+from nucliadb_protos import knowledgebox_pb2, utils_pb2
 from nucliadb_telemetry import errors
 from nucliadb_utils.settings import is_onprem_nucliadb, nuclia_settings
 
@@ -48,48 +53,150 @@ WHITELISTED_HEADERS = {
 }
 
 
-class LearningService(str, Enum):
+class LearningService(Enum):
     CONFIG = "config"
-    COLLECTOR = "collector-api"
 
 
+class SimilarityFunction(IntEnum):
+    # Keep this in sync with learning config repo
+    # It's an IntEnum to match the protobuf definition
+    DOT = 0
+    COSINE = 1
+
+
+class SemanticConfig(BaseModel):
+    # Keep this in sync with learning config repo
+    similarity: SimilarityFunction
+    size: int
+    threshold: float
+    matryoshka_dims: list[int] = []
+
+    def into_semantic_model_metadata(self) -> knowledgebox_pb2.SemanticModelMetadata:
+        semantic_model = knowledgebox_pb2.SemanticModelMetadata()
+        LEARNING_SIMILARITY_FUNCTION_TO_PROTO = {
+            SimilarityFunction.COSINE: utils_pb2.VectorSimilarity.COSINE,
+            SimilarityFunction.DOT: utils_pb2.VectorSimilarity.DOT,
+        }
+        semantic_model.similarity_function = LEARNING_SIMILARITY_FUNCTION_TO_PROTO[self.similarity]
+        semantic_model.vector_dimension = self.size
+        semantic_model.matryoshka_dimensions.extend(self.matryoshka_dims)
+        return semantic_model
+
+
+# Subset of learning configuration of nucliadb's interest. Look at
+# learning_config models for more fields
 class LearningConfiguration(BaseModel):
     semantic_model: str
+    # aka similarity function
     semantic_vector_similarity: str
-    semantic_vector_size: Optional[int]
-    semantic_threshold: Optional[float]
+    # aka vector_dimension
+    semantic_vector_size: Optional[int] = None
+    # aka min_score
+    semantic_threshold: Optional[float] = None
+    # List of possible subdivisions of the matryoshka embeddings (if the model
+    # supports it)
+    semantic_matryoshka_dimensions: Optional[list[int]] = Field(
+        default=None, alias="semantic_matryoshka_dims"
+    )
+
+    semantic_models: list[str] = Field(default_factory=list)
+
+    # This is where the config for each semantic model (aka vectorsets) is returned
+    semantic_model_configs: dict[str, SemanticConfig] = Field(default_factory=dict)
+
+    @model_validator(mode="before")
+    @classmethod
+    def maintain_bw_compatibility_with_single_model_configs(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            if not data.get("semantic_model", None) and len(data.get("semantic_models", [])) > 0:
+                data["semantic_model"] = data["semantic_models"][0]
+        return data
+
+    @model_validator(mode="after")
+    def validate_matryoshka_and_vector_dimension_consistency(self) -> Self:
+        vector_size = self.semantic_vector_size
+        matryoshka_dimensions = self.semantic_matryoshka_dimensions or []
+        if (
+            len(matryoshka_dimensions) > 0
+            and vector_size is not None
+            and vector_size not in matryoshka_dimensions
+        ):
+            raise ValueError("Semantic vector size is inconsistent with matryoshka dimensions")
+        return self
+
+    def into_semantic_models_metadata(
+        self,
+    ) -> dict[str, knowledgebox_pb2.SemanticModelMetadata]:
+        result = {}
+        for model_name, config in self.semantic_model_configs.items():
+            result[model_name] = config.into_semantic_model_metadata()
+        return result
+
+    def into_semantic_model_metadata(self) -> knowledgebox_pb2.SemanticModelMetadata:
+        semantic_model = knowledgebox_pb2.SemanticModelMetadata()
+
+        LEARNING_SIMILARITY_FUNCTION_TO_PROTO = {
+            "cosine": utils_pb2.VectorSimilarity.COSINE,
+            "dot": utils_pb2.VectorSimilarity.DOT,
+        }
+        semantic_model.similarity_function = LEARNING_SIMILARITY_FUNCTION_TO_PROTO[
+            self.semantic_vector_similarity.lower()
+        ]
+
+        if self.semantic_vector_size is not None:
+            semantic_model.vector_dimension = self.semantic_vector_size
+        else:
+            logger.warning("Vector dimension not set!")
+
+        if self.semantic_matryoshka_dimensions is not None:
+            semantic_model.matryoshka_dimensions.extend(self.semantic_matryoshka_dimensions)
+
+        return semantic_model
+
+
+class ProxiedLearningConfigError(Exception):
+    def __init__(self, status_code: int, content: bytes, content_type: str):
+        self.status_code = status_code
+        self.content = content
+        self.content_type = content_type
+
+
+def raise_for_status(response: httpx.Response) -> None:
+    try:
+        response.raise_for_status()
+    except httpx.HTTPStatusError as err:
+        content_type = err.response.headers.get("Content-Type", "application/json")
+        raise ProxiedLearningConfigError(
+            status_code=err.response.status_code,
+            content=err.response.content,
+            content_type=content_type,
+        )
 
 
 async def get_configuration(
     kbid: str,
 ) -> Optional[LearningConfiguration]:
-    async with learning_config_client() as client:
-        resp = await client.get(f"config/{kbid}")
-        try:
-            resp.raise_for_status()
-        except httpx.HTTPStatusError as err:
-            if err.response.status_code == 404:
-                return None
-            raise
-        return LearningConfiguration.parse_obj(resp.json())
+    return await learning_config_service().get_configuration(kbid)
 
 
 async def set_configuration(
     kbid: str,
     config: dict[str, Any],
 ) -> LearningConfiguration:
-    async with learning_config_client() as client:
-        resp = await client.post(f"config/{kbid}", json=config)
-        resp.raise_for_status()
-        return LearningConfiguration.parse_obj(resp.json())
+    return await learning_config_service().set_configuration(kbid, config)
+
+
+async def update_configuration(
+    kbid: str,
+    config: dict[str, Any],
+) -> None:
+    return await learning_config_service().update_configuration(kbid, config)
 
 
 async def delete_configuration(
     kbid: str,
 ) -> None:
-    async with learning_config_client() as client:
-        resp = await client.delete(f"config/{kbid}")
-        resp.raise_for_status()
+    return await learning_config_service().delete_configuration(kbid)
 
 
 async def learning_config_proxy(
@@ -107,21 +214,6 @@ async def learning_config_proxy(
     )
 
 
-async def learning_collector_proxy(
-    request: Request,
-    method: str,
-    url: str,
-    extra_headers: Optional[dict[str, str]] = None,
-) -> Union[Response, StreamingResponse]:
-    return await proxy(
-        service=LearningService.COLLECTOR,
-        request=request,
-        method=method,
-        url=url,
-        extra_headers=extra_headers,
-    )
-
-
 def is_white_listed_header(header: str) -> bool:
     return header.lower() in WHITELISTED_HEADERS
 
@@ -213,13 +305,9 @@ async def proxy(
 
 def get_base_url(service: LearningService) -> str:
     if is_onprem_nucliadb():
-        nuclia_public_url = nuclia_settings.nuclia_public_url.format(
-            zone=nuclia_settings.nuclia_zone
-        )
+        nuclia_public_url = nuclia_settings.nuclia_public_url.format(zone=nuclia_settings.nuclia_zone)
         return f"{nuclia_public_url}/api/v1"
-    learning_svc_base_url = nuclia_settings.learning_internal_svc_base_url.format(
-        service=service.value
-    )
+    learning_svc_base_url = nuclia_settings.learning_internal_svc_base_url.format(service=service.value)
     return f"{learning_svc_base_url}/api/v1/internal"
 
 
@@ -248,9 +336,7 @@ async def service_client(
         # This is a workaround to be able to run integration tests that start nucliadb with docker.
         # The learning APIs are not available in the docker setup, so we use a dummy client.
         client = DummyClient(base_url=base_url, headers=headers)
-        logger.warning(
-            "Using dummy client. If you see this in production, something is wrong."
-        )
+        logger.warning("Using dummy client. If you see this in production, something is wrong.")
     else:
         client = httpx.AsyncClient(base_url=base_url, headers=headers)  # type: ignore
     try:
@@ -299,13 +385,31 @@ class DummyClient(httpx.AsyncClient):
         return self._handle_request("DELETE", *args, **kwargs)
 
     def get_config(self, *args: Any, **kwargs: Any):
+        size = 768 if os.environ.get("TEST_SENTENCE_ENCODER") == "multilingual-2023-02-21" else 512
         lconfig = LearningConfiguration(
             semantic_model="multilingual",
             semantic_vector_similarity="cosine",
-            semantic_vector_size=None,
+            semantic_vector_size=size,
             semantic_threshold=None,
+            semantic_matryoshka_dims=[],
+            semantic_model_configs={
+                "multilingual": SemanticConfig(
+                    similarity=SimilarityFunction.COSINE,
+                    size=size,
+                    threshold=0,
+                    matryoshka_dims=[],
+                )
+            },
         )
-        return self._response(content=lconfig.dict())
+        return self._response(content=lconfig.model_dump())
+
+    def post_config(self, *args: Any, **kwargs: Any):
+        # simulate post that returns the created config
+        return self.get_config(*args, **kwargs)
+
+    def patch_config(self, *args: Any, **kwargs: Any):
+        # simulate patch that returns the updated config
+        return self.get_config(*args, **kwargs)
 
     async def request(  # type: ignore
         self,
@@ -315,9 +419,7 @@ class DummyClient(httpx.AsyncClient):
         content=None,
         headers=None,
     ) -> httpx.Response:
-        return self._handle_request(
-            method, url, params=params, content=content, headers=headers
-        )
+        return self._handle_request(method, url, params=params, content=content, headers=headers)
 
     def _handle_request(self, *args: Any, **kwargs: Any) -> httpx.Response:
         """
@@ -331,3 +433,114 @@ class DummyClient(httpx.AsyncClient):
             return getattr(self, method)(*args, **kwargs)
         else:
             return self._response()
+
+
+class LearningConfigService(ABC):
+    @abstractmethod
+    async def get_configuration(self, kbid: str) -> Optional[LearningConfiguration]: ...
+
+    @abstractmethod
+    async def set_configuration(self, kbid: str, config: dict[str, Any]) -> LearningConfiguration: ...
+
+    @abstractmethod
+    async def update_configuration(self, kbid: str, config: dict[str, Any]) -> None: ...
+
+    @abstractmethod
+    async def delete_configuration(self, kbid: str) -> None: ...
+
+
+class ProxiedLearningConfig(LearningConfigService):
+    async def get_configuration(self, kbid: str) -> Optional[LearningConfiguration]:
+        async with self._client() as client:
+            resp = await client.get(f"config/{kbid}")
+            try:
+                raise_for_status(resp)
+            except ProxiedLearningConfigError as err:
+                if err.status_code == 404:
+                    return None
+                raise
+            return LearningConfiguration.model_validate(resp.json())
+
+    async def set_configuration(self, kbid: str, config: dict[str, Any]) -> LearningConfiguration:
+        async with self._client() as client:
+            resp = await client.post(f"config/{kbid}", json=config)
+            raise_for_status(resp)
+            return LearningConfiguration.model_validate(resp.json())
+
+    async def update_configuration(self, kbid: str, config: dict[str, Any]) -> None:
+        async with self._client() as client:
+            resp = await client.patch(f"config/{kbid}", json=config)
+            raise_for_status(resp)
+            return
+
+    async def delete_configuration(self, kbid: str) -> None:
+        async with self._client() as client:
+            resp = await client.delete(f"config/{kbid}")
+            raise_for_status(resp)
+
+    @contextlib.asynccontextmanager
+    async def _client(self) -> AsyncIterator[httpx.AsyncClient]:
+        async with httpx.AsyncClient(
+            base_url=get_base_url(LearningService.CONFIG),
+            headers=get_auth_headers(),
+        ) as client:
+            yield client
+
+
+_IN_MEMORY_CONFIGS: dict[str, LearningConfiguration]
+_IN_MEMORY_CONFIGS = LRU(50)  # type: ignore
+
+
+class InMemoryLearningConfig(LearningConfigService):
+    def __init__(self):
+        self.in_memory_configs = {}
+
+    async def get_configuration(self, kbid: str) -> Optional[LearningConfiguration]:
+        return _IN_MEMORY_CONFIGS.get(kbid, None)
+
+    async def set_configuration(self, kbid: str, config: dict[str, Any]) -> LearningConfiguration:
+        if not config:
+            # generate a default config
+            default_model = os.environ.get("TEST_SENTENCE_ENCODER", "multilingual")
+            size = 768 if default_model == "multilingual-2023-02-21" else 512
+            # XXX for some reason, we override the model name and set this one
+            # default_model = "multilingual"
+            learning_config = LearningConfiguration(
+                semantic_model=default_model,
+                semantic_vector_similarity="cosine",
+                semantic_vector_size=size,
+                semantic_threshold=None,
+                semantic_matryoshka_dims=[],
+                semantic_models=[default_model],
+                semantic_model_configs={
+                    default_model: SemanticConfig(
+                        similarity=SimilarityFunction.COSINE,
+                        size=size,
+                        threshold=0,
+                        matryoshka_dims=[],
+                    )
+                },
+            )
+
+        else:
+            learning_config = LearningConfiguration.model_validate(config)
+
+        _IN_MEMORY_CONFIGS[kbid] = learning_config
+        return learning_config
+
+    async def update_configuration(self, kbid: str, config: dict[str, Any]) -> None:
+        if kbid not in _IN_MEMORY_CONFIGS:
+            raise ValueError(f"Configuration for kbid {kbid} not found")
+        learning_config = _IN_MEMORY_CONFIGS[kbid]
+        learning_config = learning_config.model_copy(update=config)
+        _IN_MEMORY_CONFIGS[kbid] = learning_config
+
+    async def delete_configuration(self, kbid: str) -> None:
+        _IN_MEMORY_CONFIGS.pop(kbid, None)
+
+
+def learning_config_service() -> LearningConfigService:
+    if nuclia_settings.dummy_learning_services:
+        return InMemoryLearningConfig()
+    else:
+        return ProxiedLearningConfig()

nucliadb/metrics_exporter.py

@@ -20,12 +20,14 @@
 from __future__ import annotations
 
 import asyncio
-from typing import AsyncGenerator, Callable
+from typing import AsyncGenerator, Callable, Tuple, cast
 
 from nucliadb import logger
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster import manager as cluster_manager
 from nucliadb.common.context import ApplicationContext
+from nucliadb.common.maindb.pg import PGDriver
+from nucliadb.common.maindb.utils import get_driver
 from nucliadb.migrator.datamanager import MigrationsDataManager
 from nucliadb_telemetry import metrics
 from nucliadb_telemetry.logs import setup_logging
@@ -34,9 +36,9 @@ from nucliadb_utils.fastapi.run import serve_metrics
 
 SHARD_COUNT = metrics.Gauge("nucliadb_node_shard_count", labels={"node": ""})
 
-MIGRATION_COUNT = metrics.Gauge(
-    "nucliadb_migration", labels={"type": "", "version": ""}
-)
+MIGRATION_COUNT = metrics.Gauge("nucliadb_migration", labels={"type": "", "version": ""})
+
+PENDING_RESOURCE_COUNT = metrics.Gauge("nucliadb_pending_resources_count")
 
 
 async def update_node_metrics(context: ApplicationContext):
@@ -57,7 +59,7 @@ async def iter_kbids(context: ApplicationContext) -> AsyncGenerator[str, None]:
     """
     Return a list of all KB ids.
     """
-    async with context.kv_driver.transaction() as txn:
+    async with context.kv_driver.transaction(read_only=True) as txn:
         async for kbid, _ in datamanagers.kb.get_kbs(txn):
             yield kbid
 
@@ -72,9 +74,7 @@ async def update_migration_metrics(context: ApplicationContext):
     mdm = MigrationsDataManager(context.kv_driver)
     global_info = await mdm.get_global_info()
     if global_info is not None:
-        MIGRATION_COUNT.set(
-            1, labels=dict(type="global", version=str(global_info.current_version))
-        )
+        MIGRATION_COUNT.set(1, labels=dict(type="global", version=str(global_info.current_version)))
 
     version_count: dict[str, int] = {}
     async for kbid in iter_kbids(context):
@@ -88,9 +88,25 @@ async def update_migration_metrics(context: ApplicationContext):
         MIGRATION_COUNT.set(count, labels=dict(type="kb", version=version))
 
 
-async def run_exporter_task(
-    context: ApplicationContext, exporter_task: Callable, interval: int
-):
+async def update_resource_metrics(context: ApplicationContext):
+    """
+    Report the number of pending resources older than some estimated processing time
+    """
+    driver = get_driver()
+    if not isinstance(driver, PGDriver):
+        return
+
+    async with driver._get_connection() as conn, conn.cursor() as cur:
+        await cur.execute(
+            "SELECT COUNT(*) FROM catalog "
+            "WHERE labels @> '{/n/s/PENDING}' "
+            "AND COALESCE(modified_at, created_at) BETWEEN NOW() - INTERVAL '1 month' AND NOW() - INTERVAL '6 hours'"
+        )
+        count = cast(Tuple[int], await cur.fetchone())[0]
+        PENDING_RESOURCE_COUNT.set(count)
+
+
+async def run_exporter_task(context: ApplicationContext, exporter_task: Callable, interval: int):
     """
     Run coroutine infinitely, catching exceptions and logging them.
     It will wait for the interval before running again.
@@ -100,9 +116,7 @@ async def run_exporter_task(
             try:
                 await exporter_task(context)
             except Exception:
-                logger.error(
-                    f"Error on exporter task {exporter_task.__name__}", exc_info=True
-                )
+                logger.error(f"Error on exporter task {exporter_task.__name__}", exc_info=True)
             await asyncio.sleep(interval)
     except asyncio.CancelledError:
         pass
@@ -114,12 +128,9 @@ async def run_exporter(context: ApplicationContext):
     for export_task, interval in [
         (update_node_metrics, 10),
         (update_migration_metrics, 60 * 3),
+        (update_resource_metrics, 60 * 5),
     ]:
-        tasks.append(
-            asyncio.create_task(
-                run_exporter_task(context, export_task, interval=interval)
-            )
-        )
+        tasks.append(asyncio.create_task(run_exporter_task(context, export_task, interval=interval)))
     try:
         while True:
             await asyncio.sleep(10)

nucliadb/middleware/__init__.py

@@ -39,9 +39,7 @@ class ProcessTimeHeaderMiddleware(BaseHTTPMiddleware):
             exposed_headers.append(PROCESS_TIME_HEADER)
             response.headers[ACCESS_CONTROL_EXPOSE_HEADER] = ",".join(exposed_headers)
 
-    async def dispatch(
-        self, request: Request, call_next: RequestResponseEndpoint
-    ) -> Response:
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
         response = None
         start = time.perf_counter()
         try:

nucliadb/migrator/command.py

@@ -53,9 +53,7 @@ def validate():
     versions = set()
     for migration in migrations:
         if migration.version in versions:
-            raise MigrationValidationError(
-                f"Migration {migration.version} is duplicated"
-            )
+            raise MigrationValidationError(f"Migration {migration.version} is duplicated")
         versions.add(migration.version)