[monarch] Refactor pickling to use thread-local Rust storage#2718
Open
zdevito wants to merge 1 commit intogh/zdevito/174/basefrom
Open
[monarch] Refactor pickling to use thread-local Rust storage#2718zdevito wants to merge 1 commit intogh/zdevito/174/basefrom
zdevito wants to merge 1 commit intogh/zdevito/174/basefrom
Conversation
This diff refactors Monarch's pickling system by moving from a Python-based `persistent_id`/`flatten`/`unflatten` approach to a simpler Rust-based thread-local storage mechanism. **Key Changes:** 1. **New `pickle.rs` module** - Introduces thread-local `ACTIVE_PICKLING_STATE` storage for tracking out-of-band pickling information during cloudpickle operations. Provides `PicklingState`, `PendingMessage`, and `pickle()` function. 2. **Simplified `PythonMessage`** - Removed `pending_pickle_state` field entirely. Constructor now takes `FrozenBuffer` directly instead of `Buffer | bytes`. 3. **Removed mailbox handling from references** - The `local_state` for message dispatch changed from `itertools.repeat(mailbox)` to an empty list. Mailboxes are no longer passed through this mechanism. 4. **Deleted `PendingPickle` and `PendingPickleState`** from `pytokio.rs` - These Python-side classes handled deferred pickling via `flatten`/`unflatten`. Replaced by Rust-side `PicklingState.resolve()` and `PendingMessage.resolve()`. 5. **`PyShared` now has `__reduce__`** - Added pickle protocol support directly via `reduce_shared()`. Also optimized `block_on` to check if value is already available before calling into tokio runtime. 6. **New `cast_unresolved()` method** - Trait method for casting messages with unresolved async values. `AsyncActorMesh` provides async implementation. 7. **Python-side simplifications** - Removed helper functions (`_is_mailbox`, `_flatten_with_pending_pickle`, `_pickle`), `_SingletonActorAdapator` class, and `allow_pending_pickle_mesh()` context manager usage. **Benefits:** - Reduced Python overhead: No more Python-side `persistent_id` callbacks or `flatten`/`unflatten` traversals during pickling - Cleaner architecture: Pickling state handled via thread-local Rust storage that `__reduce__` implementations can access directly - Simplified message type: `PythonMessage` no longer carries pending pickle state - ~200 lines of Python removed, ~175 lines of Rust removed from pytokio.rs Differential Revision: [D92435072](https://our.internmc.facebook.com/intern/diff/D92435072/) [ghstack-poisoned]
zdevito
added a commit
that referenced
this pull request
Feb 21, 2026
This diff refactors Monarch's pickling system by moving from a Python-based `persistent_id`/`flatten`/`unflatten` approach to a simpler Rust-based thread-local storage mechanism. **Key Changes:** 1. **New `pickle.rs` module** - Introduces thread-local `ACTIVE_PICKLING_STATE` storage for tracking out-of-band pickling information during cloudpickle operations. Provides `PicklingState`, `PendingMessage`, and `pickle()` function. 2. **Simplified `PythonMessage`** - Removed `pending_pickle_state` field entirely. Constructor now takes `FrozenBuffer` directly instead of `Buffer | bytes`. 3. **Removed mailbox handling from references** - The `local_state` for message dispatch changed from `itertools.repeat(mailbox)` to an empty list. Mailboxes are no longer passed through this mechanism. 4. **Deleted `PendingPickle` and `PendingPickleState`** from `pytokio.rs` - These Python-side classes handled deferred pickling via `flatten`/`unflatten`. Replaced by Rust-side `PicklingState.resolve()` and `PendingMessage.resolve()`. 5. **`PyShared` now has `__reduce__`** - Added pickle protocol support directly via `reduce_shared()`. Also optimized `block_on` to check if value is already available before calling into tokio runtime. 6. **New `cast_unresolved()` method** - Trait method for casting messages with unresolved async values. `AsyncActorMesh` provides async implementation. 7. **Python-side simplifications** - Removed helper functions (`_is_mailbox`, `_flatten_with_pending_pickle`, `_pickle`), `_SingletonActorAdapator` class, and `allow_pending_pickle_mesh()` context manager usage. **Benefits:** - Reduced Python overhead: No more Python-side `persistent_id` callbacks or `flatten`/`unflatten` traversals during pickling - Cleaner architecture: Pickling state handled via thread-local Rust storage that `__reduce__` implementations can access directly - Simplified message type: `PythonMessage` no longer carries pending pickle state - ~200 lines of Python removed, ~175 lines of Rust removed from pytokio.rs Differential Revision: [D92435072](https://our.internmc.facebook.com/intern/diff/D92435072/) ghstack-source-id: 343347059 Pull Request resolved: #2718
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
This diff refactors Monarch's pickling system by moving from a Python-based
persistent_id/flatten/unflattenapproach to a simpler Rust-basedthread-local storage mechanism.
Key Changes:
New
pickle.rsmodule - Introduces thread-localACTIVE_PICKLING_STATEstorage for tracking out-of-band pickling information during cloudpickle
operations. Provides
PicklingState,PendingMessage, andpickle()function.Simplified
PythonMessage- Removedpending_pickle_statefield entirely.Constructor now takes
FrozenBufferdirectly instead ofBuffer | bytes.Removed mailbox handling from references - The
local_statefor messagedispatch changed from
itertools.repeat(mailbox)to an empty list. Mailboxesare no longer passed through this mechanism.
Deleted
PendingPickleandPendingPickleStatefrompytokio.rs- ThesePython-side classes handled deferred pickling via
flatten/unflatten.Replaced by Rust-side
PicklingState.resolve()andPendingMessage.resolve().PySharednow has__reduce__- Added pickle protocol support directlyvia
reduce_shared(). Also optimizedblock_onto check if value is alreadyavailable before calling into tokio runtime.
New
cast_unresolved()method - Trait method for casting messages withunresolved async values.
AsyncActorMeshprovides async implementation.Python-side simplifications - Removed helper functions (
_is_mailbox,_flatten_with_pending_pickle,_pickle),_SingletonActorAdapatorclass,and
allow_pending_pickle_mesh()context manager usage.Benefits:
persistent_idcallbacks orflatten/unflattentraversals during picklingthat
__reduce__implementations can access directlyPythonMessageno longer carries pending pickle stateDifferential Revision: D92435072