Channel benchmarks (pytorch-labs#627)

Summary: Pull Request resolved: pytorch-labs#627

Reviewed By: vidhyav

Differential Revision: D78755894

Author: eliothedeman

hyperactor/Cargo.toml

@@ -1,4 +1,4 @@
-# @generated by autocargo from //monarch/hyperactor:[hyperactor,hyperactor-example-derive]
+# @generated by autocargo from //monarch/hyperactor:[channel_benchmarks,hyperactor,hyperactor-example-derive]
 
 [package]
 name = "hyperactor"
@@ -9,6 +9,10 @@ description = "a high-performance, scalable actor framework for cluster computin
 repository = "https://github.com/pytorch-labs/monarch/"
 license = "BSD-3-Clause"
 
+[[bin]]
+name = "channel_benchmarks"
+path = "benches/channel_benchmarks.rs"
+
 [[bin]]
 name = "hyperactor_example_derive"
 path = "example/derive.rs"
@@ -21,6 +25,7 @@ bincode = "1.3.3"
 bytes = { version = "1.9.0", features = ["serde"] }
 cityhasher = "0.1.0"
 crc32fast = "1.4"
+criterion = { version = "0.5.1", features = ["async_tokio", "csv_output"] }
 dashmap = { version = "5.5.3", features = ["rayon", "serde"] }
 derivative = "2.2"
 dns-lookup = "1.0"

hyperactor/benches/README.md

@@ -0,0 +1,89 @@
+# Hyperactor Channel Benchmarks
+
+This directory contains performance benchmarks for the Hyperactor channel system, specifically testing message passing throughput across different transport types and message sizes.
+
+## Overview
+
+The benchmark suite measures the performance of Hyperactor's channel communication system by testing:
+
+- **Transport Types**: Local, TCP, MetaTLS, and Unix socket transports
+- **Message Sizes**: Ranging from 10 bytes to 1 GB (10^1 to 10^9 bytes)
+- **Metrics**: Throughput measured in bytes per second
+
+## Benchmark Details
+
+### Transport Types Tested
+
+1. **Local**: In-memory transport for same-process communication
+2. **TCP**: Network transport using TCP sockets
+3. **MetaTLS**: Secure transport using Meta's TLS implementation
+4. **Unix**: Unix domain socket transport for inter-process communication
+
+### Message Size Range
+
+The benchmark tests message sizes across multiple orders of magnitude:
+- Small messages: 10B, 100B, 1KB
+- Medium messages: 10KB, 100KB, 1MB
+- Large messages: 10MB, 100MB, 1GB
+
+Each message contains:
+- An ID field (`u64`)
+- A payload of the specified size (filled with zeros)
+
+## Running the Benchmarks
+
+### Prerequisites
+
+Ensure you have the Rust toolchain and required dependencies installed.
+
+### Running All Benchmarks
+
+```bash
+# From the hyperactor directory
+buck run @//mode/opt //monarch/hyperactor/benches:channel_benchmarks -- --bench
+```
+
+### Running with Cargo (if available)
+
+```bash
+# From the hyperactor directory
+cargo bench --bench channel_benchmarks
+```
+
+## Understanding the Results
+
+The benchmark output will show results in the format:
+```
+message_sizes/message_<transport>_<size><unit>
+```
+
+For example:
+- `message_sizes/message_local_1kb`: Local transport with 1KB messages
+- `message_sizes/message_tcp_10mb`: TCP transport with 10MB messages
+- `message_sizes/message_metatls_100b`: MetaTLS transport with 100 byte messages
+
+### Interpreting Performance
+
+- **Higher throughput** (bytes/sec) indicates better performance
+- **Local transport** typically shows the highest throughput
+- **Larger messages** generally achieve higher throughput due to amortized overhead
+- **Secure transports** (MetaTLS) may show lower throughput due to encryption overhead
+
+## Benchmark Implementation
+
+The benchmark uses the [Criterion](https://docs.rs/criterion/) framework for statistical analysis and uses Tokio's async runtime for handling the asynchronous channel operations.
+
+Each benchmark iteration:
+1. Sets up a server listening on the specified transport
+2. Creates a client connection to the server
+3. Sends messages of the specified size
+4. Measures the time taken for round-trip communication
+5. Calculates throughput based on bytes transferred and time elapsed
+
+## Use Cases
+
+These benchmarks are useful for:
+- **Performance regression testing**: Ensuring changes don't degrade performance
+- **Transport selection**: Choosing the appropriate transport for your use case
+- **Capacity planning**: Understanding throughput limits for different message sizes
+- **Optimization validation**: Verifying that performance improvements are effective

hyperactor/benches/channel_benchmarks.rs

@@ -0,0 +1,92 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::time::Instant;
+
+use criterion::Criterion;
+use criterion::Throughput;
+use criterion::criterion_group;
+use criterion::criterion_main;
+use hyperactor::Named;
+use hyperactor::channel::ChannelAddr;
+use hyperactor::channel::ChannelTransport;
+use hyperactor::channel::Rx;
+use hyperactor::channel::Tx;
+use hyperactor::channel::dial;
+use hyperactor::channel::serve;
+use serde::Deserialize;
+use serde::Serialize;
+use tokio::runtime::Runtime;
+
+#[derive(Debug, Clone, Serialize, Deserialize, Named, PartialEq)]
+struct Message {
+    id: u64,
+    payload: Vec<u8>,
+}
+
+impl Message {
+    fn new(id: u64, size: usize) -> Self {
+        Self {
+            id,
+            payload: vec![0; size],
+        }
+    }
+}
+
+// Benchmark message sizes
+fn bench_message_sizes(c: &mut Criterion) {
+    let mut group = c.benchmark_group("message_sizes");
+
+    let transports = vec![
+        ("local", ChannelTransport::Local),
+        ("tcp", ChannelTransport::Tcp),
+        ("metatls", ChannelTransport::MetaTls),
+        ("unix", ChannelTransport::Unix),
+    ];
+
+    for size_exp in 5..10 {
+        let size = 10_usize.pow(size_exp);
+        let fsize = size as f64;
+        let (nice_size, postfix) = match size {
+            1_000..=999_999 => (fsize / 1000f64, "kb"),
+            1_000_000..=999_999_999 => (fsize / 1_000_000f64, "mb"),
+            1_000_000_000..=999_999_999_999 => (fsize / 1_000_000_000f64, "gb"),
+            _ => (fsize, "b"),
+        };
+
+        for (transport_name, transport) in &transports {
+            let transport = transport.clone();
+            group.throughput(Throughput::Bytes(size as u64));
+            group.bench_function(
+                format!("message_{}_{}{}", transport_name, nice_size, postfix),
+                move |b| {
+                    let mut b = b.to_async(Runtime::new().unwrap());
+                    let tt = &transport;
+                    b.iter_custom(|iters| async move {
+                        let addr = ChannelAddr::any(tt.clone());
+                        let (listen_addr, mut rx) = serve::<Message>(addr).await.unwrap();
+                        let tx = dial::<Message>(listen_addr).unwrap();
+                        let msg = Message::new(0, size);
+                        let start = Instant::now();
+                        for _ in 0..iters {
+                            tx.post(msg.clone());
+                            rx.recv().await.unwrap();
+                        }
+                        start.elapsed()
+                    });
+                },
+            );
+        }
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_message_sizes);
+
+criterion_main!(benches);