Change accumulator releated signatures (#216)

Summary:
Pull Request resolved: #216

This diff is a prep diff for D76156297. Specifically, it does the following:

1. Add a `struct ReducerSpec` to replace the `64` reducer typehash. The main purpose of this struct is to add a new `builder_params` field, so we can argument the accumulator, with it, serialize it, and then send to comm actor, so comm actor can instantiate the reducer with this info.
2. Make some methods in `trait Accumulator` and `trait CommReducer` fallible. This is required by the python accumulator, because it does pickling and thus fallible.
3. Add a `typehash_f` field to `ReducerFactory`, which can be used to calculate the typehash directly.

Reviewed By: mariusae

Differential Revision: D76357417

Author: pzhan9

hyperactor/src/accum.rs

@@ -33,6 +33,7 @@ pub(crate) mod sealed {
     use async_trait::async_trait;
 
     use crate::PortId;
+    use crate::accum::ReducerSpec;
     use crate::actor::Actor;
     use crate::actor::ActorHandle;
     use crate::data::Serialized;
@@ -47,7 +48,11 @@ pub(crate) mod sealed {
     }
 
     pub trait CanSplitPort: Send + Sync {
-        fn split(&self, port_id: PortId, reducer: Option<u64>) -> PortId;
+        fn split(
+            &self,
+            port_id: PortId,
+            reducer_spec: Option<ReducerSpec>,
+        ) -> anyhow::Result<PortId>;
     }
 
     #[async_trait]

hyperactor/src/cap.rs

@@ -102,6 +102,7 @@ use crate::OncePortRef;
 use crate::PortRef;
 use crate::accum;
 use crate::accum::Accumulator;
+use crate::accum::ReducerSpec;
 use crate::actor::Signal;
 use crate::actor::remote::USER_PORT_OFFSET;
 use crate::cap;
@@ -1016,10 +1017,10 @@ impl Mailbox {
         let (sender, receiver) = mpsc::unbounded_channel::<A::State>();
         let port_id = PortId(self.state.actor_id.clone(), port_index);
         let state = Mutex::new(A::State::default());
-        let reducer_typehash = accum.reducer_typehash();
+        let reducer_spec = accum.reducer_spec();
         let enqueue = move |update: A::Update| {
             let mut state = state.lock().unwrap();
-            accum.accumulate(&mut state, update);
+            accum.accumulate(&mut state, update)?;
             let _ = sender.send(state.clone());
             Ok(())
         };
@@ -1029,7 +1030,7 @@ impl Mailbox {
                 port_index,
                 sender: UnboundedPortSender::Func(Arc::new(enqueue)),
                 bound: Arc::new(OnceLock::new()),
-                reducer_typehash: Some(reducer_typehash),
+                reducer_spec,
             },
             PortReceiver::new(
                 receiver,
@@ -1052,7 +1053,7 @@ impl Mailbox {
             port_index: self.state.allocate_port(),
             sender: UnboundedPortSender::Func(Arc::new(enqueue)),
             bound: Arc::new(OnceLock::new()),
-            reducer_typehash: None,
+            reducer_spec: None,
         }
     }
 
@@ -1240,7 +1241,7 @@ impl SplitPortBuffer {
 }
 
 impl cap::sealed::CanSplitPort for Mailbox {
-    fn split(&self, port_id: PortId, reducer_typehash: Option<u64>) -> PortId {
+    fn split(&self, port_id: PortId, reducer_spec: Option<ReducerSpec>) -> anyhow::Result<PortId> {
         fn post(mailbox: &Mailbox, port_id: PortId, msg: Serialized) {
             mailbox.post(
                 MessageEnvelope::new(mailbox.actor_id().clone(), port_id, msg),
@@ -1255,7 +1256,15 @@ impl cap::sealed::CanSplitPort for Mailbox {
         let port_index = self.state.allocate_port();
         let split_port = self.actor_id().port_id(port_index);
         let mailbox = self.clone();
-        let reducer = reducer_typehash.and_then(accum::resolve_reducer);
+        let reducer = reducer_spec
+            .map(
+                |ReducerSpec {
+                     typehash,
+                     builder_params,
+                 }| { accum::resolve_reducer(typehash, builder_params) },
+            )
+            .transpose()?
+            .flatten();
         let enqueue: Box<
             dyn Fn(Serialized) -> Result<(), (Serialized, anyhow::Error)> + Send + Sync,
         > = match reducer {
@@ -1291,7 +1300,7 @@ impl cap::sealed::CanSplitPort for Mailbox {
                 port_id: split_port.clone(),
             },
         );
-        split_port
+        Ok(split_port)
     }
 }
 
@@ -1317,7 +1326,7 @@ pub struct PortHandle<M: Message> {
     bound: Arc<OnceLock<PortId>>,
     // Typehash of an optional reducer. When it's defined, we include it in port
     /// references to optionally enable incremental accumulation.
-    reducer_typehash: Option<u64>,
+    reducer_spec: Option<ReducerSpec>,
 }
 
 impl<M: Message> PortHandle<M> {
@@ -1327,7 +1336,7 @@ impl<M: Message> PortHandle<M> {
             port_index,
             sender,
             bound: Arc::new(OnceLock::new()),
-            reducer_typehash: None,
+            reducer_spec: None,
         }
     }
 
@@ -1357,7 +1366,7 @@ impl<M: RemoteMessage> PortHandle<M> {
             self.bound
                 .get_or_init(|| self.mailbox.bind(self).port_id().clone())
                 .clone(),
-            self.reducer_typehash.clone(),
+            self.reducer_spec.clone(),
         )
     }
 
@@ -1375,7 +1384,7 @@ impl<M: Message> Clone for PortHandle<M> {
             port_index: self.port_index,
             sender: self.sender.clone(),
             bound: self.bound.clone(),
-            reducer_typehash: self.reducer_typehash.clone(),
+            reducer_spec: self.reducer_spec.clone(),
         }
     }
 }
@@ -2253,18 +2262,18 @@ mod tests {
         // accum port could have reducer typehash
         {
             let accumulator = accum::max::<u64>();
-            let reducer_typehash = accumulator.reducer_typehash();
+            let reducer_spec = accumulator.reducer_spec().unwrap();
             let (port, _) = mbox.open_accum_port(accum::max::<u64>());
-            assert_eq!(port.reducer_typehash, Some(reducer_typehash),);
+            assert_eq!(port.reducer_spec, Some(reducer_spec.clone()));
             let port_ref = port.bind();
-            assert_eq!(port_ref.reducer_typehash(), &Some(reducer_typehash));
+            assert_eq!(port_ref.reducer_spec(), &Some(reducer_spec));
         }
         // normal port should not have reducer typehash
         {
             let (port, _) = mbox.open_port::<u64>();
-            assert_eq!(port.reducer_typehash, None);
+            assert_eq!(port.reducer_spec, None);
             let port_ref = port.bind();
-            assert_eq!(port_ref.reducer_typehash(), &None);
+            assert_eq!(port_ref.reducer_spec(), &None);
         }
     }
 
@@ -2822,7 +2831,7 @@ mod tests {
         port_id2_1: PortId,
     }
 
-    async fn setup_split_port_ids(reducer_typehash: Option<u64>) -> Setup {
+    async fn setup_split_port_ids(reducer_spec: Option<ReducerSpec>) -> Setup {
         let muxer = MailboxMuxer::new();
         let actor0 = Mailbox::new(id!(test[0].actor), BoxedMailboxSender::new(muxer.clone()));
         let actor1 = Mailbox::new(id!(test[1].actor1), BoxedMailboxSender::new(muxer.clone()));
@@ -2834,11 +2843,11 @@ mod tests {
         let port_id = port_handle.bind().port_id().clone();
 
         // Split it twice on actor1
-        let port_id1 = port_id.split(&actor1, reducer_typehash.clone());
-        let port_id2 = port_id.split(&actor1, reducer_typehash.clone());
+        let port_id1 = port_id.split(&actor1, reducer_spec.clone()).unwrap();
+        let port_id2 = port_id.split(&actor1, reducer_spec.clone()).unwrap();
 
         // A split port id can also be split
-        let port_id2_1 = port_id2.split(&actor1, reducer_typehash);
+        let port_id2_1 = port_id2.split(&actor1, reducer_spec).unwrap();
 
         Setup {
             receiver,
@@ -2911,7 +2920,7 @@ mod tests {
         let _config_guard = config.override_key(crate::config::SPLIT_MAX_BUFFER_SIZE, 1);
 
         let sum_accumulator = accum::sum::<u64>();
-        let reducer_typehash = sum_accumulator.reducer_typehash();
+        let reducer_spec = sum_accumulator.reducer_spec();
         let Setup {
             mut receiver,
             actor0,
@@ -2920,7 +2929,7 @@ mod tests {
             port_id1,
             port_id2,
             port_id2_1,
-        } = setup_split_port_ids(Some(reducer_typehash)).await;
+        } = setup_split_port_ids(reducer_spec).await;
         post(&actor0, port_id.clone(), 4);
         post(&actor1, port_id1.clone(), 2);
         post(&actor1, port_id2.clone(), 3);
@@ -2948,8 +2957,8 @@ mod tests {
         let (port_handle, mut receiver) = actor.open_port::<u64>();
         let port_id = port_handle.bind().port_id().clone();
         // Split it
-        let reducer_typehash = accum::sum::<u64>().reducer_typehash();
-        let split_port_id = port_id.split(&actor, Some(reducer_typehash));
+        let reducer_spec = accum::sum::<u64>().reducer_spec();
+        let split_port_id = port_id.split(&actor, reducer_spec).unwrap();
 
         // Send 9 messages.
         for msg in [1, 5, 3, 4, 2, 91, 92, 93, 94] {

hyperactor/src/mailbox/mod.rs

@@ -313,6 +313,7 @@ mod tests {
     use hyperactor::id;
 
     use super::*;
+    use crate::accum::ReducerSpec;
     use crate::reference::UnboundPort;
 
     // Used to demonstrate a user defined reply type.
@@ -349,7 +350,13 @@ mod tests {
     #[test]
     fn test_castable() {
         let original_port0 = PortRef::attest(id!(world[0].actor[0][123]));
-        let original_port1 = PortRef::attest_reducible(id!(world[1].actor1[0][456]), Some(123));
+        let original_port1 = PortRef::attest_reducible(
+            id!(world[1].actor1[0][456]),
+            Some(ReducerSpec {
+                typehash: 123,
+                builder_params: None,
+            }),
+        );
         let my_message = MyMessage {
             arg0: true,
             arg1: 42,
@@ -398,7 +405,13 @@ mod tests {
             .unwrap();
 
         let new_port0 = PortRef::<String>::attest(new_port_id0);
-        let new_port1 = PortRef::<MyReply>::attest_reducible(new_port_id1, Some(123));
+        let new_port1 = PortRef::<MyReply>::attest_reducible(
+            new_port_id1,
+            Some(ReducerSpec {
+                typehash: 123,
+                builder_params: None,
+            }),
+        );
         let new_bindings = Bindings(
             [
                 (

hyperactor/src/message.rs

@@ -52,6 +52,7 @@ use crate::Handler;
 use crate::Message;
 use crate::Named;
 use crate::RemoteMessage;
+use crate::accum::ReducerSpec;
 use crate::actor::ActorError;
 use crate::actor::ActorErrorKind;
 use crate::actor::ActorHandle;
@@ -1198,8 +1199,8 @@ impl<A: Actor> cap::sealed::CanOpenPort for Instance<A> {
 }
 
 impl<A: Actor> cap::sealed::CanSplitPort for Instance<A> {
-    fn split(&self, port_id: PortId, reducer_typehash: Option<u64>) -> PortId {
-        self.mailbox.split(port_id, reducer_typehash)
+    fn split(&self, port_id: PortId, reducer_spec: Option<ReducerSpec>) -> anyhow::Result<PortId> {
+        self.mailbox.split(port_id, reducer_spec)
     }
 }
 

hyperactor/src/proc.rs

@@ -42,6 +42,7 @@ use crate as hyperactor;
 use crate::Named;
 use crate::RemoteHandles;
 use crate::RemoteMessage;
+use crate::accum::ReducerSpec;
 use crate::actor::RemoteActor;
 use crate::cap;
 use crate::data::Serialized;
@@ -728,8 +729,12 @@ impl PortId {
 
     /// Split this port, returning a new port that relays messages to the port
     /// through a local proxy, which may coalesce messages.
-    pub fn split(&self, caps: &impl cap::CanSplitPort, reducer_typehash: Option<u64>) -> PortId {
-        caps.split(self.clone(), reducer_typehash)
+    pub fn split(
+        &self,
+        caps: &impl cap::CanSplitPort,
+        reducer_spec: Option<ReducerSpec>,
+    ) -> anyhow::Result<PortId> {
+        caps.split(self.clone(), reducer_spec)
     }
 }
 
@@ -763,7 +768,7 @@ pub struct PortRef<M: RemoteMessage> {
         Ord = "ignore",
         Hash = "ignore"
     )]
-    reducer_typehash: Option<u64>,
+    reducer_spec: Option<ReducerSpec>,
     phantom: PhantomData<M>,
 }
 
@@ -773,17 +778,17 @@ impl<M: RemoteMessage> PortRef<M> {
     pub fn attest(port_id: PortId) -> Self {
         Self {
             port_id,
-            reducer_typehash: None,
+            reducer_spec: None,
             phantom: PhantomData,
         }
     }
 
     /// The caller attests that the provided PortId can be
     /// converted to a reachable, typed port reference.
-    pub(crate) fn attest_reducible(port_id: PortId, reducer_typehash: Option<u64>) -> Self {
+    pub(crate) fn attest_reducible(port_id: PortId, reducer_spec: Option<ReducerSpec>) -> Self {
         Self {
             port_id,
-            reducer_typehash,
+            reducer_spec,
             phantom: PhantomData,
         }
     }
@@ -796,8 +801,8 @@ impl<M: RemoteMessage> PortRef<M> {
 
     /// The typehash of this port's reducer, if any. Reducers
     /// may be used to coalesce messages sent to a port.
-    pub fn reducer_typehash(&self) -> &Option<u64> {
-        &self.reducer_typehash
+    pub fn reducer_spec(&self) -> &Option<ReducerSpec> {
+        &self.reducer_spec
     }
 
     /// This port's ID.
@@ -846,7 +851,7 @@ impl<M: RemoteMessage> Clone for PortRef<M> {
     fn clone(&self) -> Self {
         Self {
             port_id: self.port_id.clone(),
-            reducer_typehash: self.reducer_typehash.clone(),
+            reducer_spec: self.reducer_spec.clone(),
             phantom: PhantomData,
         }
     }
@@ -866,7 +871,7 @@ impl<M: RemoteMessage> Named for PortRef<M> {
 
 /// The parameters extracted from [`PortRef`] to [`Bindings`].
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Named)]
-pub struct UnboundPort(pub PortId, pub Option<u64>);
+pub struct UnboundPort(pub PortId, pub Option<ReducerSpec>);
 
 impl UnboundPort {
     /// Update the port id of this binding.
@@ -877,7 +882,7 @@ impl UnboundPort {
 
 impl<M: RemoteMessage> From<&PortRef<M>> for UnboundPort {
     fn from(port_ref: &PortRef<M>) -> Self {
-        UnboundPort(port_ref.port_id.clone(), port_ref.reducer_typehash.clone())
+        UnboundPort(port_ref.port_id.clone(), port_ref.reducer_spec.clone())
     }
 }
 
@@ -891,7 +896,7 @@ impl<M: RemoteMessage> Bind for PortRef<M> {
     fn bind(&mut self, bindings: &mut Bindings) -> anyhow::Result<()> {
         let bound = bindings.try_pop_front::<UnboundPort>()?;
         self.port_id = bound.0;
-        self.reducer_typehash = bound.1;
+        self.reducer_spec = bound.1;
         Ok(())
     }
 }

hyperactor/src/reference.rs

@@ -18,6 +18,7 @@ use hyperactor::Handler;
 use hyperactor::Instance;
 use hyperactor::Named;
 use hyperactor::PortRef;
+use hyperactor::accum::ReducerSpec;
 use hyperactor::message::Bind;
 use hyperactor::message::Bindings;
 use hyperactor::message::IndexedErasedUnbound;

hyperactor_mesh/examples/dining_philosophers.rs

@@ -19,6 +19,7 @@ use hyperactor::Named;
 use hyperactor::PortHandle;
 use hyperactor::RemoteHandles;
 use hyperactor::RemoteMessage;
+use hyperactor::accum::ReducerSpec;
 use hyperactor::actor::RemoteActor;
 use hyperactor::mailbox::MailboxSenderError;
 use hyperactor::mailbox::PortReceiver;