feature(datashare-python): refresh shared resources

datashare-python/datashare_python/config.py

@@ -114,6 +114,7 @@ class LoggingConfig(BaseModel):
 
 class ResourceCacheConfig(BaseModel):
     size: int = 1
+    refresh_interval: int | None = None
     exit_context_managers: bool = True
 
     def to_resource_cache(self) -> "SharedResources":
@@ -123,7 +124,9 @@ def to_resource_cache(self) -> "SharedResources":
         if self.exit_context_managers:
             eviction_callback = close_cm_callback
         return SharedResources(
-            cache_size=self.size, eviction_callback=eviction_callback
+            cache_size=self.size,
+            eviction_callback=eviction_callback,
+            refresh_interval=self.refresh_interval,
         )
 
 

datashare-python/datashare_python/utils.py

@@ -822,9 +822,11 @@ def __init__(
         self,
         cache_size: int = 1,
         eviction_callback: Callable[[Any, Any], None] | None = None,
+        refresh_interval: int | None = None,
     ) -> None:
         self._eviction_callback = eviction_callback
         self._cache = LRU(cache_size, self._on_evict)
+        self._refresh_interval = refresh_interval
         self._sentinel = object()
         # Leased resources are tracked by id, evicting a leased resource defers the
         # eviction callback until its last lease is released
@@ -854,6 +856,14 @@ def get_or_cache_resource(
         """
         value = self._get_and_acquire(key)
         if value is self._sentinel:
+            # 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_and_acquire(key, value)
         try:
@@ -873,6 +883,14 @@ async def async_get_or_cache_resource(
         """
         value = self._get_and_acquire(key)
         if value is self._sentinel:
+            # 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 = await default_factory()
             self._cache_and_acquire(key, value)
         try:
@@ -884,28 +902,31 @@ def _get_and_acquire(self, key: str) -> Any:
         with self._lock:
             value = self._cache.get(key, self._sentinel)
             if value is not self._sentinel:
-                self._acquire(value)
+                # If refresh needs to happen, we set the value to the sentinel.
+                # The old value will be release with the last user call to _release.
+                # Returning the sentinel will trigger the call to the factory
+                if (
+                    self._refresh_interval is not None
+                    and self._n_leases[self._vkey(value)] >= self._refresh_interval
+                ):
+                    self._pending_evictions[self._vkey(value)] = (key, value)
+                    value = self._sentinel
+                else:
+                    self._acquire(value)
         return value
 
     def _cache_and_acquire(self, key: str, value: Any) -> None:
-        # 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
         with self._lock:
             self._acquire(value)
         # Can evict other resources, must be called without holding the lock
         self._cache[key] = value
 
     def _acquire(self, value: Any) -> None:
-        value_id = id(value)
+        value_id = self._vkey(value)
         self._n_leases[value_id] = self._n_leases.get(value_id, 0) + 1
 
     def _release(self, value: Any) -> None:
-        value_id = id(value)
+        value_id = self._vkey(value)
         with self._lock:
             self._n_leases[value_id] -= 1
             if self._n_leases[value_id]:
@@ -917,15 +938,18 @@ def _release(self, value: Any) -> None:
 
     def _on_evict(self, key: Any, value: Any) -> None:
         with self._lock:
-            if id(value) in self._n_leases:
-                self._pending_evictions[id(value)] = (key, value)
+            if self._vkey(value) in self._n_leases:
+                self._pending_evictions[self._vkey(value)] = (key, value)
                 return
         self._evict(key, value)
 
     def _evict(self, key: Any, value: Any) -> None:
         if self._eviction_callback is not None:
             self._eviction_callback(key, value)
 
+    def _vkey(self, value: Any) -> int:
+        return id(value)
+
 
 def close_cm_callback(key: str, value: Any) -> None:  # noqa: ARG001
     if hasattr(value, "__exit__"):

datashare-python/tests/test_utils.py

@@ -8,7 +8,7 @@
 from datetime import timedelta
 from pathlib import Path
 from typing import ClassVar
-from unittest.mock import MagicMock, PropertyMock
+from unittest.mock import AsyncMock, MagicMock, PropertyMock
 
 import pytest
 import temporalio
@@ -558,6 +558,69 @@ async def other_factory() -> str:
     eviction_callback.assert_called_once_with("k0", "leased")
 
 
+async def test_eviction_callback_should_be_called_on_shared_resources_exit() -> None:
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    shared = SharedResources(cache_size=10, eviction_callback=eviction_callback)
+
+    async def factory() -> str:
+        return "leased"
+
+    # When
+    with shared, shared.get_or_cache_resource("k0", factory):
+        pass
+    eviction_callback.assert_called_once()
+
+
+async def test_get_or_cache_resource_with_refresh_should_trigger_eviction() -> None:
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    refresh_interval = 1
+    cache_size = refresh_interval + 1
+    shared = SharedResources(
+        cache_size=cache_size,
+        eviction_callback=eviction_callback,
+        refresh_interval=refresh_interval,
+    )
+    factory = MagicMock(side_effect=object)
+
+    # When
+    with shared.get_or_cache_resource("k0", factory):
+        eviction_callback.assert_not_called()
+        with shared.get_or_cache_resource("k0", factory):
+            eviction_callback.assert_not_called()
+    # Then
+    assert factory.call_count == 2
+    assert eviction_callback.call_count == 1
+
+
+async def test_async_get_or_cache_resource_with_refresh_should_trigger_eviction() -> (
+    None
+):
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    refresh_interval = 1
+    cache_size = refresh_interval + 1
+    shared = SharedResources(
+        cache_size=cache_size,
+        eviction_callback=eviction_callback,
+        refresh_interval=refresh_interval,
+    )
+    factory = AsyncMock(side_effect=object)
+
+    # When
+    async with shared.async_get_or_cache_resource("k0", factory):
+        eviction_callback.assert_not_called()
+        async with shared.async_get_or_cache_resource("k0", factory):
+            eviction_callback.assert_not_called()
+    # Then
+    assert factory.call_count == 2
+    assert eviction_callback.call_count == 1
+
+
 @pytest.mark.parametrize(
     ("wf_context", "act_context", "run_context", "expected_work_dir"),
     [

workers/passport-worker/passport_worker/inference.py

@@ -190,7 +190,7 @@ async def _count_pages(batch: Path) -> int:
 
 async def _read_images(
     batch: Path, passport_detector: PassportDetector, paths: WorkerPaths, errors: list
-) -> AsyncIterable[tuple[DatashareFile, "np.ndarray", "DetectionInputs"]]:
+) -> AsyncIterable[tuple[ProcessedPage, "np.ndarray", "DetectionInputs"]]:
     import cv2  # noqa: PLC0415
 
     for page in read_jsonl_as(paths.workdir / batch, ProcessedPage):
@@ -210,7 +210,7 @@ async def _read_images(
 
 
 async def _detect_passport_pages(
-    batch: Iterable[tuple[DatashareFile, "np.ndarray", "DetectionInputs"]],
+    batch: Iterable[tuple[ProcessedPage, "np.ndarray", "DetectionInputs"]],
     passport_detector: PassportDetector,
     *,
     read_mrz: bool,