mooncake_bench.rs 23.6 KB
Newer Older
Yan Ru Pei's avatar
Yan Ru Pei committed
1
2
3
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

4
5
6
7
#[path = "common/mod.rs"]
mod common;
use common::*;

Yan Ru Pei's avatar
Yan Ru Pei committed
8
9
use clap::{Parser, Subcommand};
use dynamo_kv_router::LocalBlockHash;
10
use dynamo_kv_router::indexer::{KvIndexer, KvIndexerInterface, KvIndexerMetrics};
11
use dynamo_kv_router::protocols::{KvCacheEvent, KvCacheEventData, RouterEvent};
12
13
14
use dynamo_kv_router::{
    ConcurrentRadixTree, ConcurrentRadixTreeCompressed, PositionalIndexer, ThreadPoolIndexer,
};
15
use dynamo_mocker::loadgen::Trace;
16
use serde::Serialize;
Yan Ru Pei's avatar
Yan Ru Pei committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
use std::sync::Arc;
use tokio::time::{Duration, Instant};
use tokio_util::sync::CancellationToken;

/// Indexer backend selection and its backend-specific parameters.
#[derive(Subcommand, Debug, Clone)]
enum IndexerArgs {
    /// Single-threaded radix tree indexer.
    RadixTree {},

    /// Position-based nested map indexer with jump search.
    NestedMap {
        /// Number of positions to skip during jump search before scanning back.
        #[clap(long, default_value = "8")]
        jump_size: usize,

        /// Number of OS threads that consume and apply KV cache events.
        #[clap(long, default_value = "16")]
        num_event_workers: usize,
    },

    /// Lock-based concurrent radix tree indexer.
    ConcurrentRadixTree {
        /// Number of OS threads that consume and apply KV cache events.
        #[clap(long, default_value = "16")]
        num_event_workers: usize,
    },
44
45
46
47
48
49
50

    /// Compressed concurrent radix tree indexer (compressed edges).
    ConcurrentRadixTreeCompressed {
        /// Number of OS threads that consume and apply KV cache events.
        #[clap(long, default_value = "16")]
        num_event_workers: usize,
    },
Yan Ru Pei's avatar
Yan Ru Pei committed
51
52
53
54
}

