mesh: key streams off of sender only (#397)

Summary:
Pull Request resolved: #397

Each stream was previously associated with `(sender, <shape of mesh>)`, the shape being a stand-in for the mesh itself. In practice, each proc is reachable through exactly one mesh, so this is unnecessary.

This change simplifies this to just the sender.

In the future, we'll likely want to support multiple sessions per sender (e.g., for persistent meshes), but that is a straightforward addition we can make as needed.
ghstack-source-id: 293881496
exported-using-ghexport

Reviewed By: shayne-fletcher, vidhyav

Differential Revision: D77609169

fbshipit-source-id: 4babc95eba0d4123f9308b43b5bf37a9dfab2747

Author: mariusae

hyperactor_mesh/src/comm.rs

@@ -83,8 +83,8 @@ struct ReceiveState {
 pub struct CommActor {
     /// Each world will use its own seq num from this caster.
     send_seq: HashMap<Slice, usize>,
-    /// Each world/caster uses its own stream.
-    recv_state: HashMap<(Slice, ActorId), ReceiveState>,
+    /// Each sender is a unique stream.
+    recv_state: HashMap<ActorId, ReceiveState>,
 
     /// The comm actor's mode.
     mode: CommActorMode,
@@ -340,13 +340,12 @@ impl Handler<ForwardMessage> for CommActor {
 
         // Resolve/dedup routing frames.
         let rank = self.mode.self_rank(this.self_id());
-        let slice = dests[0].slice.as_ref().clone();
         let (deliver_here, next_steps) =
             ndslice::selection::routing::resolve_routing(rank, dests, &mut |_| {
                 panic!("Choice encountered in CommActor routing")
             })?;
 
-        let recv_state = self.recv_state.entry((slice, sender.clone())).or_default();
+        let recv_state = self.recv_state.entry(sender.clone()).or_default();
         match recv_state.seq.cmp(&last_seq) {
             // We got the expected next message to deliver to this host.
             Ordering::Equal => {

ndslice/src/selection/routing.rs

@@ -421,7 +421,7 @@ impl RoutingFrame {
                 }
 
                 use rand::Rng;
-                let mut rng = rand::thread_rng();
+                let mut rng: rand::prelude::ThreadRng = rand::thread_rng();
                 let i = rng.gen_range(0..size);
                 let mut coord = self.here.clone();
                 coord[self.dim] = i;