Refactoring: move verify_casting out so it can be re-used by other tests

Differential Revision: D83001963

Author: pzhan9

hyperactor_mesh/src/v1/actor_mesh.rs

@@ -34,6 +34,7 @@ use serde::Deserializer;
 use serde::Serialize;
 use serde::Serializer;
 
+use crate::CommActor;
 use crate::actor_mesh as v0_actor_mesh;
 use crate::comm::multicast;
 use crate::reference::ActorMeshId;
@@ -135,33 +136,7 @@ impl<A: Actor + RemoteActor> ActorMeshRef<A> {
         M: Castable + RemoteMessage + Clone, // Clone is required until we are fully onto comm actor
     {
         if let Some(root_comm_actor) = self.proc_mesh.root_comm_actor() {
-            let cast_mesh_shape = view::Ranked::region(self).into();
-            let actor_mesh_id = ActorMeshId::V1(self.name.clone());
-            match &self.proc_mesh.root_region {
-                Some(root_region) => {
-                    let root_mesh_shape = root_region.into();
-                    v0_actor_mesh::cast_to_sliced_mesh::<A, M>(
-                        cx,
-                        actor_mesh_id,
-                        root_comm_actor,
-                        &sel!(*),
-                        message,
-                        &cast_mesh_shape,
-                        &root_mesh_shape,
-                    )
-                    .map_err(|e| Error::CastingError(self.name.clone(), e.into()))
-                }
-                None => v0_actor_mesh::actor_mesh_cast::<A, M>(
-                    cx,
-                    actor_mesh_id,
-                    root_comm_actor,
-                    sel!(*),
-                    &cast_mesh_shape,
-                    &cast_mesh_shape,
-                    message,
-                )
-                .map_err(|e| Error::CastingError(self.name.clone(), e.into())),
-            }
+            self.cast_v0(cx, message, root_comm_actor)
         } else {
             for (point, actor) in self.iter() {
                 let mut headers = Attrs::new();
@@ -179,6 +154,45 @@ impl<A: Actor + RemoteActor> ActorMeshRef<A> {
         }
     }
 
+    fn cast_v0<M>(
+        &self,
+        cx: &impl context::Actor,
+        message: M,
+        root_comm_actor: &ActorRef<CommActor>,
+    ) -> v1::Result<()>
+    where
+        A: RemoteHandles<M> + RemoteHandles<IndexedErasedUnbound<M>>,
+        M: Castable + RemoteMessage + Clone, // Clone is required until we are fully onto comm actor
+    {
+        let cast_mesh_shape = view::Ranked::region(self).into();
+        let actor_mesh_id = ActorMeshId::V1(self.name.clone());
+        match &self.proc_mesh.root_region {
+            Some(root_region) => {
+                let root_mesh_shape = root_region.into();
+                v0_actor_mesh::cast_to_sliced_mesh::<A, M>(
+                    cx,
+                    actor_mesh_id,
+                    root_comm_actor,
+                    &sel!(*),
+                    message,
+                    &cast_mesh_shape,
+                    &root_mesh_shape,
+                )
+                .map_err(|e| Error::CastingError(self.name.clone(), e.into()))
+            }
+            None => v0_actor_mesh::actor_mesh_cast::<A, M>(
+                cx,
+                actor_mesh_id,
+                root_comm_actor,
+                sel!(*),
+                &cast_mesh_shape,
+                &cast_mesh_shape,
+                message,
+            )
+            .map_err(|e| Error::CastingError(self.name.clone(), e.into())),
+        }
+    }
+
     pub async fn supervision_events(
         &self,
         cx: &impl context::Actor,

hyperactor_mesh/src/v1/proc_mesh.rs

@@ -591,7 +591,9 @@ impl view::RankedSliceable for ProcMeshRef {
 #[cfg(test)]
 mod tests {
     use std::collections::HashSet;
+    use std::time::Duration;
 
+    use hyperactor::Instance;
     use hyperactor::clock::Clock;
     use hyperactor::clock::RealClock;
     use hyperactor::mailbox;
@@ -626,6 +628,34 @@ mod tests {
         );
     }
 
+    async fn verify_casting(
+        actor_mesh: &ActorMeshRef<testactor::TestActor>,
+        instance: &Instance<()>,
+    ) {
+        let (port, mut rx) = mailbox::open_port(&instance);
+        actor_mesh
+            .cast(&instance, testactor::GetActorId(port.bind()))
+            .unwrap();
+
+        let mut expected_actor_ids: HashSet<_> = actor_mesh
+            .values()
+            .map(|actor_ref| actor_ref.actor_id().clone())
+            .collect();
+
+        while !expected_actor_ids.is_empty() {
+            let actor_id = rx.recv().await.unwrap();
+            assert!(
+                expected_actor_ids.remove(&actor_id),
+                "got {actor_id}, expect {expected_actor_ids:?}"
+            );
+        }
+
+        // No more messages
+        RealClock.sleep(Duration::from_secs(1)).await;
+        let result = rx.try_recv();
+        assert!(result.as_ref().unwrap().is_none(), "got {result:?}");
+    }
+
     #[async_timed_test(timeout_secs = 30)]
     async fn test_spawn_actor() {
         hyperactor_telemetry::initialize_logging(hyperactor::clock::ClockKind::default());
@@ -636,58 +666,20 @@ mod tests {
             let actor_mesh: ActorMesh<testactor::TestActor> =
                 proc_mesh.spawn(instance, "test", &()).await.unwrap();
 
-            // Verify casting to the root actor mesh
-            {
-                let (port, mut rx) = mailbox::open_port(&instance);
-                actor_mesh
-                    .cast(instance, testactor::GetActorId(port.bind()))
-                    .unwrap();
-
-                let mut expected_actor_ids: HashSet<_> = actor_mesh
-                    .values()
-                    .map(|actor_ref| actor_ref.actor_id().clone())
-                    .collect();
-
-                while !expected_actor_ids.is_empty() {
-                    let actor_id = rx.recv().await.unwrap();
-                    assert!(
-                        expected_actor_ids.remove(&actor_id),
-                        "got {actor_id}, expect {expected_actor_ids:?}"
-                    );
-                }
-
-                // No more messages
-                RealClock.sleep(Duration::from_secs(1)).await;
-                let result = rx.try_recv();
-                assert!(result.as_ref().unwrap().is_none(), "got {result:?}");
-            }
+            // First cast. The seq should be 1 for all actors.
+            verify_casting(&actor_mesh, &instance).await;
 
             // Verify casting to the sliced actor mesh
             let sliced_actor_mesh = actor_mesh.range("replicas", 1..3).unwrap();
-            {
-                let (port, mut rx) = mailbox::open_port(instance);
-                sliced_actor_mesh
-                    .cast(instance, testactor::GetActorId(port.bind()))
-                    .unwrap();
-
-                let mut expected_actor_ids: HashSet<_> = sliced_actor_mesh
-                    .values()
-                    .map(|actor_ref| actor_ref.actor_id().clone())
-                    .collect();
-
-                while !expected_actor_ids.is_empty() {
-                    let actor_id = rx.recv().await.unwrap();
-                    assert!(
-                        expected_actor_ids.remove(&actor_id),
-                        "got {actor_id}, expect {expected_actor_ids:?}"
-                    );
-                }
-
-                // No more messages
-                RealClock.sleep(Duration::from_secs(1)).await;
-                let result = rx.try_recv();
-                assert!(result.as_ref().unwrap().is_none(), "got {result:?}");
-            }
+            // Second cast. The seq should be 2 for actors in the sliced mesh.
+            verify_casting(&sliced_actor_mesh, &instance).await;
+
+            // Verify casting to a different sliced actor mesh
+            let sliced_actor_mesh = actor_mesh.range("replicas", 0..2).unwrap();
+            // For actors in the previous sliced mesh, the seq should be 3 since
+            // this is the third cast for them. For other actors, the seq should
+            // be 2.
+            verify_casting(&sliced_actor_mesh, &instance).await;
         }
     }
 }