gosh_model/
blackbox.rs

1// [[file:../models.note::*header][header:1]]
2//! Represents an universal blackbox (external) model defined by user scripts
3//!
4//! # Usage
5//!
6//! ```ignore
7//! use gosh::models::*;
8//! 
9//! // initialize blackbox model from directory
10//! let dir = "/share/apps/mopac/sp";
11//! let bbm = BlackBoxModel::from_dir(dir)?;
12//! 
13//! // calculate one molecule
14//! let mp = bbm.compute(&mol)?;
15//! 
16//! // calculate a list of molecules
17//! let mp_all = bbm.compute_bunch(&mols)?;
18//! ```
19// header:1 ends here
20
21// [[file:../models.note::c3765387][c3765387]]
22use std::path::{Path, PathBuf};
23use tempfile::TempDir;
24
25use super::*;
26use gchemol::Molecule;
27// c3765387 ends here
28
29// [[file:../models.note::bd430804][bd430804]]
30mod cmd;
31// bd430804 ends here
32
33// [[file:../models.note::*base][base:1]]
34pub struct BlackBoxModel {
35    /// Set the run script file for calculation.
36    run_file: PathBuf,
37
38    /// Set the template file for rendering molecule.
39    tpl_file: PathBuf,
40
41    /// The script for interaction with the main process
42    int_file: Option<PathBuf>,
43
44    /// Set the root directory for scratch files.
45    scr_dir: Option<PathBuf>,
46
47    /// Job starting directory
48    job_dir: Option<PathBuf>,
49
50    // the field order matters
51    // https://stackoverflow.com/questions/41053542/forcing-the-order-in-which-struct-fields-are-dropped
52    task: Option<Task>,
53
54    /// unique temporary working directory
55    temp_dir: Option<TempDir>,
56
57    /// Record the number of potential evalulations.
58    ncalls: usize,
59}
60// base:1 ends here
61
62// [[file:../models.note::045f62c4][045f62c4]]
63// NOTE: There is no implementation of Drop for std::process::Child
64/// A simple wrapper for killing child process on drop
65struct Task(std::process::Child);
66
67impl Drop for Task {
68    // NOTE: There is no implementation of Drop for std::process::Child
69    fn drop(&mut self) {
70        info!("Task dropped. Kill external commands in session.");
71        let child = &mut self.0;
72
73        if let Ok(Some(x)) = child.try_wait() {
74            info!("child process exited gracefully with status {x:?}.");
75        } else {
76            // inform child to exit gracefully
77            if let Err(e) = send_signal_term(child.id()) {
78                error!("Kill child process failure: {:?}", e);
79            }
80            std::thread::sleep(std::time::Duration::from_secs_f64(0.1));
81            // wait a few seconds for child process to exit, or the scratch
82            // directory will be removed immediately.
83            if let Ok(None) = child.try_wait() {
84                info!("Child process is still running, one second for clean up ...",);
85                std::thread::sleep(std::time::Duration::from_secs(1));
86            }
87        }
88    }
89}
90
91fn send_signal_term(pid: u32) -> Result<()> {
92    use nix::sys::signal::{kill, Signal};
93
94    let pid = nix::unistd::Pid::from_raw(pid as i32);
95    let signal = Signal::SIGTERM;
96    info!("Inform child process {} to exit by sending signal {:?}.", pid, signal);
97    kill(pid, signal).with_context(|| format!("kill process {:?}", pid))?;
98
99    Ok(())
100}
101// 045f62c4 ends here
102
103// [[file:../models.note::6cc8ead1][6cc8ead1]]
104mod env {
105    use super::*;
106    use tempfile::{tempdir, tempdir_in};
107
108    /// Return a temporary directory under `BBM_SCR_ROOT` for safe calculation.
109    fn new_scratch_directory(scr_root: Option<&Path>) -> Result<TempDir> {
110        // create leading directories
111        if let Some(d) = &scr_root {
112            if !d.exists() {
113                std::fs::create_dir_all(d).context("create scratch root dir")?;
114            }
115        }
116        scr_root.map_or_else(
117            || tempdir().context("create temp scratch dir"),
118            |d| tempdir_in(d).with_context(|| format!("create temp scratch dir under {:?}", d)),
119        )
120    }
121
122    impl BlackBoxModel {
123        /// Create a temporary working directory and prepare running script
124        pub(super) fn prepare_compute_env(&mut self) -> Result<PathBuf> {
125            let run = "run";
126
127            // create run script if it is not ready
128            let runfile = if let Some(tdir) = &self.temp_dir {
129                tdir.path().join(run)
130            } else {
131                let tdir = new_scratch_directory(self.scr_dir.as_deref())?;
132                debug!("BBM scratching directory: {:?}", tdir);
133
134                // copy run script to work/scratch directory
135                let dest = tdir.path().join(run);
136                let txt = gut::fs::read_file(&self.run_file)?;
137                gut::fs::write_script_file(&dest, &txt)?;
138
139                // save temp dir for next execution
140                self.temp_dir = tdir.into();
141                dest.canonicalize()?
142            };
143
144            Ok(runfile)
145        }
146
147        pub(super) fn from_dotenv(dir: &Path) -> Result<Self> {
148            // canonicalize the file paths
149            let dir = dir
150                .canonicalize()
151                .with_context(|| format!("invalid template directory: {:?}", dir))?;
152
153            // read environment variables from .env config if any
154            let envfile = envfile::EnvFile::new(dir.join(".env")).unwrap();
155            for (key, value) in &envfile.store {
156                debug!("found env var from {:?}: {}={}", &envfile.path, key, value);
157            }
158
159            let run_file = envfile.get("BBM_RUN_FILE").unwrap_or("submit.sh");
160            let tpl_file = envfile.get("BBM_TPL_FILE").unwrap_or("input.hbs");
161            let int_file_opt = envfile.get("BBM_INT_FILE");
162            let bbm = BlackBoxModel {
163                run_file: dir.join(run_file),
164                tpl_file: dir.join(tpl_file),
165                int_file: int_file_opt.map(|f| dir.join(f)),
166                scr_dir: envfile.get("BBM_SCR_DIR").map(|x| x.into()),
167                job_dir: std::env::current_dir()?.into(),
168                temp_dir: None,
169                task: None,
170                ncalls: 0,
171            };
172            Ok(bbm)
173        }
174
175        // Construct from environment variables
176        // 2020-09-05: it is dangerous if we have multiple BBMs in the sample process
177        // fn from_env() -> Self {
178        //     match envy::prefixed("BBM_").from_env::<BlackBoxModel>() {
179        //         Ok(bbm) => bbm,
180        //         Err(error) => panic!("{:?}", error),
181        //     }
182        // }
183    }
184
185    #[test]
186    fn test_env() -> Result<()> {
187        let d = new_scratch_directory(Some("/scratch/test".as_ref()))?;
188        assert!(d.path().exists());
189        let d = new_scratch_directory(None)?;
190        assert!(d.path().exists());
191        Ok(())
192    }
193}
194// 6cc8ead1 ends here
195
196// [[file:../models.note::360435b0][360435b0]]
197impl BlackBoxModel {
198    fn compute_normal(&mut self, mol: &Molecule) -> Result<Computed> {
199        // 1. render input text with the template
200        let txt = self.render_input(&mol)?;
201
202        // 2. call external engine
203        let output = self.submit_cmd(&txt)?;
204
205        // 3. collect model properties
206        let mp = output
207            .parse()
208            .with_context(|| format!("failed to parse computed results: {:?}", output))?;
209
210        self.ncalls += 1;
211        Ok(mp)
212    }
213
214    fn compute_normal_bunch(&mut self, mols: &[Molecule]) -> Result<Vec<Computed>> {
215        // 1. render input text with the template
216        let txt = self.render_input_bunch(mols)?;
217
218        // 2. call external engine
219        let output = self.submit_cmd(&txt)?;
220
221        // 3. collect model properties
222        let all = Computed::parse_all(&output)?;
223
224        self.ncalls += 1;
225        Ok(all)
226    }
227}
228// 360435b0 ends here
229
230// [[file:../models.note::5f8be97b][5f8be97b]]
231impl BlackBoxModel {
232    /// Render `mol` to input string using this template.
233    pub fn render_input(&self, mol: &Molecule) -> Result<String> {
234        // check NaN values in positions
235        for (i, a) in mol.atoms() {
236            let p = a.position();
237            if p.iter().any(|x| x.is_nan()) {
238                error!("Invalid position of atom {}: {:?}", i, p);
239                bail!("Molecule has invalid data in positions.");
240            }
241        }
242        // render input text with external template file
243        let txt = mol.render_with(&self.tpl_file)?;
244
245        Ok(txt)
246    }
247
248    /// Render input using template in bunch mode.
249    pub fn render_input_bunch(&self, mols: &[Molecule]) -> Result<String> {
250        let mut txt = String::new();
251        for mol in mols.iter() {
252            let part = self.render_input(&mol)?;
253            txt.push_str(&part);
254        }
255
256        Ok(txt)
257    }
258}
259// 5f8be97b ends here
260
261// [[file:../models.note::*pub/methods][pub/methods:1]]
262impl BlackBoxModel {
263    /// Construct BlackBoxModel model under directory context.
264    pub fn from_dir<P: AsRef<Path>>(dir: P) -> Result<Self> {
265        Self::from_dotenv(dir.as_ref()).context("Initialize BlackBoxModel failure.")
266    }
267
268    /// keep scratch files for user inspection of failure.
269    pub fn keep_scratch_files(self) {
270        if let Some(tdir) = self.temp_dir {
271            let path = tdir.into_path();
272            println!("Directory for scratch files: {}", path.display());
273        } else {
274            warn!("No temp dir found.");
275        }
276    }
277
278    /// Return the number of potentail evaluations
279    pub fn number_of_evaluations(&self) -> usize {
280        self.ncalls
281    }
282}
283// pub/methods:1 ends here
284
285// [[file:../models.note::5ff4e3f1][5ff4e3f1]]
286impl ChemicalModel for BlackBoxModel {
287    fn compute(&mut self, mol: &Molecule) -> Result<Computed> {
288        let mp = self.compute_normal(mol)?;
289
290        // sanity checking: the associated structure should have the same number
291        // of atoms
292        debug_assert!({
293            let n = mol.natoms();
294            if let Some(pmol) = mp.get_molecule() {
295                pmol.natoms() == n
296            } else {
297                true
298            }
299        });
300
301        Ok(mp)
302    }
303
304    fn compute_bunch(&mut self, mols: &[Molecule]) -> Result<Vec<Computed>> {
305        let all = self.compute_normal_bunch(mols)?;
306
307        // one-to-one mapping
308        debug_assert_eq!(mols.len(), all.len());
309        Ok(all)
310    }
311}
312// 5ff4e3f1 ends here
313
314// [[file:../models.note::ba896ae9][ba896ae9]]
315#[test]
316fn test_bbm() -> Result<()> {
317    // setup two BBMs
318    let bbm_vasp = "./tests/files/vasp-sp";
319    let bbm_siesta = "./tests/files/siesta-sp";
320    let vasp = BlackBoxModel::from_dir(bbm_vasp)?;
321    let siesta = BlackBoxModel::from_dir(bbm_siesta)?;
322
323    // VASP uses input.tera as the input template
324    assert!(vasp.tpl_file.ends_with("input.tera"));
325    // VASP uses input.hbs as the input template
326    assert!(siesta.tpl_file.ends_with("input.hbs"));
327
328    Ok(())
329}
330// ba896ae9 ends here
gosh_model/blackbox.rs

gosh_model/
blackbox.rs
