Variable message size cast to actor mesh (#763)

Summary:
Pull Request resolved: #763

X-link: #763

Benchmark how long it takes to send a message of size X to an actor mesh of 1 and 10 actors on local transport

Benchmark shows throughput starts scaling from 10kb -> 1MB. After that system seems to have reached max throughput (Bottleneck TBD):
 - ~500Mb/s for 1 Actor
- ~2Gb/s for 10 Actors

| Benchmark                | Time [Min, Median, Max]           | Throughput [Min, Median, Max]           | Throughput Change % |
|--------------------------|----------------------------------|-----------------------------------------|--------------------|
| actors/1/size/10kb       | [1.1249 ms, 1.1324 ms, 1.1401 ms] | [8.3649 MiB/s, 8.4220 MiB/s, 8.4779 MiB/s] | N/A                |
| actors/1/size/100kb      | [1.1926 ms, 1.1984 ms, 1.2043 ms] | [79.187 MiB/s, 79.577 MiB/s, 79.969 MiB/s] | +845.3%            |
| actors/1/size/1mb        | [1.7636 ms, 1.7798 ms, 1.7979 ms] | [530.44 MiB/s, 535.84 MiB/s, 540.77 MiB/s] | +573.6%            |
| actors/1/size/10mb       | [18.867 ms, 19.286 ms, 19.640 ms] | [485.57 MiB/s, 494.50 MiB/s, 505.48 MiB/s] | -7.7%              |
| actors/1/size/100mb      | [180.03 ms, 183.91 ms, 187.25 ms] | [509.29 MiB/s, 518.56 MiB/s, 529.73 MiB/s] | +4.9%              |
| actors/1/size/1gb        | [1.7496 s, 1.7891 s, 1.8314 s]    | [520.73 MiB/s, 533.05 MiB/s, 545.10 MiB/s] | +2.8%              |
| actors/10/size/10kb      | [1.2557 ms, 1.2662 ms, 1.2758 ms] | [74.751 MiB/s, 75.318 MiB/s, 75.948 MiB/s] | N/A            |
| actors/10/size/100kb     | [1.3862 ms, 1.4058 ms, 1.4244 ms] | [669.55 MiB/s, 678.38 MiB/s, 688.00 MiB/s] | +801.1%            |
| actors/10/size/1mb       | [3.3235 ms, 3.3544 ms, 3.3854 ms] | [2.7510 GiB/s, 2.7764 GiB/s, 2.8022 GiB/s] | +309.3%            |
| actors/10/size/10mb      | [41.557 ms, 42.532 ms, 43.215 ms] | [2.1551 GiB/s, 2.1897 GiB/s, 2.2411 GiB/s] | -21.1%             |
| actors/10/size/100mb     | [427.53 ms, 443.07 ms, 460.93 ms] | [2.0205 GiB/s, 2.1020 GiB/s, 2.1784 GiB/s] | -4.0%              |
| actors/10/size/1gb       | [4.4195 s, 4.6372 s, 4.8398 s]    | [1.9243 GiB/s, 2.0084 GiB/s, 2.1073 GiB/s] | -4.5%              |

Comparison to raw channel benchmark  https://fburl.com/code/gyryo1mc seems to indicate that there is a bottleneck on the actor cast path capping throughput .

| Benchmark                    | Time [Min, Median, Max]            | Throughput [Min, Median, Max]           | Throughput Change % |
|------------------------------|------------------------------------|-----------------------------------------|---------------------|
| send_receive/local/10000      | [511.62 ns, 514.23 ns, 517.90 ns]  | [17.983 GiB/s, 18.111 GiB/s, 18.203 GiB/s] | N/A                |
| send_receive/local/1000000    | [131.12 µs, 132.97 µs, 134.57 µs]  | [6.9205 GiB/s, 7.0040 GiB/s, 7.1029 GiB/s]  | -61.3%             |
| send_receive/local/10000000   | [3.0627 ms, 3.0889 ms, 3.1061 ms]  | [2.9984 GiB/s, 3.0150 GiB/s, 3.0408 GiB/s]  | -56.9%             |
| send_receive/local/100000000  | [27.317 ms, 28.661 ms, 29.659 ms]  | [3.1400 GiB/s, 3.2494 GiB/s, 3.4094 GiB/s]  | +7.8%              |
| send_receive/local/1000000000 | [244.17 ms, 247.15 ms, 250.02 ms]  | [3.7250 GiB/s, 3.7682 GiB/s, 3.8142 GiB/s]  | +16.0%             |

Reviewed By: pzhan9

Differential Revision: D79577855

fbshipit-source-id: 32f9a0c7fc1ee67b8fceaa0d6ef50211c4044694

Author: pablorfb-meta

hyperactor_mesh/benches/bench_actor.rs

@@ -36,14 +36,17 @@ pub struct BenchMessage {
         BenchMessage { cast = true },
     ],
 )]
