Skip to main content

aster_bench/runners/
bench_runner.rs

1use 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        // split models that must run serial from those that can be run in parallel
47        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        // exec parallel models
55        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        // exec serial models
64        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}