skm_cli/commands/
bench.rs1use 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: Vec<String>,
14
15 #[arg(short, long)]
17 file: Option<PathBuf>,
18
19 #[arg(short, long, default_value = ".")]
21 skills: Vec<PathBuf>,
22
23 #[arg(short, long, default_value = "100")]
25 iterations: usize,
26
27 #[arg(long, default_value = "10")]
29 warmup: usize,
30
31 #[arg(long, default_value = "text")]
33 format: String,
34}
35
36pub async fn bench(args: BenchArgs) -> anyhow::Result<()> {
37 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 let registry = SkillRegistry::new(&args.skills).await?;
50 let trigger = TriggerStrategy::from_registry(®istry).await?;
51 let selector = CascadeSelector::builder().with_triggers(trigger).build();
52 let ctx = SelectionContext::new();
53
54 println!("Warming up ({} iterations)...", args.warmup);
56 for _ in 0..args.warmup {
57 for query in &queries {
58 let _ = selector.select(query, ®istry, &ctx).await;
59 }
60 }
61
62 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, ®istry, &ctx).await;
71 latencies.push(start.elapsed());
72 }
73 }
74 let total_time = bench_start.elapsed();
75
76 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}