Skip to main content

somatize_runtime/runner/
mod.rs

1//! Runner module — trait-based execution contracts.
2//!
3//! A [`Runner`] defines the contract for executing plans (fit + forward).
4//! [`LocalRunner`] executes locally using the Executor.
5//! The worker's `RemoteRunner` prepares the environment and delegates to `LocalRunner`.
6
7pub mod local;
8pub mod remote;
9
10use somatize_compiler::ExecutionPlan;
11use somatize_core::cache::CacheStore;
12use somatize_core::error::Result;
13use somatize_core::value::Value;
14use std::collections::HashMap;
15
16use crate::EventBus;
17use crate::executor::GraphInfo;
18use crate::node_catalog::NodeCatalog;
19use std::sync::Arc;
20
21/// Everything a runner needs besides the plan and the data.
22///
23/// A struct rather than six more parameters, and one of them is the point:
24/// `graph_info`. Both runner methods used to build
25/// `GraphInfo::for_linear(plan.node_ids())` — chaining the plan's nodes in
26/// flattened order as if every graph were a chain. On a diamond that is
27/// simply wrong: `GraphSession::forward` on `a → {b, c} → d` answered with
28/// `d(c(...))`, `d` never seeing `b` and `a` never seeing the input.
29///
30/// The caller supplies the real topology now. A caller that genuinely has
31/// only a plan can still pass `GraphInfo::for_linear`, but it has to say so.
32pub struct RunContext<'a> {
33    /// Implementations and trained states for every node in the plan.
34    pub catalog: &'a NodeCatalog,
35    /// Output cache consulted and filled by `run_node`.
36    pub cache: &'a dyn CacheStore,
37    /// Bus the run emits its node events on.
38    pub events: &'a Arc<EventBus>,
39    /// Tags every node event of this run — callers that emit a
40    /// `RunStarted`/`RunCompleted` bracket pass the same id so readers can
41    /// group a run's events.
42    pub run_id: &'a str,
43    /// The real topology for input resolution — the reason this struct
44    /// exists; see the type docs.
45    pub graph_info: GraphInfo,
46    /// The run's experiment seed, folded into every cache key.
47    ///
48    /// Without it two seeds share a state cache line, so the second one
49    /// trains on the first one's recorded state and the sweep measures
50    /// one seed N times. Only the Python fit path used to salt.
51    pub seed: Option<i64>,
52    /// Performs and journals step effects.
53    ///
54    /// Needed only when the plan contains a step. It lives here rather than
55    /// being built inside the runner because a driver carries the journal —
56    /// which is what makes a resumed run replay instead of re-calling a
57    /// model — and only the caller knows where that journal lives.
58    pub driver: Option<crate::effects::EffectDriver>,
59}
60
61impl<'a> RunContext<'a> {
62    /// A context over the real topology; use [`Self::linear`] only when a
63    /// plan is genuinely all you have.
64    pub fn new(
65        catalog: &'a NodeCatalog,
66        cache: &'a dyn CacheStore,
67        events: &'a Arc<EventBus>,
68        run_id: &'a str,
69        graph_info: GraphInfo,
70    ) -> Self {
71        Self {
72            catalog,
73            cache,
74            events,
75            run_id,
76            graph_info,
77            seed: None,
78            driver: None,
79        }
80    }
81
82    /// Fold this run's seed into the cache keys.
83    pub fn with_seed(mut self, seed: Option<i64>) -> Self {
84        self.seed = seed;
85        self
86    }
87
88    /// Register the effect driver a plan containing steps needs.
89    ///
90    /// The driver should already carry its catalog
91    /// ([`crate::effects::EffectDriver::with_catalog`]) if a step may fan
92    /// out dynamically — the same rule as
93    /// [`crate::executor::Context::with_driver`], so the two entry points
94    /// cannot drift apart on who attaches it.
95    pub fn with_driver(mut self, driver: crate::effects::EffectDriver) -> Self {
96        self.driver = Some(driver);
97        self
98    }
99
100    /// For a caller that has only a plan: treat it as a chain.
101    ///
102    /// Correct for a linear pipeline and a fabrication for anything else,
103    /// which is why it is spelled out at the call site rather than being
104    /// what you get by default.
105    pub fn linear(
106        catalog: &'a NodeCatalog,
107        cache: &'a dyn CacheStore,
108        events: &'a Arc<EventBus>,
109        run_id: &'a str,
110        plan: &ExecutionPlan,
111    ) -> Self {
112        let ids = plan.node_ids();
113        Self::new(catalog, cache, events, run_id, GraphInfo::for_linear(&ids))
114    }
115
116    /// Clone the driver for a run's own context.
117    pub(crate) fn driver(&self) -> Option<crate::effects::EffectDriver> {
118        self.driver.clone()
119    }
120}
121
122/// Contract for executing plans. Every execution mode (local, remote, stream)
123/// implements this trait. One interface, polymorphic dispatch.
124pub trait Runner: Send + Sync {
125    /// Train: fit each filter, forward to propagate outputs.
126    /// Returns (last output, all node outputs).
127    fn fit(
128        &self,
129        plan: &ExecutionPlan,
130        ctx: &RunContext<'_>,
131        input: &Value,
132        y: Option<&Value>,
133    ) -> Result<(Value, HashMap<String, Value>)>;
134
135    /// Inference: forward data through the compiled plan.
136    fn forward(&self, plan: &ExecutionPlan, ctx: &RunContext<'_>, input: &Value) -> Result<Value>;
137}
138
139pub use local::LocalRunner;
140pub use remote::{RemoteRunner, Transport};