-
Notifications
You must be signed in to change notification settings - Fork 1.8k
make RedisLockManager
picklable for ray/dask compat
#18018
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
09fbec2
to
7bb6ed4
Compare
RedisLockManager
picklable for ray/dask compat
3ced681
to
b2802e4
Compare
remove unused code typing one more
b2802e4
to
6108ee2
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR introduces support for pickling in the RedisLockManager so that it can be used with Ray/Dask. It does so by refactoring client initialization into a dedicated _init_clients method, adding getstate/setstate to exclude live connections and reset locks, and extending the test suite to verify both synchronous and asynchronous lock behavior post-unpickle.
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
File | Description |
---|---|
src/integrations/prefect-redis/tests/test_locking.py | Adds tests to ensure correct reinitialization of RedisLockManager clients and lock behavior after pickling/unpickling. |
src/integrations/prefect-redis/prefect_redis/locking.py | Refactors client initialization into _init_clients, implements getstate/setstate, and updates lock acquisition logic for asynchronous operations. |
Comments suppressed due to low confidence (1)
src/integrations/prefect-redis/prefect_redis/locking.py:147
- [nitpick] Consider renaming 'aacquire_lock' to 'acquire_lock_async' for improved clarity and consistency with typical asynchronous method naming conventions.
async def aacquire_lock(
update inaccurate comment
de0d83e
to
d8c44c8
Compare
resolves #18017 where unpicklable
_thread.lock
inRedisLockManager
caused issues when used with Ray/Daskthis PR:
__init__
: Now calls a newself._init_clients()
method to instantiateself.client
andself.async_client
immediately_init_clients()
: private method containing the client instantiation logic__getstate__
: pickle only connection parameters (host, port, etc.), excluding live clients and_locks
.__setstate__
: after restoring state, now callsself._init_clients()
to ensure fresh client connections are established, resetsself._locks = {}
.