Skip to main content

skm_cli/commands/
bench.rs

1//! Bench command: benchmark selection performance.
2
3use clap::Args;
4use std::path::PathBuf;
5use std::time::{Duration, Instant};
6
7use skm_core::SkillRegistry;
8use skm_select::{CascadeSelector, SelectionContext, TriggerStrategy};
9
10#[derive(Args)]
11pub struct BenchArgs {
12    /// Queries to benchmark (one per line in file, or comma-separated)
13    queries: Vec<String>,
14
15    /// File containing queries (one per line)
16    #[arg(short, long)]
17    file: Option<PathBuf>,
18
19    /// Skill directories
20    #[arg(short, long, default_value = ".")]
21    skills: Vec<PathBuf>,
22
23    /// Number of iterations
24    #[arg(short, long, default_value = "100")]
25    iterations: usize,
26
27    /// Warmup iterations
28    #[arg(long, default_value = "10")]
29    warmup: usize,
30
31    /// Output format
32    #[arg(long, default_value = "text")]
33    format: String,
34}
35
36pub async fn bench(args: BenchArgs) -> anyhow::Result<()> {
37    // Collect queries
38    let mut queries = args.queries.clone();
39    if let Some(ref file) = args.file {
40        let content = std::fs::read_to_string(file)?;
41        queries.extend(content.lines().map(|l| l.trim().to_string()).filter(|l| !l.is_empty()));
42    }
43
44    if queries.is_empty() {
45        anyhow::bail!("No queries provided. Use positional args or --file.");
46    }
47
48    // Load skills
49    let registry = SkillRegistry::new(&args.skills).await?;
50    let trigger = TriggerStrategy::from_registry(&registry).await?;
51    let selector = CascadeSelector::builder().with_triggers(trigger).build();
52    let ctx = SelectionContext::new();
53
54    // Warmup
55    println!("Warming up ({} iterations)...", args.warmup);
56    for _ in 0..args.warmup {
57        for query in &queries {
58            let _ = selector.select(query, &registry, &ctx).await;
59        }
60    }
61
62    // Benchmark
63    println!("Benchmarking ({} iterations)...", args.iterations);
64    let mut latencies = Vec::with_capacity(args.iterations * queries.len());
65
66    let bench_start = Instant::now();
67    for _ in 0..args.iterations {
68        for query in &queries {
69            let start = Instant::now();
70            let _ = selector.select(query, &registry, &ctx).await;
71            latencies.push(start.elapsed());
72        }
73    }
74    let total_time = bench_start.elapsed();
75
76    // Calculate stats
77    latencies.sort();
78    let p50 = latencies[latencies.len() / 2];
79    let p95 = latencies[(latencies.len() as f64 * 0.95) as usize];
80    let p99 = latencies[(latencies.len() as f64 * 0.99) as usize];
81    let min = latencies[0];
82    let max = latencies[latencies.len() - 1];
83    let avg: Duration = latencies.iter().sum::<Duration>() / latencies.len() as u32;
84    let throughput = latencies.len() as f64 / total_time.as_secs_f64();
85
86    match args.format.as_str() {
87        "json" => {
88            let output = serde_json::json!({
89                "queries": queries.len(),
90                "iterations": args.iterations,
91                "total_selections": latencies.len(),
92                "total_time_ms": total_time.as_millis(),
93                "throughput_per_sec": throughput,
94                "latency": {
95                    "min_us": min.as_micros(),
96                    "max_us": max.as_micros(),
97                    "avg_us": avg.as_micros(),
98                    "p50_us": p50.as_micros(),
99                    "p95_us": p95.as_micros(),
100                    "p99_us": p99.as_micros()
101                }
102            });
103            println!("{}", serde_json::to_string_pretty(&output)?);
104        }
105        _ => {
106            println!("\nBenchmark Results:");
107            println!("  Queries: {}", queries.len());
108            println!("  Iterations: {}", args.iterations);
109            println!("  Total selections: {}", latencies.len());
110            println!("  Total time: {:?}", total_time);
111            println!("  Throughput: {:.1} selections/sec", throughput);
112            println!("\nLatency:");
113            println!("  Min: {:?}", min);
114            println!("  Max: {:?}", max);
115            println!("  Avg: {:?}", avg);
116            println!("  P50: {:?}", p50);
117            println!("  P95: {:?}", p95);
118            println!("  P99: {:?}", p99);
119        }
120    }
121
122    Ok(())
123}