aster_bench/runners/
bench_runner.rs1use crate::bench_config::{BenchModel, BenchRunConfig};
2use crate::bench_work_dir::BenchmarkWorkDir;
3use crate::eval_suites::EvaluationSuite;
4use crate::runners::model_runner::ModelRunner;
5use crate::utilities::{await_process_exits, parallel_bench_cmd};
6use anyhow::Context;
7use std::path::PathBuf;
8
9#[derive(Clone)]
10pub struct BenchRunner {
11 config: BenchRunConfig,
12}
13
14impl BenchRunner {
15 pub fn new(config_path: PathBuf) -> anyhow::Result<BenchRunner> {
16 let config = BenchRunConfig::from(config_path.clone())?;
17
18 let resolved_output_dir = match &config.output_dir {
19 Some(path) => {
20 if !path.is_absolute() {
21 anyhow::bail!(
22 "Config Error in '{}': 'output_dir' must be an absolute path, but found relative path: {}",
23 config_path.display(),
24 path.display()
25 );
26 }
27 path.clone()
28 }
29 None => std::env::current_dir().context(
30 "Failed to get current working directory to use as default output directory",
31 )?,
32 };
33
34 BenchmarkWorkDir::init_experiment(resolved_output_dir)?;
35
36 config.save("config.cfg".to_string());
37 Ok(BenchRunner { config })
38 }
39
40 pub fn from(config: String) -> anyhow::Result<BenchRunner> {
41 let config = BenchRunConfig::from_string(config)?;
42 Ok(BenchRunner { config })
43 }
44
45 pub fn run(&mut self) -> anyhow::Result<()> {
46 let (parallel_models, serial_models): &(Vec<BenchModel>, Vec<BenchModel>) = &self
48 .config
49 .models
50 .clone()
51 .into_iter()
52 .partition(|model| model.parallel_safe);
53
54 let mut parallel_models_handle = Vec::new();
56 for model in parallel_models {
57 self.config.models = vec![model.clone()];
58 let cfg = self.config.to_string()?;
59 let model_handle = parallel_bench_cmd("eval-model".to_string(), cfg, Vec::new());
60 parallel_models_handle.push(model_handle);
61 }
62
63 for model in serial_models {
65 self.config.models = vec![model.clone()];
66 ModelRunner::from(self.config.to_string()?)?.run()?;
67 }
68
69 await_process_exits(&mut parallel_models_handle, Vec::new());
70
71 Ok(())
72 }
73
74 pub fn list_selectors(_config: Option<PathBuf>) -> anyhow::Result<()> {
75 let selector_eval_counts = EvaluationSuite::available_selectors();
76 let mut keys: Vec<_> = selector_eval_counts.keys().collect();
77 keys.sort();
78 let max_key_len = keys.iter().map(|k| k.len()).max().unwrap_or(0);
79 println!(
80 "selector {} => Eval Count",
81 " ".repeat(max_key_len - "selector".len())
82 );
83 println!("{}", "-".repeat(max_key_len + 6));
84 for selector in keys {
85 println!(
86 "{} {} => {}",
87 selector,
88 " ".repeat(max_key_len - selector.len()),
89 selector_eval_counts.get(selector).unwrap()
90 );
91 }
92 Ok(())
93 }
94}