mesh: streams are keyed on (ActorMeshId, Sender) (#471)

Summary:
Pull Request resolved: #471

This change aligns all of the stream keys along the comm actor delivery path. We formalize streams to be per-(ActorMeshId, Sender), and enforce this everywhere.

After this change, we can begin to simplify the comm actor delivery code.

with shayne-fletcher

Reviewed By: shayne-fletcher

Differential Revision: D77963903

fbshipit-source-id: d906e5384b3cace8d900c0b1acf7c9509f249491

Author: mariusae

controller/src/lib.rs

@@ -39,6 +39,8 @@ use hyperactor_mesh::comm::multicast::CastMessage;
 use hyperactor_mesh::comm::multicast::CastMessageEnvelope;
 use hyperactor_mesh::comm::multicast::DestinationPort;
 use hyperactor_mesh::comm::multicast::Uslice;
+use hyperactor_mesh::reference::ActorMeshId;
+use hyperactor_mesh::reference::ProcMeshId;
 use hyperactor_multiprocess::proc_actor::ProcActor;
 use hyperactor_multiprocess::proc_actor::spawn;
 use hyperactor_multiprocess::supervision::WorldSupervisionMessageClient;
@@ -422,7 +424,12 @@ impl ControllerMessageHandler for ControllerActor {
                 dsl::union(sel, slice_to_selection(slice))
             }),
         };
