diff --git a/Cargo.lock b/Cargo.lock index a605ce5d..4d1d4b42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -254,6 +260,12 @@ dependencies = [ "tokio", ] +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" @@ -346,6 +358,12 @@ version = "0.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.11.1" @@ -579,6 +597,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + [[package]] name = "criterion" version = "0.8.2" @@ -885,6 +912,16 @@ version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1140,6 +1177,20 @@ dependencies = [ "hashbrown 0.16.1", ] +[[package]] +name = "hdrhistogram" +version = "7.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d" +dependencies = [ + "base64 0.21.7", + "byteorder", + "crossbeam-channel", + "flate2", + "nom", + "num-traits", +] + [[package]] name = "heck" version = "0.5.0" @@ -1271,7 +1322,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1591,7 +1642,7 @@ version = "0.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05a6d6f3611ad1d21732adbd7a2e921f598af6c92d71ae6e2620da4b67ee1f0d" dependencies = [ - "base64", + "base64 0.22.1", "jiff", "serde", "serde_json", @@ -1615,7 +1666,7 @@ version = "3.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af97b8b696eb737e5694f087c498ca725b172c2a5bc3a6916328d160225537ee" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "either", "futures", @@ -1894,6 +1945,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.1.1" @@ -2189,12 +2250,14 @@ version = "0.18.0" dependencies = [ "ahash", "bytesize", + "clap", "criterion", "cudarc", "dashmap", "fastrace", "futures", "hashlink", + "hdrhistogram", "io-uring", "k8s-openapi", "kube", @@ -2316,7 +2379,7 @@ version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ - "base64", + "base64 0.22.1", "serde_core", ] @@ -2904,7 +2967,7 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2938,7 +3001,7 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -3355,6 +3418,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + [[package]] name = "slab" version = "0.4.12" @@ -3745,7 +3814,7 @@ checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" dependencies = [ "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", "h2", "http", @@ -3830,7 +3899,7 @@ version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ - "base64", + "base64 0.22.1", "bitflags 2.11.0", "bytes", "futures-util", diff --git a/pegaflow-core/Cargo.toml b/pegaflow-core/Cargo.toml index 74618aa1..835ba9b9 100644 --- a/pegaflow-core/Cargo.toml +++ b/pegaflow-core/Cargo.toml @@ -47,6 +47,8 @@ criterion.workspace = true cudarc.workspace = true tokio.workspace = true tempfile = "3" +clap.workspace = true +hdrhistogram = "7.5.4" [[bench]] name = "transfer_batch" @@ -60,5 +62,9 @@ harness = false name = "uds_latency" harness = false +[[bench]] +name = "uring_bench" +harness = false + [[example]] name = "pinned_alloc" diff --git a/pegaflow-core/benches/uring_bench.rs b/pegaflow-core/benches/uring_bench.rs new file mode 100644 index 00000000..32b996f5 --- /dev/null +++ b/pegaflow-core/benches/uring_bench.rs @@ -0,0 +1,260 @@ +/* +======================================== +io_uring SSD Benchmark (pegaflow) +======================================== + +fio equivalent example: + +fio --name=uring-test \ + --filename=testfile \ + --rw=randrw \ + --rwmixwrite=50 \ + --bs=4k \ + --iodepth=64 \ + --numjobs=1 \ + --size=1G \ + --group_reporting + +Parameter mapping: + +| Benchmark Param | fio Param | +|----------------|----------------| +| block_size | bs | +| io_depth | iodepth | +| threads | numjobs | +| concurrency | iodepth*numjobs| +| write_ratio | rwmixwrite | +| total_bytes | size | + +Notes: +- Random IO (uniform distribution, block-aligned) +- Latency = submission → completion (includes queueing) +- Buffers reused (no allocator noise) +- File SHOULD be preallocated for realistic SSD results +*/ + +use std::{ + fs::OpenOptions, + io, + os::unix::{fs::OpenOptionsExt, io::AsRawFd}, + time::{Duration, Instant}, +}; + +use clap::Parser; +use hdrhistogram::Histogram; +use rand::{RngExt, SeedableRng, rngs::StdRng}; +use tokio::sync::oneshot; + +use pegaflow_core::backing::uring::{UringConfig, UringIoEngine}; + +#[derive(Parser, Debug)] +struct Config { + #[arg(long)] + file: String, + + #[arg(long, default_value = "1073741824")] + total_bytes: u64, + + #[arg(long, default_value = "4096")] + block_size: usize, + + #[arg(long, default_value = "1")] + threads: usize, + + #[arg(long, default_value = "64")] + io_depth: usize, + + #[arg(long)] + concurrency: Option, + + #[arg(long, default_value = "0.5")] + write_ratio: f64, + + #[arg(long, default_value = "42")] + seed: u64, + + #[arg(long, default_value = "false")] + direct: bool, +} + +struct Pending { + start: Instant, + rx: oneshot::Receiver>, + buf: Vec, +} + +fn main() -> io::Result<()> { + let cfg = Config::parse(); + + println!("==== uring benchmark ===="); + println!("{:#?}", cfg); + + if cfg.block_size == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "block_size must be > 0", + )); + } + + if cfg.total_bytes < cfg.block_size as u64 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "total_bytes must be >= block_size", + )); + } + + if cfg.threads == 0 || cfg.io_depth == 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "threads and io_depth must be > 0", + )); + } + + let engine_capacity = cfg.threads * cfg.io_depth; + + if let Some(c) = cfg.concurrency + && (c == 0 || c != engine_capacity) + { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "concurrency must equal threads * io_depth and be > 0", + )); + } + + let max_inflight = cfg.concurrency.unwrap_or(engine_capacity); + + let mut opts = OpenOptions::new(); + opts.read(true).write(true).create(true); + + if cfg.direct { + opts.custom_flags(libc::O_DIRECT); + } + + let file = opts.open(&cfg.file)?; + let fd = file.as_raw_fd(); + + let file_size = file.metadata()?.len(); + + if file_size < cfg.total_bytes { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "file size ({}) < total_bytes ({}). Preallocate file first.", + file_size, cfg.total_bytes + ), + )); + } + + let total_ops = cfg.total_bytes / cfg.block_size as u64; + + let max_blocks = file_size / cfg.block_size as u64; + + let engine = UringIoEngine::new( + fd, + UringConfig { + threads: cfg.threads, + io_depth: cfg.io_depth, + ..Default::default() + }, + )?; + + let mut rng = StdRng::seed_from_u64(cfg.seed); + + let mut histogram = Histogram::::new_with_bounds(1, 60_000_000, 3).unwrap(); + + let mut inflight: Vec = Vec::with_capacity(max_inflight); + + let mut buffer_pool: Vec> = (0..max_inflight) + .map(|_| vec![0u8; cfg.block_size]) + .collect(); + + let mut submitted = 0u64; + let mut completed = 0u64; + + let start = Instant::now(); + + while completed < total_ops { + // Submit + while submitted < total_ops && inflight.len() < max_inflight { + let block_idx = rng.random_range(0..max_blocks); + let offset = block_idx * cfg.block_size as u64; + + let is_write = rng.random_bool(cfg.write_ratio); + + let mut buf = buffer_pool.pop().expect("buffer pool empty"); + + let rx = if is_write { + engine.writev_at_async(vec![(buf.as_ptr(), buf.len())], offset)? + } else { + engine.readv_at_async(vec![(buf.as_mut_ptr(), buf.len())], offset)? + }; + + inflight.push(Pending { + start: Instant::now(), + rx, + buf, + }); + + submitted += 1; + } + + // Poll completions + let mut i = 0; + while i < inflight.len() { + match inflight[i].rx.try_recv() { + Ok(res) => { + let pending = inflight.swap_remove(i); + + match res { + Ok(n) if n == cfg.block_size => { + let latency = pending.start.elapsed().as_micros() as u64; + let latency = latency.max(1); + + histogram.record(latency).unwrap(); + } + Ok(n) => { + return Err(io::Error::other(format!( + "Partial IO detected: {} bytes", + n + ))); + } + Err(e) => { + return Err(io::Error::other(format!("IO error: {:?}", e))); + } + } + + buffer_pool.push(pending.buf); + completed += 1; + } + Err(oneshot::error::TryRecvError::Empty) => { + i += 1; + } + Err(oneshot::error::TryRecvError::Closed) => { + return Err(io::Error::other("Completion channel closed unexpectedly")); + } + } + } + + if inflight.len() == max_inflight { + std::thread::sleep(Duration::from_micros(10)); + } + } + + let elapsed = start.elapsed().as_secs_f64(); + + let iops = completed as f64 / elapsed; + let throughput_mb = (completed * cfg.block_size as u64) as f64 / (1024.0 * 1024.0) / elapsed; + + println!("\n==== Results ===="); + println!("Time: {:.2} s", elapsed); + println!("Total Ops: {}", completed); + println!("IOPS: {:.2}", iops); + println!("Throughput: {:.2} MB/s", throughput_mb); + + println!("\nLatency (µs):"); + println!("p50: {}", histogram.value_at_quantile(0.50)); + println!("p95: {}", histogram.value_at_quantile(0.95)); + println!("p99: {}", histogram.value_at_quantile(0.99)); + + Ok(()) +} diff --git a/pegaflow-core/src/backing/mod.rs b/pegaflow-core/src/backing/mod.rs index 5d73da33..e4a13771 100644 --- a/pegaflow-core/src/backing/mod.rs +++ b/pegaflow-core/src/backing/mod.rs @@ -2,7 +2,7 @@ pub(super) mod rdma; pub(super) mod rdma_fetch; pub(super) mod ssd; pub(super) mod ssd_cache; -pub(super) mod uring; +pub mod uring; use std::sync::Arc; diff --git a/pegaflow-core/src/backing/uring.rs b/pegaflow-core/src/backing/uring.rs index 67c1aa65..8c339275 100644 --- a/pegaflow-core/src/backing/uring.rs +++ b/pegaflow-core/src/backing/uring.rs @@ -24,7 +24,7 @@ use tokio::sync::oneshot; /// Configuration for io_uring engine. #[derive(Debug, Clone)] -pub(super) struct UringConfig { +pub struct UringConfig { pub threads: usize, pub io_depth: usize, /// Enable SQ polling; requires kernel support. Off by default. @@ -196,14 +196,14 @@ impl UringShard { } /// io_uring based engine for single-file read/write. -pub(super) struct UringIoEngine { +pub struct UringIoEngine { txs: Vec>, #[allow(dead_code)] handles: Vec>, } impl UringIoEngine { - pub(super) fn new(fd: RawFd, cfg: UringConfig) -> io::Result { + pub fn new(fd: RawFd, cfg: UringConfig) -> io::Result { if cfg.threads == 0 { return Err(io::Error::new( io::ErrorKind::InvalidInput, @@ -253,7 +253,7 @@ impl UringIoEngine { /// /// # Safety /// Caller must ensure all buffer pointers remain valid until the returned receiver completes. - pub(super) fn readv_at_async( + pub fn readv_at_async( &self, iovecs: Vec<(*mut u8, usize)>, offset: u64, @@ -300,7 +300,7 @@ impl UringIoEngine { /// /// # Safety /// Caller must ensure all buffer pointers remain valid until the returned receiver completes. - pub(super) fn writev_at_async( + pub fn writev_at_async( &self, iovecs: Vec<(*const u8, usize)>, offset: u64, diff --git a/pegaflow-core/src/lib.rs b/pegaflow-core/src/lib.rs index 36107e6a..f74d9031 100644 --- a/pegaflow-core/src/lib.rs +++ b/pegaflow-core/src/lib.rs @@ -11,7 +11,7 @@ mod trace; mod allocator; -mod backing; +pub mod backing; mod block; mod cache; mod gpu_worker;