kermit_bench/
manager.rs

1use {
2    super::{benchmark::Benchmark, downloader::Downloader},
3    std::path::PathBuf,
4};
5
6pub struct DatasetManager {
7    dir: PathBuf,
8    datasets: Vec<Box<dyn Benchmark + 'static>>,
9}
10
11impl DatasetManager {
12    pub fn new<P: Into<PathBuf>>(dataset_dir: P) -> Self {
13        let path = dataset_dir.into();
14        if !path.exists() {
15            std::fs::create_dir_all(&path).expect("Failed to create dataset directory");
16        }
17        Self {
18            dir: path,
19            datasets: vec![],
20        }
21    }
22
23    pub fn init_dataset(
24        &mut self, dataset: impl Benchmark + 'static,
25    ) -> Result<(), Box<dyn std::error::Error>> {
26        let dl_spec = &dataset.metadata().download_spec;
27        let source = Downloader::download(dl_spec)?;
28        dataset.load(&source, self.dir.as_path())?;
29        Downloader::clean(dl_spec);
30        self.datasets.push(Box::new(dataset));
31        Ok(())
32    }
33}