Use PyBytes as backing for Buffer

Summary:
Currently our pickle is still not truly zero copy because the Pickler calls `Buffer::write()` which is copying bytes from `PyBytes` to `BytesMut` via `extend_from_slice()`. 

To avoid copies, we can just make `Buffer` backed by a `Vec<PyBytes>` with each call to `Buffer::write()` pushing the PyBytes to the Vec.

The following figures show a round trip produced from
```
await am.echo.call(b"x" * 1024 * 1024)
```
Before: 
Send path: 600us pickle (purple), write frames (dark green), receive frames (light green), 130 us unpickle, 
Reply path: 600us pickle, write frames, receive frames, 130us unpickle
 {F1983399784}



After:
Send path: 20us pickle (purple), write frames (dark green), receive frames (light green), 200us unpickle, 
Reply path: 20us pickle, write frames, receive frames, 150us unpickle
{F1983399596}

Differential Revision: D86696391

Author: thomasywang

monarch_hyperactor/src/actor.rs

@@ -47,11 +47,13 @@ use pyo3::types::PyType;
 use serde::Deserialize;
 use serde::Serialize;
 use serde_bytes::ByteBuf;
+use serde_multipart::FragmentedPart;
 use serde_multipart::Part;
 use tokio::sync::Mutex;
 use tokio::sync::oneshot;
 use tracing::Instrument;
 
+use crate::buffers::Buffer;
 use crate::buffers::FrozenBuffer;
 use crate::config::SHARED_ASYNCIO_RUNTIME;
 use crate::context::PyInstance;
@@ -261,7 +263,7 @@ fn mailbox<'py, T: Actor>(py: Python<'py>, cx: &Context<'_, T>) -> Bound<'py, Py
 #[derive(Clone, Serialize, Deserialize, Named, PartialEq, Default)]
 pub struct PythonMessage {
     pub kind: PythonMessageKind,
-    pub message: Part,
+    pub message: FragmentedPart,
 }
 
 struct ResolvedCallMethod {
@@ -277,7 +279,14 @@ impl PythonMessage {
     pub fn new_from_buf(kind: PythonMessageKind, message: impl Into<Part>) -> Self {
         Self {
             kind,
-            message: message.into(),
+            message: FragmentedPart::Contiguous(message.into()),
+        }
+    }
+
+    pub fn new_from_fragmented(kind: PythonMessageKind, fragmented_part: FragmentedPart) -> Self {
+        Self {
+            kind,
+            message: fragmented_part,
         }
     }
 
@@ -332,7 +341,7 @@ impl PythonMessage {
                     Ok(ResolvedCallMethod {
                         method: name,
                         bytes: FrozenBuffer {
-                            inner: self.message.into_inner(),
+                            inner: self.message.into_bytes(),
                         },
                         local_state,
                         response_port,
@@ -371,7 +380,7 @@ impl PythonMessage {
                 Ok(ResolvedCallMethod {
                     method: name,
                     bytes: FrozenBuffer {
-                        inner: self.message.into_inner(),
+                        inner: self.message.into_bytes(),
                     },
                     local_state,
                     response_port,
@@ -390,7 +399,7 @@ impl std::fmt::Debug for PythonMessage {
             .field("kind", &self.kind)
             .field(
                 "message",
-                &hyperactor::data::HexFmt(&(*self.message)[..]).to_string(),
+                &hyperactor::data::HexFmt(&(*self.message.as_bytes())[..]).to_string(),
             )
             .finish()
     }
@@ -419,9 +428,11 @@ impl PythonMessage {
     #[new]
     #[pyo3(signature = (kind, message))]
     pub fn new<'py>(kind: PythonMessageKind, message: Bound<'py, PyAny>) -> PyResult<Self> {
-        if let Ok(buff) = message.extract::<Bound<'py, FrozenBuffer>>() {
-            let frozen = buff.borrow_mut();
-            return Ok(PythonMessage::new_from_buf(kind, frozen.inner.clone()));
+        if let Ok(mut buff) = message.extract::<PyRefMut<'py, Buffer>>() {
+            return Ok(PythonMessage::new_from_fragmented(
+                kind,
+                buff.into_fragmented_part(),
+            ));
         } else if let Ok(buff) = message.extract::<Bound<'py, PyBytes>>() {
             return Ok(PythonMessage::new_from_buf(
                 kind,
@@ -442,7 +453,7 @@ impl PythonMessage {
     #[getter]
     fn message(&self) -> FrozenBuffer {
         FrozenBuffer {
-            inner: self.message.clone().into_inner(),
+            inner: self.message.as_bytes(),
         }
     }
 }
@@ -971,7 +982,7 @@ mod tests {
                 },
                 response_port: Some(EitherPortRef::Unbounded(port_ref.clone().into())),
             },
-            message: Part::from(vec![1, 2, 3]),
+            message: FragmentedPart::Contiguous(Part::from(vec![1, 2, 3])),
         };
         {
             let mut erased = ErasedUnbound::try_from_message(message.clone()).unwrap();

monarch_hyperactor/src/buffers.rs

@@ -12,86 +12,92 @@ use std::ffi::c_int;
 use std::ffi::c_void;
 
 use bytes::Buf;
-use bytes::BytesMut;
 use hyperactor::Named;
 use pyo3::buffer::PyBuffer;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
 use pyo3::types::PyBytesMethods;
 use serde::Deserialize;
 use serde::Serialize;
+use serde_multipart::FragmentedPart;
+use serde_multipart::Part;
+
+/// Wrapper that keeps Py<PyBytes> alive while allowing zero-copy access to its memory
+struct PyBytesWrapper {
+    _py_bytes: Py<PyBytes>,
+    ptr: *const u8,
+    len: usize,
+}
+
+impl PyBytesWrapper {
+    fn new(py_bytes: Py<PyBytes>) -> Self {
+        let (ptr, len) = Python::with_gil(|py| {
+            let bytes_ref = py_bytes.as_bytes(py);
+            (bytes_ref.as_ptr(), bytes_ref.len())
+        });
+        Self {
+            _py_bytes: py_bytes,
+            ptr,
+            len,
+        }
+    }
+}
+
+impl AsRef<[u8]> for PyBytesWrapper {
+    fn as_ref(&self) -> &[u8] {
+        // SAFETY: ptr is valid as long as py_bytes is alive (kept alive by Py<PyBytes>)
+        // Python won't free the memory until the Py<PyBytes> refcount reaches 0
+        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
+    }
+}
+
+// SAFETY: Py<PyBytes> is Send/Sync for immutable bytes
+unsafe impl Send for PyBytesWrapper {}
+// SAFETY: Py<PyBytes> is Send/Sync for immutable bytes
+unsafe impl Sync for PyBytesWrapper {}
 
 /// A mutable buffer for reading and writing bytes data.
 ///
-/// The `Buffer` struct provides an interface for accumulating byte data that can be written to
-/// and then frozen into an immutable `FrozenBuffer` for reading. It uses the `bytes::BytesMut`
-/// internally for efficient memory management.
+/// The `Buffer` struct provides an interface for accumulating byte data from Python `bytes` objects
+/// that can be converted into a `FragmentedPart` for zero-copy multipart message serialization.
+/// It accumulates references to Python bytes objects without copying.
 ///
 /// # Examples
 ///
 /// ```python
 /// from monarch._rust_bindings.monarch_hyperactor.buffers import Buffer
 ///
-/// # Create a new buffer with default capacity (4096 bytes)
+/// # Create a new buffer
 /// buffer = Buffer()
 ///
 /// # Write some data
 /// data = b"Hello, World!"
 /// bytes_written = buffer.write(data)
 ///
-/// # Check length
-/// print(len(buffer))  # 13
-///
-/// # Freeze for reading
-/// frozen = buffer.freeze()
-/// content = frozen.read()
+/// # Use in multipart serialization
+/// # The buffer accumulates multiple writes as separate fragments
 /// ```
 #[pyclass(subclass, module = "monarch._rust_bindings.monarch_hyperactor.buffers")]
-#[derive(Clone, Serialize, Deserialize, Named, PartialEq, Default)]
+#[derive(Clone, Default)]
 pub struct Buffer {
-    pub(crate) inner: bytes::BytesMut,
-}
-
-impl Buffer {
-    /// Consumes the Buffer and returns the underlying BytesMut.
-    /// This allows zero-copy access to the raw buffer data.
-    pub fn into_inner(self) -> bytes::BytesMut {
-        self.inner
-    }
-}
-
-impl<T> From<T> for Buffer
-where
-    T: Into<BytesMut>,
-{
-    fn from(value: T) -> Self {
-        Self {
-            inner: value.into(),
-        }
-    }
+    inner: Vec<Py<PyBytes>>,
 }
 
 #[pymethods]
 impl Buffer {
     /// Creates a new empty buffer with specified initial capacity.
     ///
-    /// # Arguments
-    /// * `size` - Initial capacity in bytes (default: 4096)
     ///
     /// # Returns
     /// A new empty `Buffer` instance with the specified capacity.
     #[new]
-    #[pyo3(signature=(size=4096))]
-    fn new(size: usize) -> Self {
-        Self {
-            inner: bytes::BytesMut::with_capacity(size),
-        }
+    fn new() -> Self {
+        Self { inner: Vec::new() }
     }
 
     /// Writes bytes data to the buffer.
     ///
-    /// Appends the provided bytes to the end of the buffer, extending its capacity
-    /// if necessary.
+    /// This keeps a reference to the Python bytes object without copying.
     ///
     /// # Arguments
     /// * `buff` - The bytes object to write to the buffer
@@ -100,26 +106,49 @@ impl Buffer {
     /// The number of bytes written (always equal to the length of input bytes)
     fn write<'py>(&mut self, buff: &Bound<'py, PyBytes>) -> usize {
         let bytes_written = buff.as_bytes().len();
-        self.inner.extend_from_slice(buff.as_bytes());
+        self.inner.push(buff.clone().unbind());
         bytes_written
     }
 
-    /// Freezes this buffer into an immutable `FrozenBuffer`.
+    /// Freezes the buffer, converting it into an immutable `FrozenBuffer` for reading.
+    ///
+    /// This consumes all accumulated PyBytes and converts them into a contiguous bytes buffer.
+    /// After freezing, the original buffer is cleared.
     ///
-    /// This operation consumes the mutable buffer's contents, transferring ownership
-    /// to a new `FrozenBuffer` that can only be read from. The original buffer
-    /// becomes empty after this operation.
+    /// This operation should avoided in hot paths as it creates a copy in order to concatenate
+    /// bytes that are fragmented in memory into a single series of contiguous bytes
     ///
     /// # Returns
-    /// A new `FrozenBuffer` containing all the data that was in this buffer
+    /// A new `FrozenBuffer` containing all the bytes that were written to this buffer
     fn freeze(&mut self) -> FrozenBuffer {
-        let buff = std::mem::take(&mut self.inner);
+        let fragmented_part = self.into_fragmented_part();
         FrozenBuffer {
-            inner: buff.freeze(),
+            inner: fragmented_part.into_bytes(),
         }
     }
 }
 
+impl Buffer {
+    /// Converts accumulated `PyBytes` objects to [`FragmentedPart`] for zero-copy multipart messages.
+    ///
+    /// Returns a `FragmentedPart::Fragmented` variant since the buffer accumulates multiple
+    /// separate PyBytes objects that remain physically fragmented.
+    pub fn into_fragmented_part(&mut self) -> FragmentedPart {
+        let inner = std::mem::take(&mut self.inner);
+
+        FragmentedPart::Fragmented(
+            inner
+                .into_iter()
+                .map(|py_bytes| {
+                    let wrapper = PyBytesWrapper::new(py_bytes);
+                    let bytes = bytes::Bytes::from_owner(wrapper);
+                    Part::from(bytes)
+                })
+                .collect(),
+        )
+    }
+}
+
 /// An immutable buffer for reading bytes data.
 ///
 /// The `FrozenBuffer` struct provides a read-only interface to byte data. Once created,

python/monarch/_rust_bindings/monarch_hyperactor/actor.pyi

@@ -23,7 +23,7 @@ from typing import (
     Union,
 )
 
-from monarch._rust_bindings.monarch_hyperactor.buffers import FrozenBuffer
+from monarch._rust_bindings.monarch_hyperactor.buffers import Buffer, FrozenBuffer
 
 from monarch._rust_bindings.monarch_hyperactor.mailbox import (
     Mailbox,
@@ -204,7 +204,7 @@ class PythonMessage:
     def __init__(
         self,
         kind: PythonMessageKind,
-        message: Union[FrozenBuffer, bytes],
+        message: Union[Buffer, bytes],
     ) -> None: ...
     @property
     def message(self) -> FrozenBuffer:

python/monarch/_rust_bindings/monarch_hyperactor/buffers.pyi

@@ -130,45 +130,37 @@ class Buffer:
     """
     A mutable buffer for reading and writing bytes data.
 
-    The `Buffer` struct provides an interface for accumulating byte data that can be written to
-    and then frozen into an immutable `FrozenBuffer` for reading. It uses the `bytes::BytesMut`
-    internally for efficient memory management.
+    The `Buffer` struct provides an interface for accumulating byte data from Python `bytes` objects
+    that can be converted into a `FragmentedPart` for zero-copy multipart message serialization.
+    It accumulates references to Python bytes objects without copying.
 
     Examples:
         ```python
         from monarch._rust_bindings.monarch_hyperactor.buffers import Buffer
 
-        # Create a new buffer with default capacity (4096 bytes)
+        # Create a new buffer
         buffer = Buffer()
 
         # Write some data
         data = b"Hello, World!"
         bytes_written = buffer.write(data)
 
-        # Check length
-        print(len(buffer))  # 13
-
-        # Freeze for reading
-        frozen = buffer.freeze()
-        content = frozen.read()
+        # Use in multipart serialization
+        # The buffer accumulates multiple writes as separate fragments
         ```
     """
 
-    def __init__(self, size: int = 4096) -> None:
+    def __init__(self) -> None:
         """
-        Create a new empty buffer with specified initial capacity.
-
-        Arguments:
-        - `size`: Initial capacity in bytes (default: 4096)
+        Create a new empty buffer.
         """
         ...
 
     def write(self, buff: bytes) -> int:
         """
         Write bytes data to the buffer.
 
-        Appends the provided bytes to the end of the buffer, extending its capacity
-        if necessary.
+        This keeps a reference to the Python bytes object without copying.
 
         Arguments:
         - `buff`: The bytes object to write to the buffer
@@ -178,15 +170,6 @@ class Buffer:
         """
         ...
 
-    def __len__(self) -> int:
-        """
-        Return the number of bytes remaining in the buffer.
-
-        Returns:
-        The number of bytes that can be read from the buffer
-        """
-        ...
-
     def freeze(self) -> FrozenBuffer:
         """
         Freeze this buffer into an immutable `FrozenBuffer`.
@@ -195,6 +178,9 @@ class Buffer:
         to a new `FrozenBuffer` that can only be read from. The original buffer
         becomes empty after this operation.
 
+        This operation should generally be avoided in hot paths as it creates copies in order to concatenate
+        bytes that are potentially fragmented in memory into a single contiguous series of bytes
+
         Returns:
         A new `FrozenBuffer` containing all the data that was in this buffer
         """