impl IndexerArgs {
    /// Construct the concrete indexer from the parsed CLI args.
55
    fn build(self, block_size: u32) -> Arc<dyn KvIndexerInterface + Send + Sync> {
Yan Ru Pei's avatar
Yan Ru Pei committed
56
57
58
59
        let cancel_token = CancellationToken::new();
        let metrics = Arc::new(KvIndexerMetrics::new_unregistered());
        match self {
            IndexerArgs::RadixTree {} => {
60
                Arc::new(KvIndexer::new(cancel_token, block_size, metrics))
Yan Ru Pei's avatar
Yan Ru Pei committed
61
62
63
64
65
66
67
            }
            IndexerArgs::NestedMap {
                jump_size,
                num_event_workers,
            } => Arc::new(ThreadPoolIndexer::new(
                PositionalIndexer::new(jump_size),
                num_event_workers,
68
                block_size,
Yan Ru Pei's avatar
Yan Ru Pei committed
69
            )),
70
71
72
            IndexerArgs::ConcurrentRadixTree { num_event_workers } => Arc::new(
                ThreadPoolIndexer::new(ConcurrentRadixTree::new(), num_event_workers, block_size),
            ),
73
74
75
76
77
78
79
            IndexerArgs::ConcurrentRadixTreeCompressed { num_event_workers } => {
                Arc::new(ThreadPoolIndexer::new(
                    ConcurrentRadixTreeCompressed::new(),
                    num_event_workers,
                    block_size,
                ))
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
80
81
        }
    }
82

83
84
    fn supports_remove(_name: &str) -> bool {
        true
85
86
87
    }

    fn is_multi_threaded(name: &str) -> bool {
88
89
90
91
        matches!(
            name,
            "nested-map" | "concurrent-radix-tree" | "concurrent-radix-tree-compressed"
        )
92
93
    }

94
    /// Construct an indexer from a short name string.
95
96
    fn from_name(
        name: &str,
97
98
        block_size: u32,
        num_event_workers: usize,
99
    ) -> anyhow::Result<Arc<dyn KvIndexerInterface + Send + Sync>> {
100
        let nw = num_event_workers;
101
102
103
104
105
106
107
108
109
        let indexer_args = match name {
            "radix-tree" => IndexerArgs::RadixTree {},
            "nested-map" => IndexerArgs::NestedMap {
                jump_size: 8,
                num_event_workers: nw,
            },
            "concurrent-radix-tree" => IndexerArgs::ConcurrentRadixTree {
                num_event_workers: nw,
            },
110
111
112
            "concurrent-radix-tree-compressed" => IndexerArgs::ConcurrentRadixTreeCompressed {
                num_event_workers: nw,
            },
113
            _ => anyhow::bail!(
114
                "Unknown indexer '{}'. Valid names: radix-tree, \
115
                 nested-map, concurrent-radix-tree, concurrent-radix-tree-compressed",
116
117
118
                name
            ),
        };
119
        Ok(indexer_args.build(block_size))
120
    }
Yan Ru Pei's avatar
Yan Ru Pei committed
121
122
123
124
125
}

#[derive(Parser, Debug)]
#[clap(version, about, long_about = None)]
struct Args {
126
127
    #[clap(flatten)]
    common: CommonArgs,
Yan Ru Pei's avatar
Yan Ru Pei committed
128

129
    /// Output path for the sweep plot SVG.
130
    #[clap(long, default_value = "sweep_plot.svg")]
131
132
133
134
    sweep_output: String,

    /// Comma-separated list of indexer names to benchmark and compare on the
    /// same plot. Overrides the subcommand indexer when present. Valid names:
135
    /// radix-tree, nested-map, concurrent-radix-tree,
136
    /// concurrent-radix-tree-compressed.
137
138
139
140
    #[clap(long, value_delimiter = ',')]
    compare: Vec<String>,

