aster_bench/
bench_work_dir.rs

1use anyhow::Context;
2use chrono::Local;
3use include_dir::{include_dir, Dir};
4use serde::{Deserialize, Serialize};
5use std::fs;
6use std::io;
7use std::path::Path;
8use std::path::PathBuf;
9use std::process::Command;
10
11pub static BUILTIN_EVAL_ASSETS: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/assets");
12
13#[derive(Clone, Serialize, Deserialize, Debug)]
14pub struct BenchmarkWorkDir {
15    pub base_path: PathBuf,
16    pub run_dir: PathBuf,
17    pub cwd: PathBuf,
18    pub run_id: Option<String>,
19}
20
21impl Default for BenchmarkWorkDir {
22    fn default() -> Self {
23        Self::new("work_dir".to_string(), Vec::new())
24    }
25}
26
27impl BenchmarkWorkDir {
28    pub fn new(work_dir_name: String, include_dirs: Vec<PathBuf>) -> Self {
29        let run_dir = std::env::current_dir().unwrap().canonicalize().unwrap();
30        let base_path = PathBuf::from(format!("./{}", work_dir_name));
31        fs::create_dir_all(&base_path).unwrap();
32
33        let base_path = PathBuf::from(&base_path).canonicalize().unwrap();
34
35        // abs paths from dir-strings
36        let dirs = Self::canonical_dirs(include_dirs);
37
38        // deep copy each dir
39        let _: Vec<_> = dirs
40            .iter()
41            .map(|d| Self::deep_copy(d.as_path(), base_path.as_path(), true))
42            .collect();
43
44        Self::copy_auto_included_dirs(&base_path);
45
46        std::env::set_current_dir(&base_path).unwrap();
47
48        BenchmarkWorkDir {
49            base_path: base_path.clone(),
50            run_dir,
51            cwd: base_path.clone(),
52            run_id: None,
53        }
54    }
55
56    pub fn init_experiment(output_dir: PathBuf) -> anyhow::Result<()> {
57        if !output_dir.is_absolute() {
58            anyhow::bail!(
59                "Internal Error: init_experiment received a non-absolute path: {}",
60                output_dir.display()
61            );
62        }
63
64        // create experiment folder
65        let current_time = Local::now().format("%H:%M:%S").to_string();
66        let current_date = Local::now().format("%Y-%m-%d").to_string();
67        let exp_folder_name = format!("benchmark-{}-{}", &current_date, &current_time);
68        let base_path = output_dir.join(exp_folder_name);
69
70        fs::create_dir_all(&base_path).with_context(|| {
71            format!(
72                "Failed to create benchmark directory: {}",
73                base_path.display()
74            )
75        })?;
76        std::env::set_current_dir(&base_path).with_context(|| {
77            format!(
78                "Failed to change working directory to: {}",
79                base_path.display()
80            )
81        })?;
82        Ok(())
83    }
84
85    pub fn canonical_dirs(include_dirs: Vec<PathBuf>) -> Vec<PathBuf> {
86        include_dirs
87            .iter()
88            .map(|d| {
89                let canon = d.canonicalize();
90                if canon.is_err() {
91                    eprintln!("{:?} can't be canonicalized", d);
92                    panic!();
93                }
94                canon.unwrap()
95            })
96            .collect::<Vec<_>>()
97    }
98    fn copy_auto_included_dirs(dest: &Path) {
99        let mut assets_dest = dest.to_path_buf();
100        assets_dest.push("assets");
101        if !assets_dest.exists() {
102            fs::create_dir_all(&assets_dest).unwrap();
103        }
104        BUILTIN_EVAL_ASSETS.extract(assets_dest).unwrap();
105    }
106    pub fn cd(&mut self, path: PathBuf) -> anyhow::Result<&mut Self> {
107        fs::create_dir_all(&path)?;
108        std::env::set_current_dir(&path)?;
109        self.cwd = path;
110        Ok(self)
111    }
112    pub(crate) fn _run_dir(&mut self) -> Option<PathBuf> {
113        if let Some(run_id) = &self.run_id {
114            let mut eval_dir = self.base_path.clone();
115            eval_dir.push(run_id);
116            return Some(eval_dir);
117        }
118        None
119    }
120
121    pub fn set_eval(&mut self, eval: &str, run_id: String) {
122        self.run_id = Some(run_id.clone());
123
124        let eval = eval.replace(":", std::path::MAIN_SEPARATOR_STR);
125        let mut eval_dir = self.base_path.clone();
126        eval_dir.push(run_id);
127        eval_dir.push(eval);
128
129        self.cd(eval_dir.clone())
130            .unwrap_or_else(|_| panic!("Failed to execute cd into {}", eval_dir.clone().display()));
131    }
132
133    fn chop_relative_base<P: AsRef<Path>>(path: P) -> anyhow::Result<PathBuf> {
134        let path = path.as_ref();
135
136        // Get the path components as an iterator
137        let mut components = path.components();
138
139        // Check the first component
140        if let Some(first) = components.next() {
141            use std::path::Component;
142
143            match first {
144                Component::ParentDir => Err(anyhow::anyhow!("RelativePathBaseError: Only paths relative to the current working directory are supported.")),
145                // If first component is "."
146                Component::CurDir => Ok(components.collect()),
147                // Otherwise, keep the full path
148                _ => {
149                    // Create a new PathBuf
150                    let mut result = PathBuf::new();
151                    // Add back the first component
152                    result.push(first);
153                    // Add all remaining components
154                    result.extend(components);
155                    Ok(result)
156                }
157            }
158        } else {
159            // Empty path
160            Ok(PathBuf::new())
161        }
162    }
163
164    pub fn fs_get(&mut self, path: String) -> anyhow::Result<PathBuf> {
165        let p = PathBuf::from(&path);
166        if p.exists() {
167            return Ok(PathBuf::from(path));
168        }
169
170        if p.is_absolute() {
171            return Err(anyhow::anyhow!("AbsolutePathError: Only paths relative to the current working directory are supported."));
172        }
173
174        let asset_rel_path = Self::chop_relative_base(p.clone())
175            .unwrap_or_else(|_| panic!("AbsolutePathError: Only paths relative to the current working directory are supported."));
176
177        let here = PathBuf::from(".").canonicalize()?;
178        let artifact_at_root = self.base_path.clone().join(asset_rel_path);
179
180        Self::deep_copy(artifact_at_root.as_path(), here.as_path(), true)?;
181        Ok(PathBuf::from(path))
182    }
183
184    pub(crate) fn deep_copy<P, Q>(src: P, dst: Q, recursive: bool) -> io::Result<()>
185    where
186        P: AsRef<Path>,
187        Q: AsRef<Path>,
188    {
189        let src = src.as_ref();
190        let dst = dst.as_ref();
191
192        let mut cmd = Command::new("cp");
193
194        // Add -r flag if recursive is true
195        if recursive {
196            cmd.arg("-r");
197        }
198
199        // Add source and destination paths
200        cmd.arg(src).arg(dst);
201
202        // Execute the command
203        let output = cmd.output()?;
204
205        if output.status.success() {
206            Ok(())
207        } else {
208            let error_message = String::from_utf8_lossy(&output.stderr).to_string();
209            Err(io::Error::other(error_message))
210        }
211    }
212
213    pub fn save(&self) {
214        let work_dir = serde_json::to_string_pretty(&self).unwrap();
215        fs::write("work_dir.json", work_dir).expect("Unable to write work-dir as file");
216    }
217}
218
219impl Drop for BenchmarkWorkDir {
220    fn drop(&mut self) {
221        std::env::set_current_dir(&self.run_dir).unwrap();
222    }
223}
aster_bench/bench_work_dir.rs

aster_bench/
bench_work_dir.rs
