feature(asr-workers,extract-worker,translation-worker): model cache

datashare-python/datashare_python/config.py

@@ -13,7 +13,7 @@
 from .objects import BaseModel
 from .task_client import DatashareTaskClient
 from .types_ import TemporalClient
-from .utils import PYDANTIC_DATA_CONVERTER
+from .utils import PYDANTIC_DATA_CONVERTER, SharedResources, close_cm_callback
 
 _ALL_LOGGERS = [datashare_python.__name__]
 
@@ -94,6 +94,19 @@ class LoggingConfig(BaseModel):
     loggers: dict[str, LogLevel]
 
 
+class ResourceCacheConfig(BaseModel):
+    size: int = 1
+    exit_context_managers: bool = True
+
+    def to_resource_cache(self) -> SharedResources:
+        eviction_callback = None
+        if self.exit_context_managers:
+            eviction_callback = close_cm_callback
+        return SharedResources(
+            cache_size=self.size, eviction_callback=eviction_callback
+        )
+
+
 _DEFAULT_LOGGERS = {datashare_python.__name__: "INFO"}
 _DEFAULT_LOGGING_CONFIG = LoggingConfig(
     format=LogFormat.DEFAULT, loggers=_DEFAULT_LOGGERS

datashare-python/datashare_python/constants.py

@@ -14,5 +14,3 @@
 MANIFEST_JSON = "manifest.json"
 
 TIKA_METADATA_RESOURCENAME = "tika_metadata_resourcename"
-
-DEFAULT_SHARED_RESOURCES_SIZE = 1

datashare-python/datashare_python/dependencies.py

@@ -12,7 +12,6 @@
 from .config import LogLevel, WorkerConfig
 from .exceptions import DependencyInjectionError
 from .logging_ import setup_worker_loggers
-from .objects import Shared
 from .task_client import DatashareTaskClient
 from .types_ import ContextManagerFactory, TemporalClient
 
@@ -24,7 +23,6 @@
 TASK_CLIENT: ContextVar[DatashareTaskClient] = ContextVar("task_client")
 TEMPORAL_CLIENT: ContextVar[TemporalClient] = ContextVar("temporal_client")
 WORKER_CONFIG: ContextVar[WorkerConfig] = ContextVar("worker_config")
-SHARED: ContextVar[Shared] = ContextVar("shared")
 
 
 def set_event_loop(event_loop: AbstractEventLoop) -> None:
@@ -101,20 +99,6 @@ def lifespan_temporal_client() -> TemporalClient:
         raise DependencyInjectionError("temporal client") from e
 
 
-# Setup shared resources
-async def set_shared_resources(shared: Shared) -> Shared:
-    SHARED.set(shared)
-    return shared
-
-
-# Return shared resources
-def lifespan_shared_resources() -> Shared:
-    try:
-        return SHARED.get()
-    except LookupError as e:
-        raise DependencyInjectionError("shared resources") from e
-
-
 @asynccontextmanager
 async def with_dependencies(
     dependencies: list[ContextManagerFactory], **kwargs

datashare-python/datashare_python/objects.py

@@ -1,9 +1,8 @@
 import logging
 import os
 from abc import ABC
-from asyncio import Lock
 from collections.abc import Awaitable, Callable
-from dataclasses import InitVar, dataclass, field
+from dataclasses import dataclass
 from datetime import UTC, datetime
 from enum import StrEnum, unique
 from io import BytesIO
@@ -12,15 +11,12 @@
 
 import langcodes
 from icij_common.registrable import Registrable
-from lru import LRU
 from pydantic_core import PydanticCustomError, ValidationError, core_schema
-from pydantic_core.core_schema import (
-    PlainValidatorFunctionSchema,
-)
+from pydantic_core.core_schema import PlainValidatorFunctionSchema
 from pydantic_extra_types.language_code import LanguageName
 from temporalio import workflow
 
-from .constants import DEFAULT_SHARED_RESOURCES_SIZE, TIKA_METADATA_RESOURCENAME
+from .constants import TIKA_METADATA_RESOURCENAME
 
 with workflow.unsafe.imports_passed_through():
     from icij_common.es import (
@@ -488,47 +484,3 @@ class TaskGroup:
     @classmethod
     def python(cls) -> Self:
         return cls(name="PYTHON")
-
-
-@dataclass(frozen=True)
-class Shared:
-    cache_size: InitVar[int] = DEFAULT_SHARED_RESOURCES_SIZE
-    eviction_callback: InitVar[Callable] = None
-    _resources: LRU = field(init=False, repr=False)
-    _lock: Lock = field(init=False, repr=False)
-
-    def __post_init__(self, cache_size: int, eviction_callback: Callable) -> None:
-        object.__setattr__(
-            self, "_resources", LRU(cache_size, callback=eviction_callback)
-        )
-        object.__setattr__(self, "_lock", Lock())
-
-    def get_resource(
-        self, key: str, default: Any = None, *, set_if_unavailable: bool = True
-    ) -> Any:
-        if key not in self._resources and set_if_unavailable:
-            self.set_resource(key, default)
-
-        return self._resources.get(key, default)
-
-    def set_resource(self, key: str, value: Any) -> None:
-        self._resources[key] = value
-
-    def pop_resource(self, key: str, default: Any = None) -> Any:
-        return self._resources.pop(key, default)
-
-    async def async_get_resource(
-        self, key: str, default: Any = None, *, set_if_unavailable: bool = True
-    ) -> Any:
-        if key not in self._resources and set_if_unavailable:
-            await self.async_set_resource(key, default)
-
-        return self._resources.get(key, default)
-
-    async def async_set_resource(self, key: str, value: Any) -> None:
-        async with self._lock:
-            self._resources[key] = value
-
-    async def async_pop_resource(self, key: str, default: Any = None) -> Any:
-        async with self._lock:
-            return self._resources.pop(key, default)

datashare-python/datashare_python/utils.py

@@ -5,6 +5,7 @@
 import json
 import os
 import shutil
+import sys
 import threading
 import time
 from collections.abc import Callable, Coroutine, Generator, Iterable, Sequence
@@ -15,10 +16,11 @@
 from hashlib import sha256
 from io import BytesIO
 from pathlib import Path
-from typing import Any, ParamSpec, TypeVar
+from typing import Any, ParamSpec, Self, TypeVar
 from uuid import uuid4
 
 import temporalio
+from lru import LRU
 from pydantic import ValidationError
 from temporalio import activity, workflow
 from temporalio.api.common.v1 import Payload
@@ -40,7 +42,7 @@
 )
 
 from .constants import MANIFEST_JSON, METADATA_JSON
-from .objects import DocArtifact, DocumentLocation, FilesystemDocument
+from .objects import BaseModel, DocArtifact, DocumentLocation, FilesystemDocument
 from .types_ import RawAsyncProgressHandler
 
 DependencyLabel = str | None
@@ -594,3 +596,68 @@ def to_payloads(self, values: Sequence[Any]) -> list[Payload]:
 PYDANTIC_DATA_CONVERTER = DataConverter(
     payload_converter_class=_PydanticPayloadConverter
 )
+
+
+class SharedResources:
+    def __init__(
+        self,
+        cache_size: int = 1,
+        eviction_callback: Callable[[Any, Any], None] | None = None,
+    ) -> None:
+        self._eviction_callback = eviction_callback
+        self._cache = LRU(cache_size, self._eviction_callback)
+        self._sentinel = object()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:  # noqa: ANN001
+        if self._eviction_callback is not None:
+            for k, v in self._cache.items():
+                self._eviction_callback(k, v)
+
+    def get_or_cache_resource(
+        self, key: str, default_factory: Callable[[], Any]
+    ) -> Any:
+        value = self._cache.get(key, self._sentinel)
+        if value is not self._sentinel:
+            return value
+        # Get the value first to be sure to return the right one in case of concurrent
+        # access.
+        #
+        # Additionally, the factory can be long to complete (that's one of the reason
+        # we want to cache its results). Because we don't lock the factory call,
+        # together with the cache update, in case of concurrent access,
+        # quicker factories will complete first and will stay longer in the cache.
+        # This is fine, the alternative is to lock factory + cache update to avoid
+        # concurrent access but this will mean waiting for the first factory call to
+        # complete, this can potentially imply longer waits than no getting the value
+        # from the cache
+        value = default_factory()
+        self._cache[key] = value
+        return value
+
+
+def close_cm_callback(key: str, value: Any) -> None:  # noqa: ARG001
+    if hasattr(value, "__exit__"):
+        value.__exit__(*sys.exc_info())
+
+
+def enter_cm(
+    cm_factory: Callable[[], contextlib.AbstractContextManager],
+) -> Callable[[], contextlib.AbstractContextManager]:  # noqa: ARG001
+    @wraps(cm_factory)
+    def factory() -> contextlib.AbstractContextManager:
+        cm = cm_factory()
+        cm.__enter__()
+        return cm
+
+    return factory
+
+
+def config_cache_key(config: BaseModel) -> str:
+    # Frozen model offer natural hash
+    if not config.__class__.model_config["frozen"]:
+        msg = f"{config.__class__} is not frozen, can't compute cache key"
+        raise ValueError(msg)
+    return config.__hash__()

datashare-python/datashare_python/worker.py

@@ -71,7 +71,7 @@ def datashare_worker(
     task_queue: str,
     # Scale horizontally be default for activities, each worker processes one activity
     # at a time
-    max_concurrent_io_activities: int = 10,
+    max_concurrent_activities: int = 1,
     sandboxed: bool = True,
 ) -> DatashareWorker:
     if workflows is None:
@@ -91,7 +91,6 @@ def datashare_worker(
         )
         logger.warning(_SEPARATE_IO_AND_CPU_ACTIVITIES)
 
-    max_concurrent_activities = max_concurrent_io_activities
     if isinstance(activity_executor, ThreadPoolExecutor):
         max_concurrent_activities = 1
         if workflows:
@@ -204,7 +203,7 @@ async def worker_context(
             workflows=workflows,
             activities=acts,
             task_queue=task_queue,
-            max_concurrent_io_activities=worker_config.max_concurrent_io_activities,
+            max_concurrent_activities=worker_config.max_concurrent_io_activities,
             sandboxed=sandboxed,
         )
         async with worker:

datashare-python/tests/test_dependencies.py

@@ -1,96 +1,7 @@
 from typing import Any
-from unittest.mock import MagicMock
 
 import pytest
-from datashare_python.dependencies import (
-    add_missing_args,
-    component_teardown,
-    lifespan_shared_resources,
-    set_shared_resources,
-)
-from datashare_python.exceptions import DependencyInjectionError
-from datashare_python.objects import Shared
-
-
-async def test_set_shared_resources_and_get() -> None:
-    shared = Shared()
-    shared.set_resource("k", "v")
-    returned = await set_shared_resources(shared)
-    assert returned is shared
-    assert lifespan_shared_resources() is shared
-
-
-async def test_set_shared_resources_overwrites() -> None:
-    first = Shared()
-    second = Shared()
-    first.set_resource("k", "v1")
-    second.set_resource("k", "v2")
-    await set_shared_resources(first)
-    await set_shared_resources(second)
-    assert lifespan_shared_resources() is second
-
-
-def test_shared_resources_raises_before_set() -> None:
-    with pytest.raises(DependencyInjectionError):
-        lifespan_shared_resources()
-
-
-def test_shared_resources_field_is_frozen() -> None:
-    shared = Shared()
-    shared.set_resource("k", "v")
-    with pytest.raises(Exception):  # noqa: B017, PT011
-        shared._resources = {}
-
-
-def test_get_resource_existing_key() -> None:
-    shared = Shared()
-    shared.set_resource("k", "v")
-    assert shared.get_resource("k") == "v"
-
-
-def test_get_resource_missing_key_with_default() -> None:
-    shared = Shared()
-    assert shared.get_resource("missing", "default") == "default"
-
-
-async def test_async_set_resource() -> None:
-    shared = Shared()
-    await shared.async_set_resource("k", "v")
-    assert shared.get_resource("k") == "v"
-
-
-async def test_async_set_resource_overwrites() -> None:
-    shared = Shared()
-    shared.set_resource("k", "old")
-    await shared.async_set_resource("k", "new")
-    assert shared.get_resource("k") == "new"
-
-
-async def test_async_pop_resource_missing_key_with_default() -> None:
-    shared = Shared()
-    assert await shared.async_pop_resource("missing", None) is None
-
-
-def test_call_component_teardown_on_key_eviction() -> None:
-    shared = Shared(1, eviction_callback=component_teardown)
-
-    first_resource = MagicMock()
-    second_resource = MagicMock()
-
-    shared.set_resource("first", first_resource)
-    shared.set_resource("second", second_resource)
-
-    assert shared._resources.keys() == ["second"]
-
-    first_resource.__exit__.assert_called_once()
-
-
-def test_dont_set_default_when_set_is_unavailable_is_false() -> None:
-    shared = Shared()
-
-    shared.get_resource("first", default="isn't there", set_if_unavailable=False)
-
-    assert "first" not in shared._resources
+from datashare_python.dependencies import add_missing_args
 
 
 @pytest.mark.parametrize(