Skip to content

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

Open
wants to merge 6 commits into
base: main
Choose a base branch
from

Conversation

zzstoatzz
Copy link
Collaborator

@zzstoatzz zzstoatzz commented May 10, 2025

resolves #18017 where unpicklable _thread.lock in RedisLockManager caused issues when used with Ray/Dask

this PR:

  • updates __init__: Now calls a new self._init_clients() method to instantiate self.client and self.async_client immediately
  • adds _init_clients(): private method containing the client instantiation logic
  • adds __getstate__: pickle only connection parameters (host, port, etc.), excluding live clients and _locks.
  • adds __setstate__: after restoring state, now calls self._init_clients() to ensure fresh client connections are established, resets self._locks = {}.

@zzstoatzz zzstoatzz force-pushed the redis-locking-serialization branch from 09fbec2 to 7bb6ed4 Compare May 10, 2025 16:41
@zzstoatzz zzstoatzz changed the title try lazy init for redis lock manager make RedisLockManager picklable for ray/dask compat May 10, 2025
@github-actions github-actions bot added the bug Something isn't working label May 10, 2025
@zzstoatzz zzstoatzz marked this pull request as ready for review May 10, 2025 17:27
@zzstoatzz zzstoatzz force-pushed the redis-locking-serialization branch 2 times, most recently from 3ced681 to b2802e4 Compare May 10, 2025 17:35
remove unused code

typing

one more
@zzstoatzz zzstoatzz force-pushed the redis-locking-serialization branch from b2802e4 to 6108ee2 Compare May 10, 2025 17:37
@zzstoatzz zzstoatzz requested a review from Copilot May 10, 2025 18:50
Copy link

@Copilot Copilot AI left a 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(

@zzstoatzz zzstoatzz force-pushed the redis-locking-serialization branch from de0d83e to d8c44c8 Compare May 11, 2025 03:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

Successfully merging this pull request may close these issues.

RedisLockManager breaks Ray and Dask flows
2 participants