Fix number of messages being waited on (#855)

Summary:
Pull Request resolved: #855

Previously we were not waiting for all messages being casted (missing GPU)

Suprisingly this caused the job to hang

Differential Revision: D80187064

Author: pablorfb-meta

hyperactor_mesh/benches/main.rs

@@ -13,6 +13,7 @@ use criterion::Criterion;
 use criterion::Throughput;
 use criterion::criterion_group;
 use criterion::criterion_main;
+use hyperactor::proc;
 use hyperactor_mesh::ProcMesh;
 use hyperactor_mesh::actor_mesh::ActorMesh;
 use hyperactor_mesh::actor_mesh::RootActorMesh;
@@ -74,8 +75,13 @@ fn bench_actor_scaling(c: &mut Criterion) {
                         .unwrap();
 
                     let mut msg_rcv = 0;
-                    while msg_rcv < host_count {
-                        let _ = rx.recv().await.unwrap();
+                    while msg_rcv < host_count * 8 {
+                        tokio::select! {
+                            _ = rx.recv() => {}
+                            _ = tokio::time::sleep(Duration::from_secs(10)) => {
+                                panic!("Timed out. Expected {} messages but got {}", host_count, msg_rcv);
+                            }
+                        }
                         msg_rcv += 1;
                     }
                 }
@@ -89,7 +95,7 @@ fn bench_actor_scaling(c: &mut Criterion) {
                     .await
                     .expect("Failed to stop allocator");
                 elapsed
-            });
+            })
         });
     }
 
@@ -154,6 +160,7 @@ fn bench_actor_mesh_message_sizes(c: &mut Criterion) {
                             let payload = vec![0u8; message_size];
 
                             actor_mesh
+
                                 .cast(
                                     client,
                                     all(true_()),
@@ -167,7 +174,12 @@ fn bench_actor_mesh_message_sizes(c: &mut Criterion) {
 
                             let mut msg_rcv = 0;
                             while msg_rcv < actor_count {
-                                let _ = rx.recv().await.unwrap();
+                                tokio::select! {
+                                    _ = rx.recv() => {}
+                                    _ = tokio::time::sleep(Duration::from_secs(10)) => {
+                                        panic!("Timed out. Expected {} messages but got {}", actor_count, msg_rcv);
+                                    }
+                                }
                                 msg_rcv += 1;
                             }
                         }