Skip to content

Comments

[monarch] Refactor pickling to use thread-local Rust storage#2704

Open
zdevito wants to merge 4 commits intogh/zdevito/172/basefrom
gh/zdevito/172/head
Open

[monarch] Refactor pickling to use thread-local Rust storage#2704
zdevito wants to merge 4 commits intogh/zdevito/172/basefrom
gh/zdevito/172/head

Conversation

@zdevito
Copy link
Contributor

@zdevito zdevito commented Feb 20, 2026

Stack from ghstack (oldest at bottom):

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

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 20, 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: 342964837
Pull Request resolved: #2704
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Feb 20, 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-poisoned]
zdevito added a commit that referenced this pull request Feb 20, 2026
Pull Request resolved: #2704

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
ghstack-source-id: 342981968
@exported-using-ghexport

Differential Revision: [D92435072](https://our.internmc.facebook.com/intern/diff/D92435072/)
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 20, 2026
Pull Request resolved: #2704

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
ghstack-source-id: 343309982
@exported-using-ghexport

Differential Revision: [D92435072](https://our.internmc.facebook.com/intern/diff/D92435072/)
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 20, 2026
Pull Request resolved: #2704

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
ghstack-source-id: 343309982
@exported-using-ghexport

Differential Revision: [D92435072](https://our.internmc.facebook.com/intern/diff/D92435072/)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot. fb-exported meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant