re-land "[monarch] Pickle directly into rust buffers" after fixing typing issues (#1365)

Summary:
Pull Request resolved: #1365

there were some typing issues in the tensor engine code which this fixes.

Reviewed By: thomasywang, pablorfb-meta

Differential Revision: D83490534

fbshipit-source-id: b14ba29980ccb220d3e471b4113baf1fb4a348e3

Author: eliothedeman

monarch_extension/src/lib.rs

@@ -74,6 +74,10 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
         runtime.handle().clone(),
         Some(::hyperactor_mesh::bootstrap::BOOTSTRAP_INDEX_ENV.to_string()),
     );
+    monarch_hyperactor::buffers::register_python_bindings(&get_or_add_new_module(
+        module,
+        "monarch_hyperactor.buffers",
+    )?)?;
 
     monarch_hyperactor::shape::register_python_bindings(&get_or_add_new_module(
         module,

monarch_extension/src/mesh_controller.rs

@@ -38,6 +38,7 @@ use hyperactor_mesh::shared_cell::SharedCell;
 use hyperactor_mesh::shared_cell::SharedCellRef;
 use monarch_hyperactor::actor::PythonMessage;
 use monarch_hyperactor::actor::PythonMessageKind;
+use monarch_hyperactor::buffers::FrozenBuffer;
 use monarch_hyperactor::local_state_broker::LocalStateBrokerActor;
 use monarch_hyperactor::mailbox::PyPortId;
 use monarch_hyperactor::ndslice::PySlice;
@@ -532,8 +533,11 @@ impl History {
             let exe = remote_exception
                 .call1((exception.backtrace, traceback, rank))
                 .unwrap();
-            let data: Vec<u8> = pickle.call1((exe,)).unwrap().extract().unwrap();
-            PythonMessage::new_from_buf(PythonMessageKind::Exception { rank: Some(rank) }, data)
+            let data: FrozenBuffer = pickle.call1((exe,)).unwrap().extract().unwrap();
+            PythonMessage::new_from_buf(
+                PythonMessageKind::Exception { rank: Some(rank) },
+                data.inner,
+            )
         }));
 
         let mut invocation = invocation.lock().unwrap();

monarch_hyperactor/Cargo.toml

@@ -20,6 +20,7 @@ anyhow = "1.0.98"
 async-once-cell = "0.4.2"
 async-trait = "0.1.86"
 bincode = "1.3.3"
+bytes = { version = "1.10", features = ["serde"] }
 clap = { version = "4.5.42", features = ["derive", "env", "string", "unicode", "wrap_help"] }
 erased-serde = "0.3.27"
 fastrand = "2.1.1"
@@ -43,6 +44,7 @@ pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone
 pyo3-async-runtimes = { version = "0.24", features = ["attributes", "tokio-runtime"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_bytes = "0.11"
+serde_multipart = { version = "0.0.0", path = "../serde_multipart" }
 tempfile = "3.22"
 thiserror = "2.0.12"
 tokio = { version = "1.47.1", features = ["full", "test-util", "tracing"] }

monarch_hyperactor/src/actor.rs

@@ -13,6 +13,7 @@ use std::sync::Arc;
 use std::sync::OnceLock;
 
 use async_trait::async_trait;
+use bytes::Bytes;
 use hyperactor::Actor;
 use hyperactor::ActorHandle;
 use hyperactor::ActorId;
@@ -32,6 +33,7 @@ use monarch_types::SerializablePyErr;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyBaseException;
 use pyo3::exceptions::PyRuntimeError;
+use pyo3::exceptions::PyTypeError;
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::PyBytes;
@@ -41,12 +43,15 @@ use pyo3::types::PyType;
 use serde::Deserialize;
 use serde::Serialize;
 use serde_bytes::ByteBuf;
+use serde_multipart::Part;
 use tokio::sync::Mutex;
 use tokio::sync::mpsc::UnboundedReceiver;
 use tokio::sync::mpsc::UnboundedSender;
 use tokio::sync::oneshot;
 use tracing::Instrument;
 
+use crate::buffers::Buffer;
+use crate::buffers::FrozenBuffer;
 use crate::config::SHARED_ASYNCIO_RUNTIME;
 use crate::local_state_broker::BrokerId;
 use crate::local_state_broker::LocalStateBrokerMessage;
@@ -236,22 +241,24 @@ 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,
-    #[serde(with = "serde_bytes")]
-    pub message: Vec<u8>,
+    pub message: Part,
 }
 
 struct ResolvedCallMethod {
     method: MethodSpecifier,
-    bytes: Vec<u8>,
+    bytes: FrozenBuffer,
     local_state: PyObject,
     /// Implements PortProtocol
     /// Concretely either a Port, DroppingPort, or LocalPort
     response_port: PyObject,
 }
 
 impl PythonMessage {
-    pub fn new_from_buf(kind: PythonMessageKind, message: Vec<u8>) -> Self {
-        Self { kind, message }
+    pub fn new_from_buf(kind: PythonMessageKind, message: impl Into<Part>) -> Self {
+        Self {
+            kind,
+            message: message.into(),
+        }
     }
 
     pub fn into_rank(self, rank: usize) -> Self {
@@ -305,7 +312,9 @@ impl PythonMessage {
                     .unwrap();
                     Ok(ResolvedCallMethod {
                         method: name,
-                        bytes: self.message,
+                        bytes: FrozenBuffer {
+                            inner: self.message.into_inner(),
+                        },
                         local_state,
                         response_port,
                     })
@@ -341,7 +350,9 @@ impl PythonMessage {
                     .unbind();
                 Ok(ResolvedCallMethod {
                     method: name,
-                    bytes: self.message,
+                    bytes: FrozenBuffer {
+                        inner: self.message.into_inner(),
+                    },
                     local_state,
                     response_port,
                 })
@@ -359,7 +370,7 @@ impl std::fmt::Debug for PythonMessage {
             .field("kind", &self.kind)
             .field(
                 "message",
-                &hyperactor::data::HexFmt(self.message.as_slice()).to_string(),
+                &hyperactor::data::HexFmt(&(*self.message)[..]).to_string(),
             )
             .finish()
     }
@@ -387,8 +398,20 @@ impl Bind for PythonMessage {
 impl PythonMessage {
     #[new]
     #[pyo3(signature = (kind, message))]
-    pub fn new(kind: PythonMessageKind, message: &[u8]) -> Self {
-        PythonMessage::new_from_buf(kind, message.to_vec())
+    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()));
+        } else if let Ok(buff) = message.extract::<Bound<'py, PyBytes>>() {
+            return Ok(PythonMessage::new_from_buf(
+                kind,
+                Vec::from(buff.as_bytes()),
+            ));
+        }
+
+        Err(PyTypeError::new_err(
+            "PythonMessage(buff) takes Buffer or bytes objects only",
+        ))
     }
 
     #[getter]
@@ -397,8 +420,10 @@ impl PythonMessage {
     }
 
     #[getter]
-    fn message<'a>(&self, py: Python<'a>) -> Bound<'a, PyBytes> {
-        PyBytes::new(py, self.message.as_ref())
+    fn message(&self) -> FrozenBuffer {
+        FrozenBuffer {
+            inner: self.message.clone().into_inner(),
+        }
     }
 }
 
@@ -842,7 +867,7 @@ mod tests {
                 },
                 response_port: Some(EitherPortRef::Unbounded(port_ref.clone().into())),
             },
-            message: vec![1, 2, 3],
+            message: Part::from(vec![1, 2, 3]),
         };
         {
             let mut erased = ErasedUnbound::try_from_message(message.clone()).unwrap();