modewise_pipeline/
batch.rs

1use std::collections::BTreeMap;
2
3use rayon::prelude::*;
4
5use crate::{Artifact, ArtifactStore, PipelineError, Plan, RunResult, run};
6
7/// Supplies pipeline inputs for many recordings, for example files found on disk.
8pub trait InputSource: Sync {
9    /// Names of the inputs each recording provides.
10    fn input_names(&self) -> Vec<String>;
11    fn recording_ids(&self) -> Vec<String>;
12    fn load(&self, id: &str) -> Result<BTreeMap<String, Artifact>, String>;
13}
14
15/// Outcome for one recording. A failure here does not stop the others.
16pub struct RecordingRun {
17    pub id: String,
18    pub result: Result<RunResult, String>,
19}
20
21/// Runs a plan once per recording, `workers` recordings at a time (0 = all cores).
22/// `on_done` sees each recording as soon as it finishes, for progress or saving.
23pub fn run_batch(
24    plan: &Plan,
25    source: &dyn InputSource,
26    ids: &[String],
27    store: &dyn ArtifactStore,
28    workers: usize,
29    on_done: &(dyn Fn(&RecordingRun) + Sync),
30) -> Result<Vec<RecordingRun>, PipelineError> {
31    let pool = rayon::ThreadPoolBuilder::new()
32        .num_threads(workers)
33        .build()
34        .map_err(|e| PipelineError::Spec(format!("cannot start {workers} workers: {e}")))?;
35    let runs = pool.install(|| {
36        ids.par_iter()
37            .map(|id| {
38                let result = source
39                    .load(id)
40                    .map_err(|e| format!("loading failed: {e}"))
41                    .and_then(|mut inputs| {
42                        // A source may offer more than this pipeline uses (such as a channel table).
43                        inputs.retain(|name, _| plan.inputs().contains(name));
44                        run(plan, &inputs, store).map_err(|e| e.to_string())
45                    });
46                let run = RecordingRun {
47                    id: id.clone(),
48                    result,
49                };
50                on_done(&run);
51                run
52            })
53            .collect()
54    });
55    Ok(runs)
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61    use crate::PipelineSpec;
62    use crate::test_support::{TestStore, ops, series};
63
64    struct Source;
65
66    impl InputSource for Source {
67        fn input_names(&self) -> Vec<String> {
68            vec!["x".into()]
69        }
70
71        fn recording_ids(&self) -> Vec<String> {
72            vec!["a".into(), "bad".into(), "b".into()]
73        }
74
75        fn load(&self, id: &str) -> Result<BTreeMap<String, Artifact>, String> {
76            match id {
77                "bad" => Err("file missing".into()),
78                _ => Ok(BTreeMap::from([(
79                    "x".into(),
80                    Artifact::Series(series(&[1.0, id.len() as f64])),
81                )])),
82            }
83        }
84    }
85
86    #[test]
87    fn runs_each_recording_and_keeps_going_after_failures() {
88        let spec = PipelineSpec::from_toml(
89            "[pipeline]\nname = \"b\"\ninputs = [\"x\"]\n[[step]]\nid = \"s\"\nop = \"scale\"\ninput = \"x\"\n",
90        )
91        .unwrap();
92        let ops = ops();
93        let plan = Plan::new(&spec, &ops.registry).unwrap();
94        let ids = Source.recording_ids();
95        let done = std::sync::atomic::AtomicUsize::new(0);
96        let count = |_: &RecordingRun| {
97            done.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
98        };
99        let runs = run_batch(&plan, &Source, &ids, &TestStore::default(), 2, &count).unwrap();
100        assert_eq!(done.into_inner(), 3);
101
102        let outcome: Vec<(&str, bool)> = runs
103            .iter()
104            .map(|r| (r.id.as_str(), r.result.is_ok()))
105            .collect();
106        assert_eq!(outcome, [("a", true), ("bad", false), ("b", true)]);
107        assert!(
108            runs[1]
109                .result
110                .as_ref()
111                .err()
112                .unwrap()
113                .contains("file missing")
114        );
115    }
116}
modewise_pipeline/batch.rs

modewise_pipeline/
batch.rs