-pub struct BenchActor {}
+pub struct BenchActor {
+    processing_time: Duration,
+}
 
 #[async_trait]
 impl Actor for BenchActor {
-    type Params = ();
-
-    async fn new(_: Self::Params) -> Result<Self, anyhow::Error> {
-        Ok(Self {})
+    type Params = Duration;
+    async fn new(params: Duration) -> Result<Self, anyhow::Error> {
+        Ok(Self {
+            processing_time: params,
+        })
     }
 }
 
@@ -55,7 +58,7 @@ impl Handler<BenchMessage> for BenchActor {
         msg: BenchMessage,
     ) -> Result<(), anyhow::Error> {
         hyperactor::clock::ClockKind::default()
-            .sleep(Duration::from_millis(100))
+            .sleep(self.processing_time.clone())
             .await;
 
         let _ = msg.reply.send(ctx, msg.step);

hyperactor_mesh/benches/main.rs

@@ -10,6 +10,7 @@ use std::time::Instant;
 
 use criterion::BenchmarkId;
 use criterion::Criterion;
+use criterion::Throughput;
 use criterion::criterion_group;
 use criterion::criterion_main;
 use hyperactor_mesh::ProcMesh;
@@ -21,6 +22,8 @@ use hyperactor_mesh::alloc::LocalAllocator;
 use hyperactor_mesh::extent;
 use hyperactor_mesh::selection::dsl::all;
 use hyperactor_mesh::selection::dsl::true_;
+use hyperactor_mesh::shape;
+use tokio::time::Duration;
 
 mod bench_actor;
 use bench_actor::BenchActor;
@@ -48,16 +51,18 @@ fn bench_actor_scaling(c: &mut Criterion) {
                     .unwrap();
 
                 let proc_mesh = ProcMesh::allocate(alloc).await.unwrap();
-                let trainer_mesh: RootActorMesh<BenchActor> =
-                    proc_mesh.spawn("trainer", &()).await.unwrap();
+                let actor_mesh: RootActorMesh<BenchActor> = proc_mesh
+                    .spawn("bench", &(Duration::from_millis(0)))
+                    .await
+                    .unwrap();
                 let client = proc_mesh.client();
 
                 let start = Instant::now();
                 for i in 0..iters {
                     let (tx, mut rx) = client.open_port();
                     let payload = vec![0u8; message_size];
 
-                    trainer_mesh
+                    actor_mesh
                         .cast(
                             all(true_()),
                             BenchMessage {
@@ -83,5 +88,90 @@ fn bench_actor_scaling(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_actor_scaling);
+fn format_size(size: usize) -> String {
+    if size >= 1_000_000_000 {
+        format!("{}GB", size / 1_000_000_000)
+    } else if size >= 1_000_000 {
+        format!("{}MB", size / 1_000_000)
+    } else if size >= 1_000 {
+        format!("{}KB", size / 1_000)
+    } else {
+        format!("{}B", size)
+    }
+}
+
+// Benchmark how long it takes to send a message of size X to an actor mesh of 10 actors
+fn bench_actor_mesh_message_sizes(c: &mut Criterion) {
+    let mut group = c.benchmark_group("actor_mesh_message_sizes");
+    group.sample_size(10);
+    let actor_counts = vec![1, 10];
+    let message_sizes: Vec<usize> = vec![
+        10_000,
+        100_000,
+        1_000_000,
+        10_000_000,
+        100_000_000,
+        1_000_000_000,
+    ];
+
+    for message_size in message_sizes {
+        for &actor_count in &actor_counts {
+            group.throughput(Throughput::Bytes((message_size * actor_count) as u64));
+            group.sampling_mode(criterion::SamplingMode::Flat);
+            group.sample_size(10);
+            group.bench_function(
+                format!("actors/{}/size/{}", actor_count, format_size(message_size)),
+                |b| {
+                    let mut b = b.to_async(Runtime::new().unwrap());
+                    b.iter_custom(|iters| async move {
+                        let alloc = LocalAllocator
+                            .allocate(AllocSpec {
+                                extent: extent!(gpus = actor_count),
+                                constraints: Default::default(),
+                            })
+                            .await
+                            .unwrap();
+
+                        let proc_mesh = ProcMesh::allocate(alloc).await.unwrap();
+                        let actor_mesh: RootActorMesh<BenchActor> = proc_mesh
+                            .spawn("bench", &(Duration::from_millis(0)))
+                            .await
+                            .unwrap();
+
+                        let client = proc_mesh.client();
+
+                        let start = Instant::now();
+                        for i in 0..iters {
+                            let (tx, mut rx) = client.open_port();
+                            let payload = vec![0u8; message_size];
+
+                            actor_mesh
+                                .cast(
+                                    all(true_()),
+                                    BenchMessage {
+                                        step: i as usize,
+                                        reply: tx.bind(),
+                                        payload,
+                                    },
+                                )
+                                .unwrap();
+
+                            let mut msg_rcv = 0;
+                            while msg_rcv < actor_count {
+                                let _ = rx.recv().await.unwrap();
+                                msg_rcv += 1;
+                            }
+                        }
+
+                        start.elapsed()
+                    });
+                },
+            );
+        }
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_actor_scaling, bench_actor_mesh_message_sizes);
 criterion_main!(benches);