Repository navigation
chore(asr-worker; datashare-python): refactor caching to prevent premature evictions - #101
Conversation
ClemDoum
left a comment
There was a problem hiding this comment.
Thanks for spotting the bug and for the fix.
I think we should make your implementation the default for get_or_cache_resource so that it's used everywhere.
The only reservation I have is about the model being dropped from the cache after each inference (if I understand correctly how the lease context manager exits).
| for k, v in self._cache.items(): | ||
| self._on_evict(k, v) | ||
|
|
||
| @contextlib.contextmanager |
There was a problem hiding this comment.
I feel like we'll always want to deal with thread concurrency/safety, so we can probably just make this one the default implem of get_or_cache_resource
| inputs: Iterable[tuple[ProcessedAudioSegment, ...]], | ||
| ) -> list[ASRResult]: | ||
| return list(inference_runner.process(inputs)) | ||
| with lease_runner() as inference_runner: |
There was a problem hiding this comment.
If I understand correctly, we'll exit the context manager and pop the value from the cache here.
This can potentially unload the model from the cache even it the next activity uses the same model and we'll have to reload it. I'm not sure it's fully equivalent to the previous behavior, but might be missing something. ?
There was a problem hiding this comment.
Is there a scenario where an activity would reuse a model? So far only inference requires them.
There was a problem hiding this comment.
You mean a different activity ?
No, only inference uses models. When I said next activity I was thinkint about new data coming to the same worker (performing only inference) with the same model/config. The goal of the shared resource was to avoid reloading the model from the filesystem when this second batch of data (and subsequent ones) comes in.
There was a problem hiding this comment.
Exiting the lease_resource context doesn't exit the resource or pop it from the cache. It only decrements the resource's lease count (_release). The value stays in the LRU, so the next activity with the same model/config gets a cache hit and the model isn't reloaded.
exit only runs when the LRU itself evicts the entry, which happens when a different key is inserted into the full cache, or when the cache exits at worker shutdown (like before). The difference is that if the evicted runner is still in use, the callback waits until its last lease is released instead of running immediately.
There was a problem hiding this comment.
OK I get it thanks for detailing. I thought eviction was systematic in _release. I missed the part about inspecting pending evictions (pending = self._pending_evictions.pop(value_id, None)), so I thought eviction would occur anytime the counter would reach 0. All good then.
Maybe 2 other questions.
I'm not sure to understand your comment 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. If an activity is cancelled / terminated, we expect what's offloaded to to_thread to be cancelled too don't we ?
What prevents from leasing higher in the stack in the main inference activity thread in infer rather than lower in the stack ? I'm asking because that would allow to preserve a inference_runner: InferenceRunner rather than a factory as argument.
There was a problem hiding this comment.
-
Don't threads that have already started keep running anyway? I thought there was no way to stop a running thread in Python. We can only stop work that hasn't started yet.
-
I don't think anything does. Only thing I can think of is that in this version, the model is loaded in the inference runner thread, which doesn't block the event loop. If we lease in
infer, we go back to blocking, and we do risk a heartbeat timeout for large-model loads (though it would have to be a significant delay I grant you). But I can make the change.
There was a problem hiding this comment.
-
agreed but I don't get the part where you say that we must keep the runner alive until it completes. If the activity gets cancelled with a
asyncio.CancelledError, we no longer care about what's running in theto_thread, so we don't need to keep the runner alive or let it complete right? Unless we want to gracefully shutdown but I don't really see how -
agreed model loading could take some time even if it should be lower than the heartbeat timeout. I'm thinking about an alternative, couldn't we do something like:
@asynccontextmanager
async def to_thread_cm(cm):
try:
yield await asyncio.to_thread(cm.__enter__)
finally:
cm.__exit__(.......)and then at the top level of infer do something like
async with to_thread_cm(llease_runner()) as runner:
inference_res = infer_act(
runner,
preprocessed_inputs,
output_dir=output_dir,
progress=progress,
...)this way activitiy function can keep native arguments ?
I think we'll generalize what you did here to all activities using cache object, it would be nice to keep native arguments.
There was a problem hiding this comment.
Or maybe add an async content manager API to lease_resource which can be used in async activites and acts like the to_thread_cm wrapper up there.
There was a problem hiding this comment.
In the next revision, the lease wraps the inference call in infer, and infer_act / _transcribe_as_list take an InferenceRunner again. But canceling the activity still won't drop the thread. It'll release the lease, and if the runner isn't evicted, the thread will run to completion. If the runner is evicted, the orphaned thread will fail on the next batch.
When
SharedResourcesevicts a resource, it calls the eviction callback (close_cm_callback, which calls the resource's__exit__) straight away, even if another activity is still using that resource. The inference runner cache is sized for a single runner; when an infer activity arrives with a different inference config, it evicts the cached runner and closes it while another thread is still insiderunner.process(), as when a worker runs more than one activity at a time, or an infer activity is cancelled or times out.lease_resource(key, factory)now wraps calls toget_or_cache_resource, counting how many callers are using it. If a resource is evicted while in use or the cache exits at worker shutdown, the eviction callback waits until the last caller releases it. Resources