From 30e0af1d6e80f78679d87498e89ae539acc739d9 Mon Sep 17 00:00:00 2001 From: Lion Summerbell Date: Thu, 24 Sep 2026 11:00:47 +0200 Subject: [PATCH 1/2] chore(asr-worker; datashare-python): refactor caching to prevent premature evictions --- datashare-python/datashare_python/utils.py | 64 ++++++++++++++- datashare-python/tests/test_utils.py | 86 +++++++++++++++++++++ workers/asr-worker/asr_worker/activities.py | 24 +++--- workers/asr-worker/tests/test_activities.py | 5 +- 4 files changed, 166 insertions(+), 13 deletions(-) diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index ff18b1b1..0a5b3f2a 100644 --- a/datashare-python/datashare_python/utils.py +++ b/datashare-python/datashare_python/utils.py @@ -766,20 +766,78 @@ 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 + for k, v in self._cache.items(): + self._on_evict(k, v) + + @contextlib.contextmanager + def lease_resource( + self, key: str, default_factory: Callable[[], Any] + ) -> Generator[Any, None, None]: + """Get or cache a resource, guaranteeing it won't be evicted while in use. + + Unlike get_or_cache_resource, the eviction callback of a resource evicted + while leased only runs once all its leases have been released. + """ + with self._lock: + value = self._cache.get(key, self._sentinel) + if value is not self._sentinel: + self._acquire(value) + if value is self._sentinel: + # Same as get_or_cache_resource, the factory is called without locking + value = default_factory() + with self._lock: + self._acquire(value) + # Can evict other resources, must be called without holding the lock + self._cache[key] = value + try: + yield value + finally: + self._release(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: - for k, v in self._cache.items(): - self._eviction_callback(k, v) + self._eviction_callback(key, value) def get_or_cache_resource( self, key: str, default_factory: Callable[[], Any] ) -> Any: + # The returned resource can be evicted and closed by the eviction callback + # while still in use. lease_resource is preferred when used concurrently or + # from another thread value = self._cache.get(key, self._sentinel) if value is not self._sentinel: return value diff --git a/datashare-python/tests/test_utils.py b/datashare-python/tests/test_utils.py index 2e52b52c..b4ea4865 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 @@ -419,6 +420,91 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: eviction_callback.assert_called_once_with(key, "value") +def test_lease_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.lease_resource("k0", lambda: "leased"): + shared.get_or_cache_resource("k1", lambda: "other") + # Then + eviction_callback.assert_not_called() + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_lease_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.lease_resource("k0", lambda: "leased") as first: + with shared.lease_resource("k0", lambda: "not_called") as second: + shared.get_or_cache_resource("k1", lambda: "other") + # Then + eviction_callback.assert_not_called() + assert first == second == "leased" + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_lease_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.lease_resource("k0", lambda: "leased"): + pass + shared.get_or_cache_resource("k1", lambda: "other") + # Then + eviction_callback.assert_called_once_with("k0", "leased") + + +def test_lease_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.lease_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_lease_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.lease_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) + shared.get_or_cache_resource("k1", Resource) + evicted.set() + thread.join() + # Then + assert closed_while_leased == [False] + + @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..c6ae8f1a 100644 --- a/workers/asr-worker/asr_worker/activities.py +++ b/workers/asr-worker/asr_worker/activities.py @@ -1,7 +1,8 @@ import asyncio import logging from asyncio import AbstractEventLoop -from collections.abc import AsyncGenerator, AsyncIterable, Iterable +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable +from contextlib import AbstractContextManager from functools import partial from itertools import tee from pathlib import Path @@ -188,19 +189,21 @@ async def infer( progress, max_progress=len(preprocessed_inputs) ) device = worker_config.devices.inference - logger.info("loading model %s on %s device", config.model, device) runner_factory = enter_cm( partial(InferenceRunner.from_config, config, device=device) ) runner_key = config_cache_key(config) cache = lifespan_inference_runner_cache() - inference_runner = cache.get_or_cache_resource(runner_key, runner_factory) + # Lease the runner so that it can't be evicted while still in use + lease_runner = partial(cache.lease_resource, runner_key, runner_factory) logger.info( - "model loaded, starting inference on %s audio chunks !", + "starting inference with %s on %s device on %s audio chunks !", + config.model, + device, len(preprocessed_inputs), ) inference_res = infer_act( - inference_runner, + lease_runner, preprocessed_inputs, output_dir=output_dir, progress=progress, @@ -321,7 +324,7 @@ def preprocess_act( async def infer_act( - inference_runner: InferenceRunner, + lease_runner: Callable[[], AbstractContextManager[InferenceRunner]], preprocessed_inputs: list[Path], output_dir: Path, event_loop: AbstractEventLoop | None = None, @@ -337,8 +340,10 @@ async def infer_act( audio_paths, inputs = tee(inputs) audio_paths = (i.path for b in audio_paths for i in b) # TODO: implement caching + # The runner is leased from the inference thread. If the activity is cancelled, + # the thread keeps running and must keep the runner alive until it completes inference_results = await asyncio.to_thread( - _transcribe_as_list, inference_runner, list(inputs) + _transcribe_as_list, lease_runner, list(inputs) ) for res_i, (path, asr_res) in enumerate( zip(audio_paths, inference_results, strict=True) @@ -356,10 +361,11 @@ async def infer_act( def _transcribe_as_list( - inference_runner: InferenceRunner, + lease_runner: Callable[[], AbstractContextManager[InferenceRunner]], inputs: Iterable[tuple[ProcessedAudioSegment, ...]], ) -> list[ASRResult]: - return list(inference_runner.process(inputs)) + with lease_runner() as inference_runner: + return list(inference_runner.process(inputs)) def postprocess_act( diff --git a/workers/asr-worker/tests/test_activities.py b/workers/asr-worker/tests/test_activities.py index 3ff0d476..9386cee5 100644 --- a/workers/asr-worker/tests/test_activities.py +++ b/workers/asr-worker/tests/test_activities.py @@ -1,5 +1,6 @@ import json from collections.abc import AsyncGenerator, Iterable +from contextlib import nullcontext from functools import partial from itertools import cycle from pathlib import Path @@ -350,7 +351,9 @@ async def test_infer_act(tmpdir: Path) -> None: paths.append(input_path) # When asr_result_paths = infer_act( - inference_runner, preprocessed_inputs=paths, output_dir=output_dir + partial(nullcontext, inference_runner), + preprocessed_inputs=paths, + output_dir=output_dir, ) # Then asr_results = [ From 7035e0178299587a6b85940f9cc0216b1afee4d1 Mon Sep 17 00:00:00 2001 From: Lion Summerbell Date: Fri, 25 Sep 2026 09:14:42 +0200 Subject: [PATCH 2/2] revert to lease wrapping inference call in infer --- datashare-python/datashare_python/utils.py | 108 ++++++++---------- datashare-python/tests/test_utils.py | 97 ++++++++++++---- workers/asr-worker/asr_worker/activities.py | 84 +++++++------- workers/asr-worker/tests/test_activities.py | 5 +- .../passport_worker/activities.py | 78 ++++++------- .../translation_worker/activities.py | 34 +++--- 6 files changed, 222 insertions(+), 184 deletions(-) diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index 0a5b3f2a..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, @@ -782,30 +783,66 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None: # noqa: ANN001 self._on_evict(k, v) @contextlib.contextmanager - def lease_resource( + def get_or_cache_resource( self, key: str, default_factory: Callable[[], Any] ) -> Generator[Any, None, None]: - """Get or cache a resource, guaranteeing it won't be evicted while in use. + """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. - Unlike get_or_cache_resource, the eviction callback of a resource evicted - while leased 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 """ - with self._lock: - value = self._cache.get(key, self._sentinel) - if value is not self._sentinel: - self._acquire(value) + value = self._get_and_acquire(key) if value is self._sentinel: - # Same as get_or_cache_resource, the factory is called without locking value = default_factory() - with self._lock: - self._acquire(value) - # Can evict other resources, must be called without holding the lock - self._cache[key] = value + 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]] + ) -> 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 @@ -832,51 +869,6 @@ def _evict(self, key: Any, value: Any) -> None: if self._eviction_callback is not None: self._eviction_callback(key, value) - def get_or_cache_resource( - self, key: str, default_factory: Callable[[], Any] - ) -> Any: - # The returned resource can be evicted and closed by the eviction callback - # while still in use. lease_resource is preferred when used concurrently or - # from another thread - 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 - - 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 - return 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 b4ea4865..8d11902b 100644 --- a/datashare-python/tests/test_utils.py +++ b/datashare-python/tests/test_utils.py @@ -386,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: @@ -400,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() @@ -414,59 +433,66 @@ 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_lease_resource_evicted_while_leased_defers_eviction_callback() -> None: +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.lease_resource("k0", lambda: "leased"): - shared.get_or_cache_resource("k1", lambda: "other") + 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_lease_resource_eviction_callback_waits_for_last_lease() -> None: +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.lease_resource("k0", lambda: "leased") as first: - with shared.lease_resource("k0", lambda: "not_called") as second: - shared.get_or_cache_resource("k1", lambda: "other") + 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_lease_resource_released_resource_is_evicted_immediately() -> None: +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.lease_resource("k0", lambda: "leased"): + with shared.get_or_cache_resource("k0", lambda: "leased"): + pass + with shared.get_or_cache_resource("k1", lambda: "other"): pass - shared.get_or_cache_resource("k1", lambda: "other") # Then eviction_callback.assert_called_once_with("k0", "leased") -def test_lease_resource_exit_defers_eviction_callback_of_leased_resource() -> None: +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.lease_resource("k0", lambda: "leased") + lease = shared.get_or_cache_resource("k0", lambda: "leased") # When with shared: lease.__enter__() @@ -476,7 +502,7 @@ def test_lease_resource_exit_defers_eviction_callback_of_leased_resource() -> No eviction_callback.assert_called_once_with("k0", "leased") -def test_lease_resource_is_not_closed_by_concurrent_eviction() -> None: +def test_get_or_cache_resource_is_not_closed_by_concurrent_eviction() -> None: # Given class Resource: def __init__(self) -> None: @@ -489,7 +515,7 @@ def __init__(self) -> None: closed_while_leased = [] def use_resource() -> None: - with shared.lease_resource("k0", Resource) as resource: + with shared.get_or_cache_resource("k0", Resource) as resource: leased.set() evicted.wait(5) closed_while_leased.append(resource.closed) @@ -498,13 +524,40 @@ def use_resource() -> None: thread = threading.Thread(target=use_resource) thread.start() leased.wait(5) - shared.get_or_cache_resource("k1", Resource) + 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 c6ae8f1a..0b05a72b 100644 --- a/workers/asr-worker/asr_worker/activities.py +++ b/workers/asr-worker/asr_worker/activities.py @@ -1,8 +1,7 @@ import asyncio import logging from asyncio import AbstractEventLoop -from collections.abc import AsyncGenerator, AsyncIterable, Callable, Iterable -from contextlib import AbstractContextManager +from collections.abc import AsyncGenerator, AsyncIterable, Iterable from functools import partial from itertools import tee from pathlib import Path @@ -153,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) @@ -189,26 +188,26 @@ async def infer( progress, max_progress=len(preprocessed_inputs) ) device = worker_config.devices.inference + logger.info("loading model %s on %s device", config.model, device) runner_factory = enter_cm( partial(InferenceRunner.from_config, config, device=device) ) runner_key = config_cache_key(config) cache = lifespan_inference_runner_cache() - # Lease the runner so that it can't be evicted while still in use - lease_runner = partial(cache.lease_resource, runner_key, runner_factory) - logger.info( - "starting inference with %s on %s device on %s audio chunks !", - config.model, - device, - len(preprocessed_inputs), - ) - inference_res = infer_act( - lease_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) @@ -241,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( @@ -324,7 +323,7 @@ def preprocess_act( async def infer_act( - lease_runner: Callable[[], AbstractContextManager[InferenceRunner]], + inference_runner: InferenceRunner, preprocessed_inputs: list[Path], output_dir: Path, event_loop: AbstractEventLoop | None = None, @@ -340,10 +339,8 @@ async def infer_act( audio_paths, inputs = tee(inputs) audio_paths = (i.path for b in audio_paths for i in b) # TODO: implement caching - # The runner is leased from the inference thread. If the activity is cancelled, - # the thread keeps running and must keep the runner alive until it completes inference_results = await asyncio.to_thread( - _transcribe_as_list, lease_runner, list(inputs) + _transcribe_as_list, inference_runner, list(inputs) ) for res_i, (path, asr_res) in enumerate( zip(audio_paths, inference_results, strict=True) @@ -361,11 +358,10 @@ async def infer_act( def _transcribe_as_list( - lease_runner: Callable[[], AbstractContextManager[InferenceRunner]], + inference_runner: InferenceRunner, inputs: Iterable[tuple[ProcessedAudioSegment, ...]], ) -> list[ASRResult]: - with lease_runner() as inference_runner: - return list(inference_runner.process(inputs)) + return list(inference_runner.process(inputs)) def postprocess_act( diff --git a/workers/asr-worker/tests/test_activities.py b/workers/asr-worker/tests/test_activities.py index 9386cee5..3ff0d476 100644 --- a/workers/asr-worker/tests/test_activities.py +++ b/workers/asr-worker/tests/test_activities.py @@ -1,6 +1,5 @@ import json from collections.abc import AsyncGenerator, Iterable -from contextlib import nullcontext from functools import partial from itertools import cycle from pathlib import Path @@ -351,9 +350,7 @@ async def test_infer_act(tmpdir: Path) -> None: paths.append(input_path) # When asr_result_paths = infer_act( - partial(nullcontext, inference_runner), - preprocessed_inputs=paths, - output_dir=output_dir, + inference_runner, preprocessed_inputs=paths, output_dir=output_dir ) # Then asr_results = [ 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