+
         let message = CastMessageEnvelope::from_serialized(
+            ActorMeshId(
+                ProcMeshId(self.worker_gang_ref.gang_id().world_id().to_string()),
+                self.worker_gang_ref.gang_id().name().to_string(),
+            ),
             cx.self_id().clone(),
             DestinationPort::new::<WorkerActor, WorkerMessage>(
                 // This is awkward, but goes away entirely with meshes.

hyperactor_mesh/src/actor_mesh.rs

@@ -60,6 +60,7 @@ use crate::reference::ProcMeshId;
 #[allow(clippy::result_large_err)] // TODO: Consider reducing the size of `CastError`.
 pub(crate) fn actor_mesh_cast<M: Castable + Clone, A>(
     caps: &impl cap::CanSend,
+    actor_mesh_id: ActorMeshId,
     actor_mesh_shape: &Shape,
     proc_mesh_shape: &Shape,
     actor_name: &str,
@@ -77,6 +78,7 @@ where
     ));
 
     let message = CastMessageEnvelope::new(
+        actor_mesh_id,
         sender.clone(),
         DestinationPort::new::<A, M>(actor_name.to_string()),
         actor_mesh_shape.clone(),
@@ -115,7 +117,7 @@ where
 }
 
 /// A mesh of actors, all of which reside on the same [`ProcMesh`].
-pub trait ActorMesh: Mesh {
+pub trait ActorMesh: Mesh<Id = ActorMeshId> {
     /// The type of actor in the mesh.
     type Actor: RemoteActor;
 
@@ -128,6 +130,7 @@ pub trait ActorMesh: Mesh {
     {
         actor_mesh_cast::<M, Self::Actor>(
             self.proc_mesh().client(),
+            self.id(),
             self.shape(),
             self.proc_mesh().shape(),
             self.name(),
@@ -288,6 +291,7 @@ impl<'a, A: RemoteActor> RootActorMesh<'a, A> {
 #[async_trait]
 impl<'a, A: RemoteActor> Mesh for RootActorMesh<'a, A> {
     type Node = ActorRef<A>;
+    type Id = ActorMeshId;
     type Sliced<'b>
         = SlicedActorMesh<'b, A>
     where
@@ -308,6 +312,10 @@ impl<'a, A: RemoteActor> Mesh for RootActorMesh<'a, A> {
     fn get(&self, rank: usize) -> Option<ActorRef<A>> {
         self.ranks.get(rank).cloned()
     }
+
+    fn id(&self) -> Self::Id {
+        ActorMeshId(self.proc_mesh.id(), self.name.clone())
+    }
 }
 
 impl<A: RemoteActor> ActorMesh for RootActorMesh<'_, A> {
@@ -337,6 +345,7 @@ impl<'a, A: RemoteActor> SlicedActorMesh<'a, A> {
 #[async_trait]
 impl<A: RemoteActor> Mesh for SlicedActorMesh<'_, A> {
     type Node = ActorRef<A>;
+    type Id = ActorMeshId;
     type Sliced<'b>
         = SlicedActorMesh<'b, A>
     where
@@ -357,6 +366,10 @@ impl<A: RemoteActor> Mesh for SlicedActorMesh<'_, A> {
     fn get(&self, _index: usize) -> Option<ActorRef<A>> {
         unimplemented!()
     }
+
+    fn id(&self) -> Self::Id {
+        self.0.id()
+    }
 }
 
 impl<A: RemoteActor> ActorMesh for SlicedActorMesh<'_, A> {

hyperactor_mesh/src/comm.rs

@@ -7,6 +7,7 @@
  */
 
 use crate::comm::multicast::CAST_ORIGINATING_SENDER;
+use crate::reference::ActorMeshId;
 pub mod multicast;
 
 use std::cmp::Ordering;
@@ -84,10 +85,10 @@ struct ReceiveState {
     ],
 )]
 pub struct CommActor {
-    /// Each world will use its own seq num from this caster.
-    send_seq: HashMap<Slice, usize>,
+    /// Sequence numbers are maintained for each (actor mesh id, sender).
+    send_seq: HashMap<(ActorMeshId, ActorId), usize>,
     /// Each sender is a unique stream.
-    recv_state: HashMap<ActorId, ReceiveState>,
+    recv_state: HashMap<(ActorMeshId, ActorId), ReceiveState>,
 
     /// The comm actor's mode.
     mode: CommActorMode,
@@ -313,7 +314,7 @@ impl Handler<CastMessage> for CommActor {
         let rank = frame.slice.location(&frame.here)?;
         let seq = self
             .send_seq
-            .entry(frame.slice.as_ref().clone())
+            .entry(cast_message.message.stream_key())
             .or_default();
         let last_seq = *seq;
         *seq += 1;
@@ -351,7 +352,7 @@ impl Handler<ForwardMessage> for CommActor {
                 panic!("Choice encountered in CommActor routing")
             })?;
 
-        let recv_state = self.recv_state.entry(sender.clone()).or_default();
+        let recv_state = self.recv_state.entry(message.stream_key()).or_default();
         match recv_state.seq.cmp(&last_seq) {
             // We got the expected next message to deliver to this host.
             Ordering::Equal => {

hyperactor_mesh/src/comm/multicast.rs

@@ -28,6 +28,8 @@ use ndslice::selection::routing::RoutingFrame;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::reference::ActorMeshId;
+
 /// A union of slices that can be used to represent arbitrary subset of
 /// ranks in a gang. It is represented by a Slice together with a Selection.
 /// This is used to define the destination of a cast message or the source of
@@ -43,6 +45,8 @@ pub struct Uslice {
 /// An envelope that carries a message destined to a group of actors.
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Named)]
 pub struct CastMessageEnvelope {
+    /// The destination actor mesh id.
+    actor_mesh_id: ActorMeshId,
     /// The sender of this message.
     sender: ActorId,
     /// The destination port of the message. It could match multiple actors with
@@ -59,6 +63,7 @@ pub struct CastMessageEnvelope {
 impl CastMessageEnvelope {
     /// Create a new CastMessageEnvelope.
     pub fn new<T: Castable + Serialize + Named>(
+        actor_mesh_id: ActorMeshId,
         sender: ActorId,
         dest_port: DestinationPort,
         shape: Shape,
@@ -67,6 +72,7 @@ impl CastMessageEnvelope {
     ) -> Result<Self, anyhow::Error> {
         let data = ErasedUnbound::try_from_message(message)?;
         Ok(Self {
+            actor_mesh_id,
             sender,
             dest_port,
             data,
@@ -79,12 +85,14 @@ impl CastMessageEnvelope {
     /// when the message do not contain reply ports. Or it does but you are okay
     /// with the destination actors reply to the client actor directly.
     pub fn from_serialized(
+        actor_mesh_id: ActorMeshId,
         sender: ActorId,
         dest_port: DestinationPort,
         shape: Shape,
         data: Serialized,
     ) -> Self {
         Self {
+            actor_mesh_id,
             sender,
             dest_port,
             data: ErasedUnbound::new(data),
@@ -112,6 +120,13 @@ impl CastMessageEnvelope {
     pub(crate) fn shape(&self) -> &Shape {
         &self.shape
     }
+
+    /// The unique key used to indicate the stream to which to deliver this message.
+    /// Concretely, the comm actors along the path should use this key to manage
+    /// sequence numbers and reorder buffers.
+    pub(crate) fn stream_key(&self) -> (ActorMeshId, ActorId) {
+        (self.actor_mesh_id.clone(), self.sender.clone())
+    }
 }
 
 /// Destination port id of a message. It is a `PortId` with the rank masked out,

hyperactor_mesh/src/mesh.rs

@@ -7,6 +7,7 @@
  */
 
 use async_trait::async_trait;
+use hyperactor::RemoteMessage;
 use ndslice::Range;
 use ndslice::Shape;
 use ndslice::ShapeError;
@@ -18,6 +19,9 @@ pub trait Mesh {
     /// The type of the node contained in the mesh.
     type Node;
 
+    /// The type of identifiers for this mesh.
+    type Id: RemoteMessage;
+
     /// The type of a slice of this mesh. Slices should not outlive their
     /// parent mesh.
     type Sliced<'a>: Mesh<Node = Self::Node> + 'a
@@ -43,6 +47,9 @@ pub trait Mesh {
             slice_iter: self.shape().slice().iter(),
         }
     }
+
+    /// The global identifier for this mesh.
+    fn id(&self) -> Self::Id;
 }
 
 /// An iterator over the nodes of a mesh.

hyperactor_mesh/src/proc_mesh.rs

@@ -54,6 +54,7 @@ use crate::assign::Ranks;
 use crate::comm::CommActorMode;
 use crate::proc_mesh::mesh_agent::MeshAgent;
 use crate::proc_mesh::mesh_agent::MeshAgentMessageClient;
+use crate::reference::ProcMeshId;
 
 pub mod mesh_agent;
 
@@ -574,6 +575,7 @@ impl<D: Deref<Target = ProcMesh> + Send + Sync + 'static> SharedSpawnable for D
 #[async_trait]
 impl Mesh for ProcMesh {
     type Node = ProcId;
+    type Id = ProcMeshId;
     type Sliced<'a> = SlicedProcMesh<'a>;
 
     fn shape(&self) -> &Shape {
@@ -591,6 +593,10 @@ impl Mesh for ProcMesh {
     fn get(&self, rank: usize) -> Option<ProcId> {
         Some(self.ranks[rank].0.clone())
     }
+
+    fn id(&self) -> Self::Id {
+        ProcMeshId(self.world_id().name().to_string())
+    }
 }
 
 impl fmt::Display for ProcMesh {
@@ -616,6 +622,7 @@ pub struct SlicedProcMesh<'a>(&'a ProcMesh, Shape);
 #[async_trait]
 impl Mesh for SlicedProcMesh<'_> {
     type Node = ProcId;
+    type Id = ProcMeshId;
     type Sliced<'b>
         = SlicedProcMesh<'b>
     where
@@ -636,6 +643,10 @@ impl Mesh for SlicedProcMesh<'_> {
     fn get(&self, _index: usize) -> Option<ProcId> {
         unimplemented!()
     }
+
+    fn id(&self) -> Self::Id {
+        self.0.id()
+    }
 }
 
 #[cfg(test)]

hyperactor_mesh/src/reference.rs

@@ -134,6 +134,7 @@ impl<A: RemoteActor> ActorMeshRef<A> {
     {
         actor_mesh_cast::<M, A>(
             caps,
+            self.mesh_id.clone(),
             self.shape(),
             self.proc_mesh_shape(),
             self.name(),