    /// Number of OS threads for event processing in compare mode. Applies to
141
142
143
    /// indexers that use a thread pool (nested-map, concurrent-radix-tree,
    /// concurrent-radix-tree-compressed).
    /// Ignored by radix-tree.
144
145
146
    #[clap(long, default_value = "16")]
    num_event_workers: usize,

Yan Ru Pei's avatar
Yan Ru Pei committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
    /// Indexer backend to benchmark (defaults to radix-tree if not specified).
    #[clap(subcommand)]
    indexer: Option<IndexerArgs>,
}

impl Args {
    /// Return the indexer config, falling back to RadixTree if none was specified.
    fn get_indexer(&self) -> IndexerArgs {
        self.indexer.clone().unwrap_or(IndexerArgs::RadixTree {})
    }
}

/// A single entry in a worker's merged benchmark timeline.
#[derive(Clone)]
enum WorkerTraceEntry {
    /// A find_matches request with pre-computed block hashes.
    Request(Vec<LocalBlockHash>),
    /// A KV cache event (store/remove/clear) to apply to the indexer.
    Event(KvCacheEvent),
}

/// A timestamped entry in a worker's benchmark trace, used to replay requests
/// and events at the correct relative timing.
#[derive(Clone)]
struct WorkerTrace {
    entry: WorkerTraceEntry,
    timestamp_us: u64,
}

/// Merge each worker's request trace and event trace into a single
/// time-ordered sequence of `WorkerTrace` entries suitable for benchmark
/// replay.
///
/// Timestamps are rescaled from the original trace / simulation durations
/// into the benchmark duration (microseconds).
fn prepare_worker_traces(
183
    artifacts: Vec<WorkerReplayArtifacts>,
184
    benchmark_duration_ms: u64,
Yan Ru Pei's avatar
Yan Ru Pei committed
185
) -> Vec<Vec<WorkerTrace>> {
186
    artifacts
Yan Ru Pei's avatar
Yan Ru Pei committed
187
        .into_iter()
188
189
190
        .map(|artifact| {
            let mut merged = artifact
                .requests
Yan Ru Pei's avatar
Yan Ru Pei committed
191
192
                .into_iter()
                .map(|request| WorkerTrace {
193
194
                    timestamp_us: request.timestamp_us,
                    entry: WorkerTraceEntry::Request(request.replay_hashes.local_block_hashes),
Yan Ru Pei's avatar
Yan Ru Pei committed
195
                })
196
197
198
199
200
                .chain(artifact.kv_events.into_iter().map(|event| WorkerTrace {
                    timestamp_us: event.timestamp_us,
                    entry: WorkerTraceEntry::Event(event.event),
                }))
                .collect::<Vec<_>>();
Yan Ru Pei's avatar
Yan Ru Pei committed
201
            merged.sort_by_key(|entry| entry.timestamp_us);
202
203
204
205
206
207
208
209
            let max_timestamp_us = merged.last().map(|entry| entry.timestamp_us).unwrap_or(0);
            for entry in &mut merged {
                entry.timestamp_us = if max_timestamp_us == 0 {
                    0
                } else {
                    entry.timestamp_us * benchmark_duration_ms * 1000 / max_timestamp_us
                };
            }
Yan Ru Pei's avatar
Yan Ru Pei committed
210
211
212
213
214
            merged
        })
        .collect()
}

215
216
217
218
219
220
221
#[derive(Serialize)]
struct SweepStepResult {
    duration_ms: u64,
    #[serde(flatten)]
    results: BenchmarkResults,
}

Yan Ru Pei's avatar
Yan Ru Pei committed
222
223
224
225
226
227
228
229
/// Run the benchmark: replay each worker's merged trace against the indexer,
/// measuring find_matches latency and event processing throughput.
///
/// Workers are spawned as tokio tasks, each replaying its trace at the
/// original inter-entry timing. After all workers finish, the event queue is
/// flushed and latency percentiles / throughput stats are printed.
async fn run_benchmark(
    indexer: Arc<dyn KvIndexerInterface + Send + Sync>,
230
    artifacts: Vec<WorkerReplayArtifacts>,
Yan Ru Pei's avatar
Yan Ru Pei committed
231
    args: &Args,
232
    benchmark_duration_ms: u64,
233
    count_events: bool,
234
) -> anyhow::Result<BenchmarkResults> {
235
    let worker_traces = prepare_worker_traces(artifacts, benchmark_duration_ms);
236
    let worker_traces = worker_traces.into_iter().map(Arc::new).collect::<Vec<_>>();
Yan Ru Pei's avatar
Yan Ru Pei committed
237
238
239
240
241
242

    let progress = make_progress_bar(Some(
        worker_traces
            .iter()
            .map(|trace| trace.len() as u64)
            .sum::<u64>()
243
            * args.common.inference_worker_duplication_factor as u64,
Yan Ru Pei's avatar
Yan Ru Pei committed
244
245
246
    ));

    let mut tasks = Vec::new();
247
    for replica in 0..args.common.inference_worker_duplication_factor {
Yan Ru Pei's avatar
Yan Ru Pei committed
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
        for (worker_id, worker_trace) in worker_traces.iter().enumerate() {
            let indexer = indexer.clone();
            let trace = worker_trace.clone();
            let progress = progress.clone();
            let worker_id = worker_id + replica * worker_traces.len();
            tasks.push(tokio::spawn(async move {
                let mut request_latencies = Vec::with_capacity(trace.len());

                let submit = |entry: WorkerTrace| async {
                    match entry.entry {
                        WorkerTraceEntry::Request(request) => {
                            let start = minstant::Instant::now();
                            indexer.find_matches(request).await?;
                            Ok::<Option<u64>, anyhow::Error>(
                                Some(start.elapsed().as_nanos() as u64),
                            )
                        }
                        WorkerTraceEntry::Event(event) => {
                            indexer
267
                                .apply_event(RouterEvent::new(worker_id as u64, event))
Yan Ru Pei's avatar
Yan Ru Pei committed
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
                                .await;
                            Ok(None)
                        }
                    }
                };

                let mut target = Instant::now();

                let mut trace = trace.iter().peekable();

                let mut local_count = 0;

                while let Some(entry) = trace.next() {
                    let mut processed = 1;
                    let entry_timestamp_us = entry.timestamp_us;

                    if let Some(latency) = submit(entry.clone()).await? {
                        request_latencies.push(latency);
                    }

                    while let Some(next) = trace.peek() {
                        if next.timestamp_us == entry_timestamp_us {
                            if let Some(latency) = submit(trace.next().unwrap().clone()).await? {
                                request_latencies.push(latency);
                            }
                            processed += 1;
                        } else {
                            break;
                        }
                    }

                    if let Some(next) = trace.peek() {
                        target += Duration::from_micros(next.timestamp_us - entry_timestamp_us);
                    }

                    if target > Instant::now() {
                        tokio::time::sleep_until(target).await;
                    }

                    local_count += processed;

                    if local_count > 100 {
                        progress.inc(local_count);
                        local_count = 0;
                    }
                }

                progress.inc(local_count);

                Ok::<_, anyhow::Error>(request_latencies)
            }));
        }
    }

    let mut latencies = Vec::new();

    for task in tasks {
        latencies.extend(task.await??);
    }

328
    if progress.elapsed() > Duration::from_millis(benchmark_duration_ms * 11 / 10) {
Yan Ru Pei's avatar
Yan Ru Pei committed
329
330
331
332
333
        eprintln!(
            "WARNING: The benchmarker is unable to keep up with the request/event generation rate. Rerun with a larger --benchmark-duration-ms."
        )
    }

334
    let total_duration = progress.elapsed();
Yan Ru Pei's avatar
Yan Ru Pei committed
335
336
337
338
339
340
341
342
343
344

    let total_events = worker_traces
        .iter()
        .map(|trace| {
            trace
                .iter()
                .filter(|trace| matches!(trace.entry, WorkerTraceEntry::Event(_)))
                .count()
        })
        .sum::<usize>()
345
        * args.common.inference_worker_duplication_factor;
Yan Ru Pei's avatar
Yan Ru Pei committed
346
347

    let total_requests = worker_traces.iter().map(|trace| trace.len()).sum::<usize>()
348
        * args.common.inference_worker_duplication_factor
Yan Ru Pei's avatar
Yan Ru Pei committed
349
350
        - total_events;

351
352
353
354
355
356
357
358
    let total_request_blocks: usize = worker_traces
        .iter()
        .flat_map(|t| t.iter())
        .filter_map(|entry| match &entry.entry {
            WorkerTraceEntry::Request(hashes) => Some(hashes.len()),
            _ => None,
        })
        .sum::<usize>()
359
        * args.common.inference_worker_duplication_factor;
Yan Ru Pei's avatar
Yan Ru Pei committed
360

361
362
363
364
365
366
367
368
369
370
371
    let total_event_blocks: usize = worker_traces
        .iter()
        .flat_map(|t| t.iter())
        .filter_map(|entry| match &entry.entry {
            WorkerTraceEntry::Event(ev) => match &ev.data {
                KvCacheEventData::Stored(s) => Some(s.blocks.len()),
                _ => Some(0),
            },
            _ => None,
        })
        .sum::<usize>()
372
        * args.common.inference_worker_duplication_factor;
Yan Ru Pei's avatar
Yan Ru Pei committed
373

374
375
    let counted_events = if count_events { total_events } else { 0 };
    let counted_event_blocks = if count_events { total_event_blocks } else { 0 };
Yan Ru Pei's avatar
Yan Ru Pei committed
376

377
378
    let total_blocks = total_request_blocks + counted_event_blocks;
    let total_ops = total_requests + counted_events;
379
380
381
382
    let offered_ops_throughput = total_ops as f32 / benchmark_duration_ms as f32 * 1000.0;
    let ops_throughput = total_ops as f32 / total_duration.as_millis() as f32 * 1000.0;
    let offered_block_throughput = total_blocks as f32 / benchmark_duration_ms as f32 * 1000.0;
    let block_throughput = total_blocks as f32 / total_duration.as_millis() as f32 * 1000.0;
Yan Ru Pei's avatar
Yan Ru Pei committed
383
384

    latencies.sort_unstable();
385
386
387
388
389
    let latency_p99_us = if latencies.is_empty() {
        0.0
    } else {
        latencies[latencies.len() * 99 / 100] as f32 / 1000.0
    };
390

Yan Ru Pei's avatar
Yan Ru Pei committed
391
    println!(
392
393
        "Ops Throughput: {} ops/s (requests + events)",
        ops_throughput
Yan Ru Pei's avatar
Yan Ru Pei committed
394
    );
395
396
397
398
399
400
401
402
403
404
405
    println!("Block Throughput: {} block ops/s", block_throughput);
    println!("Latency p99: {}us", latency_p99_us);

    Ok(BenchmarkResults {
        offered_ops_throughput,
        ops_throughput,
        offered_block_throughput,
        block_throughput,
        latency_p99_us,
    })
}
Yan Ru Pei's avatar
Yan Ru Pei committed
406

407
async fn run_tests() -> anyhow::Result<()> {
408
    use std::collections::HashSet;
409
    use std::fs::File;
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
    use std::io::Write;

    let path =
        std::env::temp_dir().join(format!("mooncake_bench_test_{}.jsonl", std::process::id()));
    {
        let mut f = File::create(&path)?;
        for (i, (hash_ids, output_length)) in
            [(&[0u64, 1, 2] as &[u64], 10u64), (&[0, 1, 3, 4], 10)]
                .iter()
                .enumerate()
        {
            writeln!(
                f,
                "{}",
                serde_json::json!({
                    "timestamp": i as u64,
426
                    "input_length": hash_ids.len(),
427
428
429
430
431
432
433
                    "hash_ids": hash_ids,
                    "output_length": output_length,
                })
            )?;
        }
    }

434
    let traces = process_mooncake_trace(path.to_str().unwrap(), 512, 2, 2, 2, 42)?;
435
436
437
438
    std::fs::remove_file(&path).ok();

    let mut all_hashes: Vec<Vec<u64>> = traces
        .into_iter()
439
440
        .flat_map(|worker| worker.sessions.into_iter())
        .flat_map(|session| session.turns.into_iter().map(|turn| turn.hash_ids))
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
        .collect();
    all_hashes.sort();

    // expand(2): [0,1,2] → [0,1,2,3,4,5], [0,1,3,4] → [0,1,2,3,6,7,8,9]
    // duplicate(2): max=9, offset=10
    let mut expected = vec![
        vec![0, 1, 2, 3, 4, 5],
        vec![10, 11, 12, 13, 14, 15],
        vec![0, 1, 2, 3, 6, 7, 8, 9],
        vec![10, 11, 12, 13, 16, 17, 18, 19],
    ];
    expected.sort();
    assert_eq!(all_hashes, expected, "hash_ids mismatch");

    // Verify prefix structure within each copy.
    let copy0: Vec<&Vec<u64>> = all_hashes.iter().filter(|h| h[0] == 0).collect();
    let copy1: Vec<&Vec<u64>> = all_hashes.iter().filter(|h| h[0] == 10).collect();
    assert_eq!(copy0.len(), 2);
    assert_eq!(copy1.len(), 2);
    assert_eq!(copy0[0][..4], copy0[1][..4], "copy 0 shared prefix broken");
    assert_eq!(copy1[0][..4], copy1[1][..4], "copy 1 shared prefix broken");

    // Verify disjointness between copies.
    let set0: HashSet<u64> = copy0.iter().flat_map(|h| h.iter().copied()).collect();
    let set1: HashSet<u64> = copy1.iter().flat_map(|h| h.iter().copied()).collect();
    assert!(set0.is_disjoint(&set1), "copies are not hash-disjoint");

468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
    let replay_trace = Trace {
        block_size: 2,
        sessions: vec![dynamo_mocker::loadgen::SessionTrace {
            session_id: "session-a".to_string(),
            first_arrival_timestamp_ms: Some(0.0),
            turns: vec![
                dynamo_mocker::loadgen::TurnTrace {
                    input_length: 4,
                    max_output_tokens: 2,
                    hash_ids: vec![1, 2],
                    delay_after_previous_ms: 0.0,
                },
                dynamo_mocker::loadgen::TurnTrace {
                    input_length: 4,
                    max_output_tokens: 2,
                    hash_ids: vec![3, 4],
                    delay_after_previous_ms: 5.0,
                },
            ],
        }],
    };
    let artifacts = generate_replay_artifacts(&[replay_trace], 1024, 2, 5).await?;
    assert_eq!(artifacts.len(), 1);
    assert_eq!(artifacts[0].requests.len(), 2);
    let first_uuid = artifacts[0].requests[0].uuid;
    let first_completion_ms = artifacts[0]
        .output_signals
        .iter()
        .find(|signal| signal.signal.uuid == first_uuid && signal.signal.completed)
        .expect("first request must complete")
        .timestamp_us as f64
        / 1000.0;
    assert!(
        artifacts[0].requests[1].scheduled_ready_at_ms + 0.1 >= first_completion_ms + 5.0,
        "expected second request to wait for completion plus delay"
    );

505
506
507
508
    println!("All tests passed.");
    Ok(())
}

Yan Ru Pei's avatar
Yan Ru Pei committed
509
510
511
512
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();

513
    if args.common.test {
514
        return run_tests().await;
515
516
    }

517
518
519
520
521
522
523
    let path = match args.common.mooncake_trace_path.as_deref() {
        Some(p) => p,
        None => {
            eprintln!("No mooncake_trace_path provided, skipping benchmark");
            return Ok(());
        }
    };
524
525
    let traces = process_mooncake_trace(
        path,
526
        args.common.block_size,
527
528
529
530
531
        args.common.trace_length_factor,
        args.common.trace_duplication_factor,
        args.common.num_unique_inference_workers,
        args.common.seed,
    )?;
532
    let artifacts = generate_replay_artifacts(
533
534
535
536
537
538
        &traces,
        args.common.num_gpu_blocks,
        args.common.block_size,
        args.common.trace_simulation_duration_ms,
    )
    .await?;
Yan Ru Pei's avatar
Yan Ru Pei committed
539

540
541
542
543
544
    let indexer_names: Vec<String> = if args.compare.is_empty() {
        let name = match args.get_indexer() {
            IndexerArgs::RadixTree {} => "radix-tree",
            IndexerArgs::NestedMap { .. } => "nested-map",
            IndexerArgs::ConcurrentRadixTree { .. } => "concurrent-radix-tree",
545
            IndexerArgs::ConcurrentRadixTreeCompressed { .. } => "concurrent-radix-tree-compressed",
546
547
548
549
550
551
        };
        vec![name.to_string()]
    } else {
        args.compare.clone()
    };

552
553
554
555
556
557
    if args.common.sweep {
        let durations_low_to_high = compute_sweep_durations(
            args.common.sweep_min_ms,
            args.common.sweep_max_ms,
            args.common.sweep_steps,
        );
558
        let durations_high_to_low: Vec<u64> = durations_low_to_high.iter().copied().rev().collect();
559
560
561
562
563
564
565
566

        let mut all_results: Vec<(&str, Vec<(u64, BenchmarkResults)>)> = Vec::new();

        for name in &indexer_names {
            println!("\n{}", "=".repeat(60));
            println!("Benchmarking indexer: {}", name);
            println!("{}", "=".repeat(60));

567
568
569
570
571
572
573
            let multi_threaded = IndexerArgs::is_multi_threaded(name);
            let durations = if multi_threaded {
                &durations_high_to_low
            } else {
                &durations_low_to_high
            };

574
            let mut results: Vec<(u64, BenchmarkResults)> = Vec::new();
575
            let mut consecutive_keeping_up = 0u32;
576

577
            for &dur_ms in durations {
578
579
                println!("\n=== Sweep: benchmark_duration_ms = {} ===", dur_ms);
                let indexer = if args.compare.is_empty() {
580
                    args.get_indexer().build(args.common.block_size)
581
                } else {
582
                    IndexerArgs::from_name(name, args.common.block_size, args.num_event_workers)?
583
                };
584
                let count_events = IndexerArgs::supports_remove(name);
585
586
                let result =
                    run_benchmark(indexer, artifacts.clone(), &args, dur_ms, count_events).await?;
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606

                if multi_threaded {
                    if result.block_throughput >= result.offered_block_throughput * 0.95 {
                        consecutive_keeping_up += 1;
                    } else {
                        consecutive_keeping_up = 0;
                    }
                    results.push((dur_ms, result));
                    if consecutive_keeping_up >= 5 {
                        println!("Early stop: achieved >= 95% offered for 5 consecutive steps");
                        break;
                    }
                } else {
                    let saturated = result.offered_block_throughput > result.block_throughput * 5.0;
                    results.push((dur_ms, result));
                    if saturated {
                        println!("Early stop: offered throughput >5x achieved throughput");
                        break;
                    }
                }
607
608
            }

609
            results.sort_by_key(|(dur, _)| std::cmp::Reverse(*dur));
610
            print_sweep_summary(name, &results);
Yan Ru Pei's avatar
Yan Ru Pei committed
611

612
613
614
615
            all_results.push((name, results));
        }

        plot_sweep(&all_results, &args.sweep_output)?;
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641

        let json_path = args
            .sweep_output
            .replace(".png", ".json")
            .replace(".svg", ".json");
        let json_map: std::collections::BTreeMap<&str, Vec<SweepStepResult>> = all_results
            .iter()
            .map(|(name, results)| {
                let steps = results
                    .iter()
                    .map(|(dur, r)| SweepStepResult {
                        duration_ms: *dur,
                        results: BenchmarkResults {
                            offered_ops_throughput: r.offered_ops_throughput,
                            ops_throughput: r.ops_throughput,
                            offered_block_throughput: r.offered_block_throughput,
                            block_throughput: r.block_throughput,
                            latency_p99_us: r.latency_p99_us,
                        },
                    })
                    .collect();
                (*name, steps)
            })
            .collect();
        std::fs::write(&json_path, serde_json::to_string_pretty(&json_map)?)?;
        println!("Sweep results saved to {}", json_path);
642
643
644
645
    } else {
        for name in &indexer_names {
            println!("\nBenchmarking indexer: {}", name);
            let indexer = if args.compare.is_empty() {
646
                args.get_indexer().build(args.common.block_size)
647
            } else {
648
                IndexerArgs::from_name(name, args.common.block_size, args.num_event_workers)?
649
            };
650
            let count_events = IndexerArgs::supports_remove(name);
651
652
            run_benchmark(
                indexer,
653
                artifacts.clone(),
654
                &args,
655
                args.common.benchmark_duration_ms,
656
                count_events,
657
658
659
660
            )
            .await?;
        }
    }
Yan Ru Pei's avatar
Yan Ru Pei committed
661
662
663

    Ok(())
}