diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index ff18b1b1..44f98f67 100644 --- a/datashare-python/datashare_python/utils.py +++ b/datashare-python/datashare_python/utils.py @@ -10,6 +10,7 @@ import threading import time from collections.abc import ( + AsyncGenerator, AsyncIterable, Awaitable, Callable, @@ -766,59 +767,108 @@ def __init__( eviction_callback: Callable[[Any, Any], None] | None = None, ) -> None: self._eviction_callback = eviction_callback - self._cache = LRU(cache_size, self._eviction_callback) + self._cache = LRU(cache_size, self._on_evict) self._sentinel = object() + # Leased resources are tracked by id, evicting a leased resource defers the + # eviction callback until its last lease is released + self._lock = threading.Lock() + self._n_leases: dict[int, int] = dict() + self._pending_evictions: dict[int, tuple[Any, Any]] = dict() 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) + for k, v in self._cache.items(): + self._on_evict(k, v) + @contextlib.contextmanager 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 + ) -> Generator[Any, None, None]: + """Get or cache a resource, leasing it for the duration of the context. + + A resource evicted while leased stays usable, its eviction callback only runs + once all its leases have been released. + + :param key: resource cache key + :param default_factory: resource factory, called on cache miss + :return: the cached resource + """ + value = self._get_and_acquire(key) + if value is self._sentinel: + value = default_factory() + self._cache_and_acquire(key, value) + try: + yield value + finally: + self._release(value) + @contextlib.asynccontextmanager async def async_get_or_cache_resource( self, key: str, default_factory: Callable[[], Awaitable[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 = await default_factory() - self._cache[key] = value + ) -> AsyncGenerator[Any, None]: + """Async version of get_or_cache_resource. + + :param key: resource cache key + :param default_factory: async resource factory, awaited on cache miss + :return: the cached resource + """ + value = self._get_and_acquire(key) + if value is self._sentinel: + value = await default_factory() + self._cache_and_acquire(key, value) + try: + yield value + finally: + self._release(value) + + 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) 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) + self._n_leases[value_id] = self._n_leases.get(value_id, 0) + 1 + + def _release(self, value: Any) -> None: + value_id = id(value) + with self._lock: + self._n_leases[value_id] -= 1 + if self._n_leases[value_id]: + return + del self._n_leases[value_id] + pending = self._pending_evictions.pop(value_id, None) + if pending is not None: + self._evict(*pending) + + 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) + 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 close_cm_callback(key: str, value: Any) -> None: # noqa: ARG001 if hasattr(value, "__exit__"): diff --git a/datashare-python/tests/test_utils.py b/datashare-python/tests/test_utils.py index 2e52b52c..8d11902b 100644 --- a/datashare-python/tests/test_utils.py +++ b/datashare-python/tests/test_utils.py @@ -2,6 +2,7 @@ import fcntl import json import os +import threading import uuid from concurrent.futures import ProcessPoolExecutor from datetime import timedelta @@ -385,9 +386,9 @@ def factory() -> str: return "cached" # When - from_cache = shared.get_or_cache_resource(new_key, factory) - # Then - assert from_cache == "cached" + with shared.get_or_cache_resource(new_key, factory) as from_cache: + # Then + assert from_cache == "cached" def test_get_or_cache_resource_existing_key() -> None: @@ -399,12 +400,31 @@ def factory() -> str: return uuid.uuid4().hex # When - first = shared.get_or_cache_resource(existing_key, factory) - second = shared.get_or_cache_resource(existing_key, factory) + with shared.get_or_cache_resource(existing_key, factory) as first: + pass + with shared.get_or_cache_resource(existing_key, factory) as second: + pass # Then assert second == first +def test_get_or_cache_resource_released_resource_stays_cached() -> None: + # Given + mock = MagicMock() + eviction_callback = mock.evict + factory = MagicMock(side_effect=object) + shared = SharedResources(cache_size=1, eviction_callback=eviction_callback) + # When + with shared.get_or_cache_resource("k", factory) as first: + pass + with shared.get_or_cache_resource("k", factory) as second: + pass + # Then + assert second is first + factory.assert_called_once() + eviction_callback.assert_not_called() + + def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: # Given mock = MagicMock() @@ -413,12 +433,131 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: factory = lambda: "value" # noqa: E731 shared = SharedResources(eviction_callback=eviction_callback) # When - with shared: - shared.get_or_cache_resource(key, factory) + with shared, shared.get_or_cache_resource(key, factory): + pass # Then eviction_callback.assert_called_once_with(key, "value") +def test_get_or_cache_resource_evicted_while_leased_defers_eviction_callback() -> None: + # Given + mock = MagicMock() + eviction_callback = mock.evict + shared = SharedResources(cache_size=1, eviction_callback=eviction_callback) + # When + with shared.get_or_cache_resource("k0", lambda: "leased"): + with shared.get_or_cache_resource("k1", lambda: "other"): + pass + # Then + eviction_callback.assert_not_called() + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_get_or_cache_resource_eviction_callback_waits_for_last_lease() -> None: + # Given + mock = MagicMock() + eviction_callback = mock.evict + shared = SharedResources(cache_size=1, eviction_callback=eviction_callback) + # When + with shared.get_or_cache_resource("k0", lambda: "leased") as first: + with ( + shared.get_or_cache_resource("k0", lambda: "not_called") as second, + shared.get_or_cache_resource("k1", lambda: "other"), + ): + pass + # Then + eviction_callback.assert_not_called() + assert first == second == "leased" + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_get_or_cache_resource_released_resource_is_evicted_immediately() -> None: + # Given + mock = MagicMock() + eviction_callback = mock.evict + shared = SharedResources(cache_size=1, eviction_callback=eviction_callback) + # When + with shared.get_or_cache_resource("k0", lambda: "leased"): + pass + with shared.get_or_cache_resource("k1", lambda: "other"): + pass + # Then + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_get_or_cache_resource_exit_defers_eviction_callback_of_leased_resource() -> ( + None +): + # Given + mock = MagicMock() + eviction_callback = mock.evict + shared = SharedResources(cache_size=1, eviction_callback=eviction_callback) + lease = shared.get_or_cache_resource("k0", lambda: "leased") + # When + with shared: + lease.__enter__() + # Then + eviction_callback.assert_not_called() + lease.__exit__(None, None, None) + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_get_or_cache_resource_is_not_closed_by_concurrent_eviction() -> None: + # Given + class Resource: + def __init__(self) -> None: + self.closed = False + + shared = SharedResources( + cache_size=1, eviction_callback=lambda _, r: setattr(r, "closed", True) + ) + leased, evicted = threading.Event(), threading.Event() + closed_while_leased = [] + + def use_resource() -> None: + with shared.get_or_cache_resource("k0", Resource) as resource: + leased.set() + evicted.wait(5) + closed_while_leased.append(resource.closed) + + # When + thread = threading.Thread(target=use_resource) + thread.start() + leased.wait(5) + with shared.get_or_cache_resource("k1", Resource): + pass + evicted.set() + thread.join() + # Then + assert closed_while_leased == [False] + + +async def test_async_get_or_cache_resource_evicted_while_leased_defers_eviction() -> ( + None +): + # Given + mock = MagicMock() + eviction_callback = mock.evict + shared = SharedResources(cache_size=1, eviction_callback=eviction_callback) + + async def factory() -> str: + return "leased" + + async def other_factory() -> str: + return "other" + + # When + async with shared.async_get_or_cache_resource("k0", factory) as first: + async with shared.async_get_or_cache_resource("k0", other_factory) as second: + pass + async with shared.async_get_or_cache_resource("k1", other_factory): + pass + # Then + eviction_callback.assert_not_called() + assert first == second == "leased" + eviction_callback.assert_called_once_with("k0", "leased") + + @pytest.mark.parametrize( ("wf_context", "act_context", "run_context", "expected_work_dir"), [ diff --git a/workers/asr-worker/asr_worker/activities.py b/workers/asr-worker/asr_worker/activities.py index 6f1a8978..0b05a72b 100644 --- a/workers/asr-worker/asr_worker/activities.py +++ b/workers/asr-worker/asr_worker/activities.py @@ -152,16 +152,16 @@ def preprocess( preprocessor_factory = enter_cm(partial(Preprocessor.from_config, config)) preprocessor_key = config_cache_key(config) cache = lifespan_preprocessor_cache() - preprocessor = cache.get_or_cache_resource( + with cache.get_or_cache_resource( preprocessor_key, preprocessor_factory - ) - batch_paths = preprocess_act( - preprocessor, - audio_batch, - worker_config=worker_config, - output_dir=output_dir, - ) - batches = [p.relative_to(workdir) for p in batch_paths] + ) as preprocessor: + batch_paths = preprocess_act( + preprocessor, + audio_batch, + worker_config=worker_config, + output_dir=output_dir, + ) + batches = [p.relative_to(workdir) for p in batch_paths] return batches @activity_defn(name=RUN_INFERENCE_ACTIVITY) @@ -194,18 +194,20 @@ async def infer( ) runner_key = config_cache_key(config) cache = lifespan_inference_runner_cache() - inference_runner = cache.get_or_cache_resource(runner_key, runner_factory) - logger.info( - "model loaded, starting inference on %s audio chunks !", - len(preprocessed_inputs), - ) - inference_res = infer_act( - inference_runner, - preprocessed_inputs, - output_dir=output_dir, - progress=progress, - ) - inference_res = [p.relative_to(workdir) async for p in inference_res] + with cache.get_or_cache_resource( + runner_key, runner_factory + ) as inference_runner: + logger.info( + "model loaded, starting inference on %s audio chunks !", + len(preprocessed_inputs), + ) + inference_res = infer_act( + inference_runner, + preprocessed_inputs, + output_dir=output_dir, + progress=progress, + ) + inference_res = [p.relative_to(workdir) async for p in inference_res] return inference_res @activity_defn(name=POSTPROCESS_ACTIVITY) @@ -238,18 +240,18 @@ def postprocess( postprocessor_factory = enter_cm(partial(Postprocessor.from_config, config)) postprocessor_key = config_cache_key(config) cache = lifespan_postprocessor_cache() - postprocessor = cache.get_or_cache_resource( + with cache.get_or_cache_resource( postprocessor_key, postprocessor_factory - ) - return postprocess_act( - inference_results, - docs, - postprocessor, - args, - artifacts_root=artifacts_root, - event_loop=self._event_loop, - progress=progress, - ) + ) as postprocessor: + return postprocess_act( + inference_results, + docs, + postprocessor, + args, + artifacts_root=artifacts_root, + event_loop=self._event_loop, + progress=progress, + ) @activity_defn(name=INDEX_TRANSCRIPTION_ACTIVITY) async def index_transcriptions( diff --git a/workers/passport-worker/passport_worker/activities.py b/workers/passport-worker/passport_worker/activities.py index 21bd2eb9..a9f13336 100644 --- a/workers/passport-worker/passport_worker/activities.py +++ b/workers/passport-worker/passport_worker/activities.py @@ -115,10 +115,6 @@ def preprocess_images( image_preprocessor_factory = enter_cm( partial(ImagePreprocessor.from_config, config.images) ) - image_preprocessor = cache.get_or_cache_resource( - cache_key, image_preprocessor_factory - ) - logger.info("loaded image preprocessor !") # We cache processing at the config level, we ignore package updates, # to discard the cache we just need to disable it in the args to overwrite pages_root = activity_workdir(workdir, project, caching_key=cache_key) @@ -126,17 +122,21 @@ def preprocess_images( executor = worker_config.to_image_preprocessing_executor() chunk_size = worker_config.preprocessing.images.chunk_size force_reprocessing = not config.use_caching - success, errors = preprocess_images_act( - batch, - worker_config.paths, - force_reprocessing=force_reprocessing, - output_root=pages_root, - image_preprocessor=image_preprocessor, - executor=executor, - chunk_size=chunk_size, - event_loop=self._event_loop, - progress=progress, - ) + with cache.get_or_cache_resource( + cache_key, image_preprocessor_factory + ) as image_preprocessor: + logger.info("loaded image preprocessor !") + success, errors = preprocess_images_act( + batch, + worker_config.paths, + force_reprocessing=force_reprocessing, + output_root=pages_root, + image_preprocessor=image_preprocessor, + executor=executor, + chunk_size=chunk_size, + event_loop=self._event_loop, + progress=progress, + ) res_root = activity_workdir(workdir, project) res_root.mkdir(parents=True, exist_ok=True) successes_path = res_root / "pages.jsonl" @@ -163,23 +163,23 @@ async def convert_to_pdfs( pdf_converter_factory = async_enter_cm( partial(PDFConverter.from_config, config.pdf_converter) ) - pdf_converter = await cache.async_get_or_cache_resource( - cache_key, pdf_converter_factory - ) workdir = worker_config.paths.workdir # We cache processing at the config level, we ignore package updates, # to discard the cache we just need to disable it in the args to overwrite pdfs_root = activity_workdir(workdir, project, caching_key=cache_key) pdfs_root.mkdir(parents=True, exist_ok=True) - successes, errors = await convert_to_pdfs_act( - batch, - pdf_converter, - worker_config.paths, - config.max_concurrency, - output_root=pdfs_root, - force_reprocessing=force_reprocessing, - progress=progress, - ) + async with cache.async_get_or_cache_resource( + cache_key, pdf_converter_factory + ) as pdf_converter: + successes, errors = await convert_to_pdfs_act( + batch, + pdf_converter, + worker_config.paths, + config.max_concurrency, + output_root=pdfs_root, + force_reprocessing=force_reprocessing, + progress=progress, + ) res_root = activity_workdir(workdir, project, act_context=True) res_root.mkdir(parents=True, exist_ok=True) pdf_paths = res_root / "pdfs.jsonl" @@ -270,21 +270,21 @@ async def detect_passports( passport_detector_factory = enter_cm( partial(PassportDetector.from_config, passport_detector_config) ) - passport_detector = cache.get_or_cache_resource( - passport_detector_key, passport_detector_factory - ) - logger.info("passport detector loaded !") workdir = worker_config.paths.workdir res_root = activity_workdir(workdir, args.project, act_context=True) res_root.mkdir(parents=True, exist_ok=True) - res = await detect_passports_act( - batch, - passport_detector, - worker_config.paths, - args, - batch_size=batch_size, - progress=progress, - ) + with cache.get_or_cache_resource( + passport_detector_key, passport_detector_factory + ) as passport_detector: + logger.info("passport detector loaded !") + res = await detect_passports_act( + batch, + passport_detector, + worker_config.paths, + args, + batch_size=batch_size, + progress=progress, + ) result_path = res_root / "inference_results.json" async with async_open(result_path, "w") as f: await f.write(res.model_dump_json()) diff --git a/workers/translation-worker/translation_worker/activities.py b/workers/translation-worker/translation_worker/activities.py index b839e281..23b088c1 100644 --- a/workers/translation-worker/translation_worker/activities.py +++ b/workers/translation-worker/translation_worker/activities.py @@ -110,29 +110,29 @@ async def translate_docs( ) translator_key = config_cache_key(config.translator) translator_cache = lifespan_translator_cache() - translator = translator_cache.get_or_cache_resource( - translator_key, translator_factory - ) # SBD, then load the splitter - logger.debug("loading %s sentence splitter...", source) splitter_factory = partial( _load_splitter_from_config, config=config.sentence_splitter, language=source ) splitter_key = config_cache_key(config.sentence_splitter) splitter_cache = lifespan_sentence_splitter_cache() - sentence_splitter = splitter_cache.get_or_cache_resource( - splitter_key, splitter_factory - ) - logger.info("translating %s batches...", len(batches)) - n_translated = await translate_docs_act( - batches, - project=project, - es_client=es_client, - progress=progress, - worker_config=worker_config, - translator=translator, - sentence_splitter=sentence_splitter, - ) + with translator_cache.get_or_cache_resource( + translator_key, translator_factory + ) as translator: + logger.debug("loading %s sentence splitter...", source) + with splitter_cache.get_or_cache_resource( + splitter_key, splitter_factory + ) as sentence_splitter: + logger.info("translating %s batches...", len(batches)) + n_translated = await translate_docs_act( + batches, + project=project, + es_client=es_client, + progress=progress, + worker_config=worker_config, + translator=translator, + sentence_splitter=sentence_splitter, + ) logger.info("done translating !") return n_translated