modewise_pipeline/
executor.rs

1use std::collections::{BTreeMap, HashMap};
2use std::sync::Arc;
3use std::time::Instant;
4
5use serde_json::json;
6
7use crate::plan::{PlannedStep, Source};
8use crate::{
9    Artifact, ArtifactKey, ArtifactKind, ArtifactStore, Inputs, PipelineError, Plan, RunManifest,
10    StepRecord,
11};
12
13pub struct RunResult {
14    pub manifest: RunManifest,
15    outputs: BTreeMap<String, ArtifactKey>,
16}
17
18impl RunResult {
19    /// All step outputs, by `step.output`.
20    pub fn outputs(&self) -> &BTreeMap<String, ArtifactKey> {
21        &self.outputs
22    }
23
24    /// Key of a step output, by `step.output`, or by `step` when it has one output.
25    pub fn output_key(&self, reference: &str) -> Option<&ArtifactKey> {
26        if let Some(key) = self.outputs.get(reference) {
27            return Some(key);
28        }
29        let prefix = format!("{reference}.");
30        let mut matches = self
31            .outputs
32            .iter()
33            .filter(|(name, _)| name.starts_with(&prefix));
34        match (matches.next(), matches.next()) {
35            (Some((_, key)), None) => Some(key),
36            _ => None,
37        }
38    }
39}
40
41/// Runs a plan. Steps whose outputs are already in the store are skipped.
42pub fn run(
43    plan: &Plan,
44    inputs: &BTreeMap<String, Artifact>,
45    store: &dyn ArtifactStore,
46) -> Result<RunResult, PipelineError> {
47    if let Some(extra) = inputs.keys().find(|k| !plan.inputs.contains(k)) {
48        return Err(PipelineError::UnexpectedInput(extra.clone()));
49    }
50    let mut keys = BTreeMap::new();
51    for name in &plan.inputs {
52        let artifact = inputs
53            .get(name)
54            .ok_or_else(|| PipelineError::MissingInput(name.clone()))?;
55        let key = match (plan.input_kind(name), artifact) {
56            (Some(ArtifactKind::Series), Artifact::Series(s)) => ArtifactKey::of_series(s),
57            (Some(ArtifactKind::Events), Artifact::Events(e)) => ArtifactKey::of_events(e),
58            (Some(ArtifactKind::Channels), Artifact::Channels(t)) => ArtifactKey::of_channels(t),
59            (expected, found) => {
60                return Err(PipelineError::InputKind {
61                    name: name.clone(),
62                    expected: expected.unwrap_or(ArtifactKind::Series),
63                    found: found.kind(),
64                });
65            }
66        };
67        if !store.contains(&key)? {
68            store.put(&key, Arc::new(artifact.clone()))?;
69        }
70        keys.insert(name.clone(), key);
71    }
72
73    let mut manifest = RunManifest {
74        pipeline: plan.name.clone(),
75        modewise_version: env!("CARGO_PKG_VERSION").to_string(),
76        inputs: keys.clone(),
77        steps: Vec::with_capacity(plan.steps.len()),
78    };
79    // Every key follows from the recipe alone, so they are all known before anything runs.
80    let mut outputs = BTreeMap::new();
81    let mut step_keys = Vec::with_capacity(plan.steps.len());
82    for step in &plan.steps {
83        let k = keys_of(step, &keys);
84        for (name, key) in &k.outputs {
85            let full = format!("{}.{name}", step.id);
86            keys.insert(full.clone(), key.clone());
87            outputs.insert(full, key.clone());
88        }
89        step_keys.push(k);
90    }
91
92    // A stored step runs when its outputs are missing; an unstored one only when a step that
93    // runs needs it. Later steps decide first, so walk backwards.
94    let consumers = |i: usize| {
95        let id = &plan.steps[i].id;
96        (i + 1..plan.steps.len()).filter(move |&j| {
97            plan.steps[j]
98                .inputs
99                .values()
100                .any(|s| matches!(s, Source::Step { step, .. } if step == id))
101        })
102    };
103    let mut needed = vec![false; plan.steps.len()];
104    for i in (0..plan.steps.len()).rev() {
105        needed[i] = if plan.steps[i].keep {
106            let mut missing = false;
107            for key in step_keys[i].outputs.values() {
108                missing |= !store.contains(key)?;
109            }
110            missing
111        } else {
112            consumers(i).any(|j| needed[j])
113        };
114    }
115
116    // Unstored outputs live here until their last consumer has run.
117    let mut scratch: HashMap<ArtifactKey, Arc<Artifact>> = HashMap::new();
118    let mut uses: HashMap<ArtifactKey, usize> = HashMap::new();
119    for (_, k) in step_keys.iter().enumerate().filter(|(i, _)| needed[*i]) {
120        for key in k.inputs.values() {
121            *uses.entry(key.clone()).or_default() += 1;
122        }
123    }
124    for (i, step) in plan.steps.iter().enumerate() {
125        let k = &step_keys[i];
126        let started = Instant::now();
127        if needed[i] {
128            execute(step, k, store, &mut scratch)?;
129            for key in k.inputs.values() {
130                if let Some(n) = uses.get_mut(key) {
131                    *n -= 1;
132                    if *n == 0 {
133                        scratch.remove(key);
134                    }
135                }
136            }
137        }
138        let descriptor = step.op.descriptor();
139        manifest.steps.push(StepRecord {
140            id: step.id.clone(),
141            op: descriptor.name.to_string(),
142            op_version: descriptor.version.to_string(),
143            params: step.params.clone(),
144            inputs: k.inputs.clone(),
145            outputs: k.outputs.clone(),
146            cached: !needed[i],
147            stored: step.keep,
148            seconds: started.elapsed().as_secs_f64(),
149        });
150    }
151    Ok(RunResult { manifest, outputs })
152}
153
154struct StepKeys {
155    inputs: BTreeMap<String, ArtifactKey>,
156    outputs: BTreeMap<String, ArtifactKey>,
157}
158
159/// A step's input and output keys: a hash of the operation, its version, its parameters and its
160/// input keys, so equal recipes share outputs.
161fn keys_of(step: &PlannedStep, keys: &BTreeMap<String, ArtifactKey>) -> StepKeys {
162    let descriptor = step.op.descriptor();
163    let inputs: BTreeMap<String, ArtifactKey> = step
164        .inputs
165        .iter()
166        .map(|(port, source)| (port.clone(), keys[&source.name()].clone()))
167        .collect();
168    let identity = json!({
169        "op": descriptor.name,
170        "version": descriptor.version,
171        "params": step.params,
172        "inputs": inputs,
173    });
174    let step_key = ArtifactKey::from_bytes(identity.to_string().as_bytes());
175    let outputs = descriptor
176        .outputs
177        .iter()
178        .map(|p| {
179            let key = ArtifactKey::from_bytes(format!("{step_key}/{}", p.name).as_bytes());
180            (p.name.to_string(), key)
181        })
182        .collect();
183    StepKeys { inputs, outputs }
184}
185
186fn execute(
187    step: &PlannedStep,
188    keys: &StepKeys,
189    store: &dyn ArtifactStore,
190    scratch: &mut HashMap<ArtifactKey, Arc<Artifact>>,
191) -> Result<(), PipelineError> {
192    let descriptor = step.op.descriptor();
193    let failed = |message: String| PipelineError::StepFailed {
194        step: step.id.clone(),
195        message,
196    };
197    let mut inputs = Inputs::new();
198    for (port, key) in &keys.inputs {
199        let artifact = match scratch.get(key) {
200            Some(a) => Arc::clone(a),
201            None => store
202                .get(key)?
203                .ok_or_else(|| failed(format!("input {port:?} is missing from the store")))?,
204        };
205        inputs.insert(port.clone(), artifact);
206    }
207    let mut produced = step.op.run(&step.params, &inputs).map_err(failed)?;
208    for port in descriptor.outputs {
209        let artifact = produced
210            .remove(port.name)
211            .ok_or_else(|| failed(format!("operation did not produce {:?}", port.name)))?;
212        let expected = step.output_kinds[port.name];
213        if artifact.kind() != expected {
214            return Err(failed(format!(
215                "output {:?} should be {}, got {}",
216                port.name,
217                expected,
218                artifact.kind()
219            )));
220        }
221        let key = &keys.outputs[port.name];
222        if step.keep {
223            store.put(key, Arc::new(artifact))?;
224        } else {
225            scratch.insert(key.clone(), Arc::new(artifact));
226        }
227    }
228    if let Some(extra) = produced.keys().next() {
229        return Err(failed(format!(
230            "operation produced undeclared output {extra:?}"
231        )));
232    }
233    Ok(())
234}
235
236#[cfg(test)]
237mod tests {
238    use std::sync::atomic::Ordering;
239
240    use super::*;
241    use crate::PipelineSpec;
242    use crate::test_support::{TestStore, ops, series};
243
244    const CHAIN: &str = r#"
245        [pipeline]
246        name = "chain"
247        inputs = ["x"]
248
249        [[step]]
250        id = "first"
251        op = "scale"
252        input = "x"
253        params = { factor = 2.0 }
254
255        [[step]]
256        id = "second"
257        op = "scale"
258        input = "first"
259        params = { factor = FACTOR }
260    "#;
261
262    fn chain(factor: f64) -> PipelineSpec {
263        PipelineSpec::from_toml(&CHAIN.replace("FACTOR", &format!("{factor:?}"))).unwrap()
264    }
265
266    fn inputs(values: &[f64]) -> BTreeMap<String, Artifact> {
267        BTreeMap::from([("x".to_string(), Artifact::Series(series(values)))])
268    }
269
270    fn cached(result: &RunResult) -> Vec<bool> {
271        result.manifest.steps.iter().map(|s| s.cached).collect()
272    }
273
274    fn unkept(second_factor: f64) -> PipelineSpec {
275        let text = CHAIN
276            .replace("FACTOR", &format!("{second_factor:?}"))
277            .replacen(
278                "params = { factor = 2.0 }",
279                "params = { factor = 2.0 }\n        keep = false",
280                1,
281            );
282        PipelineSpec::from_toml(&text).unwrap()
283    }
284
285    #[test]
286    fn unkept_outputs_stay_in_memory_and_are_recomputed_only_when_needed() {
287        let ops = ops();
288        let store = TestStore::default();
289        let plan = Plan::new(&unkept(3.0), &ops.registry).unwrap();
290        let first = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
291        assert_eq!(cached(&first), [false, false]);
292        assert!(
293            !store.contains(first.output_key("first").unwrap()).unwrap(),
294            "never written"
295        );
296        let Some(Artifact::Series(out)) = store
297            .get(first.output_key("second").unwrap())
298            .unwrap()
299            .as_deref()
300            .cloned()
301        else {
302            panic!("missing output");
303        };
304        assert_eq!(
305            out.data().row(0).to_vec(),
306            [6.0, 12.0],
307            "computed from the in-memory output"
308        );
309        assert_eq!(
310            first
311                .manifest
312                .steps
313                .iter()
314                .map(|s| s.stored)
315                .collect::<Vec<_>>(),
316            [false, true]
317        );
318
319        let again = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
320        assert_eq!(
321            cached(&again),
322            [true, true],
323            "nothing downstream is missing, so nothing runs"
324        );
325        assert_eq!(ops.scale.runs.load(Ordering::SeqCst), 2);
326
327        let changed = run(
328            &Plan::new(&unkept(4.0), &ops.registry).unwrap(),
329            &inputs(&[1.0, 2.0]),
330            &store,
331        )
332        .unwrap();
333        assert_eq!(
334            cached(&changed),
335            [false, false],
336            "the unkept step reruns because the next one must"
337        );
338        assert_eq!(ops.scale.runs.load(Ordering::SeqCst), 4);
339    }
340
341    #[test]
342    fn an_unkept_step_nothing_uses_is_refused() {
343        let ops = ops();
344        let text = CHAIN.replace("FACTOR", "3.0").replace(
345            "params = { factor = 3.0 }",
346            "params = { factor = 3.0 }\n        keep = false",
347        );
348        let err = Plan::new(&PipelineSpec::from_toml(&text).unwrap(), &ops.registry)
349            .err()
350            .unwrap();
351        assert!(
352            matches!(err, PipelineError::UnusedTransient(ref id) if id == "second"),
353            "{err}"
354        );
355    }
356
357    #[test]
358    fn an_unkept_output_serves_every_step_that_uses_it() {
359        let ops = ops();
360        let store = TestStore::default();
361        let text = r#"
362            [pipeline]
363            name = "fork"
364            inputs = ["x"]
365
366            [[step]]
367            id = "a"
368            op = "scale"
369            input = "x"
370            params = { factor = 2.0 }
371            keep = false
372
373            [[step]]
374            id = "b"
375            op = "scale"
376            input = "a"
377            params = { factor = 3.0 }
378
379            [[step]]
380            id = "c"
381            op = "scale"
382            input = "a"
383            params = { factor = 5.0 }
384        "#;
385        let plan = Plan::new(&PipelineSpec::from_toml(text).unwrap(), &ops.registry).unwrap();
386        let result = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
387        for (id, want) in [("b", [6.0, 12.0]), ("c", [10.0, 20.0])] {
388            let Some(Artifact::Series(out)) = store
389                .get(result.output_key(id).unwrap())
390                .unwrap()
391                .as_deref()
392                .cloned()
393            else {
394                panic!("missing {id}");
395            };
396            assert_eq!(out.data().row(0).to_vec(), want);
397        }
398    }
399
400    #[test]
401    fn chains_of_unkept_steps_are_computed_in_memory() {
402        let ops = ops();
403        let store = TestStore::default();
404        let text = r#"
405            [pipeline]
406            name = "three"
407            inputs = ["x"]
408
409            [[step]]
410            id = "a"
411            op = "scale"
412            input = "x"
413            params = { factor = 2.0 }
414            keep = false
415
416            [[step]]
417            id = "b"
418            op = "scale"
419            input = "a"
420            params = { factor = 3.0 }
421            keep = false
422
423            [[step]]
424            id = "c"
425            op = "scale"
426            input = "b"
427            params = { factor = 5.0 }
428        "#;
429        let plan = Plan::new(&PipelineSpec::from_toml(text).unwrap(), &ops.registry).unwrap();
430        let result = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
431        let Some(Artifact::Series(out)) = store
432            .get(result.output_key("c").unwrap())
433            .unwrap()
434            .as_deref()
435            .cloned()
436        else {
437            panic!("missing output");
438        };
439        assert_eq!(out.data().row(0).to_vec(), [30.0, 60.0]);
440        for id in ["a", "b"] {
441            assert!(!store.contains(result.output_key(id).unwrap()).unwrap());
442        }
443    }
444
445    #[test]
446    fn computes_then_reuses_everything() {
447        let ops = ops();
448        let store = TestStore::default();
449        let plan = Plan::new(&chain(3.0), &ops.registry).unwrap();
450
451        let first = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
452        assert_eq!(cached(&first), [false, false]);
453        let key = first.output_key("second").unwrap();
454        let Some(Artifact::Series(out)) = store.get(key).unwrap().as_deref().cloned() else {
455            panic!("missing output");
456        };
457        assert_eq!(out.data().row(0).to_vec(), [6.0, 12.0]);
458
459        let second = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
460        assert_eq!(cached(&second), [true, true]);
461        assert_eq!(ops.scale.runs.load(Ordering::SeqCst), 2);
462        assert_eq!(second.output_key("second"), first.output_key("second"));
463    }
464
465    #[test]
466    fn param_change_reruns_only_that_step_and_later() {
467        let ops = ops();
468        let store = TestStore::default();
469        run(
470            &Plan::new(&chain(3.0), &ops.registry).unwrap(),
471            &inputs(&[1.0, 2.0]),
472            &store,
473        )
474        .unwrap();
475
476        let changed = run(
477            &Plan::new(&chain(4.0), &ops.registry).unwrap(),
478            &inputs(&[1.0, 2.0]),
479            &store,
480        )
481        .unwrap();
482        assert_eq!(cached(&changed), [true, false]);
483    }
484
485    #[test]
486    fn new_input_data_reruns_everything() {
487        let ops = ops();
488        let store = TestStore::default();
489        let plan = Plan::new(&chain(3.0), &ops.registry).unwrap();
490        run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
491        let changed = run(&plan, &inputs(&[1.0, 2.5]), &store).unwrap();
492        assert_eq!(cached(&changed), [false, false]);
493    }
494
495    #[test]
496    fn rejects_missing_and_unexpected_inputs() {
497        let ops = ops();
498        let store = TestStore::default();
499        let plan = Plan::new(&chain(3.0), &ops.registry).unwrap();
500        let missing = run(&plan, &BTreeMap::new(), &store);
501        assert!(matches!(missing, Err(PipelineError::MissingInput(n)) if n == "x"));
502
503        let mut extra = inputs(&[1.0, 2.0]);
504        extra.insert("y".into(), Artifact::Series(series(&[1.0, 2.0])));
505        assert!(
506            matches!(run(&plan, &extra, &store), Err(PipelineError::UnexpectedInput(n)) if n == "y")
507        );
508    }
509
510    fn single_step(op: &str, params: &str) -> PipelineSpec {
511        PipelineSpec::from_toml(&format!(
512            "[pipeline]\nname = \"t\"\ninputs = [\"x\"]\n[[step]]\nid = \"s\"\nop = \"{op}\"\ninput = \"x\"\nparams = {params}\n"
513        ))
514        .unwrap()
515    }
516
517    #[test]
518    fn rejects_operations_that_break_their_contract() {
519        let ops = ops();
520        for (lie, expected) in [
521            ("missing", "did not produce"),
522            ("wrong_kind", "should be series"),
523            ("extra", "undeclared output"),
524            ("fail", "asked to fail"),
525        ] {
526            let plan = Plan::new(
527                &single_step("liar", &format!("{{ lie = \"{lie}\" }}")),
528                &ops.registry,
529            )
530            .unwrap();
531            let err = run(&plan, &inputs(&[1.0, 2.0]), &TestStore::default())
532                .err()
533                .unwrap();
534            assert!(
535                matches!(&err, PipelineError::StepFailed { step, message } if step == "s" && message.contains(expected)),
536                "{lie}: {err}"
537            );
538        }
539    }
540
541    #[test]
542    fn missing_input_in_store_fails_the_step() {
543        let ops = ops();
544        let plan = Plan::new(&chain(3.0), &ops.registry).unwrap();
545        let err = run(
546            &plan,
547            &inputs(&[1.0, 2.0]),
548            &crate::test_support::ForgetfulStore,
549        )
550        .err()
551        .unwrap();
552        assert!(
553            matches!(&err, PipelineError::StepFailed { message, .. } if message.contains("missing from the store")),
554            "{err}"
555        );
556    }
557
558    #[test]
559    fn output_keys_resolve_by_full_or_short_name() {
560        let ops = ops();
561        let spec = PipelineSpec::from_toml(
562            "[pipeline]\nname = \"t\"\ninputs = [\"x\"]\n[[step]]\nid = \"split\"\nop = \"split\"\ninput = \"x\"\n[[step]]\nid = \"one\"\nop = \"scale\"\ninput = \"x\"\n",
563        )
564        .unwrap();
565        let result = run(
566            &Plan::new(&spec, &ops.registry).unwrap(),
567            &inputs(&[1.0, 2.0]),
568            &TestStore::default(),
569        )
570        .unwrap();
571        assert_eq!(result.outputs().len(), 3);
572        assert!(result.output_key("split.low").is_some());
573        assert!(
574            result.output_key("split").is_none(),
575            "two outputs, so the short name is ambiguous"
576        );
577        assert_eq!(result.output_key("one"), result.output_key("one.signal"));
578        assert!(result.output_key("nope").is_none());
579    }
580
581    #[test]
582    fn manifest_records_resolved_params() {
583        let ops = ops();
584        let store = TestStore::default();
585        let plan = Plan::new(&chain(3.0), &ops.registry).unwrap();
586        let result = run(&plan, &inputs(&[1.0, 2.0]), &store).unwrap();
587        let outputs = result.manifest.outputs();
588        assert_eq!(
589            outputs.keys().collect::<Vec<_>>(),
590            ["first.signal", "second.signal"]
591        );
592        assert_eq!(&outputs, result.outputs());
593        let json = result.manifest.to_json();
594        assert!(json.contains("\"factor\": 3.0"), "{json}");
595        assert!(json.contains("\"pipeline\": \"chain\""));
596    }
597
598    mod properties {
599        use proptest::prelude::*;
600
601        use super::*;
602
603        proptest! {
604            #[test]
605            fn keys_are_stable_and_sensitive_to_every_sample(
606                values in prop::collection::vec(-1e6f64..1e6, 2..32),
607                at in any::<prop::sample::Index>(),
608            ) {
609                let ops = ops();
610                let plan = Plan::new(&chain(3.0), &ops.registry).unwrap();
611                let first = run(&plan, &inputs(&values), &TestStore::default()).unwrap();
612                let again = run(&plan, &inputs(&values), &TestStore::default()).unwrap();
613                prop_assert_eq!(first.outputs(), again.outputs());
614
615                let mut changed = values.clone();
616                let i = at.index(changed.len());
617                changed[i] += 1.0;
618                let other = run(&plan, &inputs(&changed), &TestStore::default()).unwrap();
619                prop_assert_ne!(first.output_key("second"), other.output_key("second"));
620            }
621        }
622    }
623}
modewise_pipeline/executor.rs

modewise_pipeline/
executor.rs
