somatize_runtime/runner/
remote.rs

1//! RemoteRunner — executes plans on remote workers via a Transport abstraction.
2//!
3//! The Transport trait abstracts HOW to communicate with workers (WS, HTTP, gRPC, etc.).
4//! RemoteRunner implements Runner by serializing fit/forward calls and sending them
5//! through the transport layer.
6
7use super::{RunContext, Runner};
8use crate::node_catalog::NodeCatalog;
9
10use crate::executor::RunMode;
11use somatize_compiler::ExecutionPlan;
12use somatize_core::error::Result;
13use somatize_core::value::Value;
14use std::collections::HashMap;
15
16/// Abstraction for communicating with remote workers.
17/// Implemented by WsTransport (WebSocket), but could be HTTP, gRPC, etc.
18pub trait Transport: Send + Sync {
19    /// Send a plan for execution and receive the output + trained states.
20    ///
21    /// `mode` says what to do with the nodes, and carries the labels when
22    /// there are any. It replaced a `fit_mode: bool` sitting beside an
23    /// `y: Option<&Value>` — a flag selecting between two operations with
24    /// differently shaped results, and a parameter that meant nothing
25    /// unless the flag was set. It is the same [`RunMode`] the local
26    /// executor reads, so the two paths cannot disagree about what a fit is.
27    ///
28    /// `seed` is the run's experiment seed, and it is a parameter rather
29    /// than something the transport digs out because the transport has no
30    /// [`RunContext`] to dig in. Without it the worker salts nothing, and a
31    /// five-seed sweep run remotely shares one cache line across all five —
32    /// the worker protocol's `SerializedPlan::seed` documents that as the
33    /// bug it exists to close, and this path was still passing `None`.
34    fn execute(
35        &self,
36        plan: &ExecutionPlan,
37        filters: &NodeCatalog,
38        input: &Value,
39        mode: &RunMode,
40        seed: Option<i64>,
41    ) -> Result<(Value, HashMap<String, Value>)>;
42
43    /// Request trained states from the remote worker.
44    fn get_state(&self, node_ids: &[String]) -> Result<HashMap<String, Value>>;
45
46    /// Load states on the remote worker.
47    fn set_state(&self, states: &HashMap<String, Value>) -> Result<()>;
48
49    /// Request gradients from the remote worker.
50    fn get_gradients(&self, node_ids: &[String]) -> Result<HashMap<String, Value>>;
51
52    /// Apply aggregated gradients on the remote worker.
53    fn apply_gradients(&self, gradients: &HashMap<String, Value>) -> Result<()>;
54
55    /// Convenience: execute a single node remotely (used by the plan executor).
56    ///
57    /// Unseeded, and it has to be: this takes a node id and nothing else,
58    /// so there is no run to take a seed from. Callers that have a
59    /// [`RunContext`] should go through [`Transport::execute`] with
60    /// `ctx.seed` instead of reaching for this.
61    fn execute_node(&self, node_id: &str, input: Option<&Value>) -> Result<Value> {
62        let plan = ExecutionPlan::Execute {
63            node_id: node_id.to_string(),
64        };
65        let input_val = input.cloned().unwrap_or(Value::Empty);
66        let filters = crate::node_catalog::NodeCatalog::new();
67        let (output, _) = self.execute(&plan, &filters, &input_val, &RunMode::Forward, None)?;
68        Ok(output)
69    }
70}
71
72/// A Runner that delegates execution to a remote worker via Transport.
73pub struct RemoteRunner {
74    transport: Box<dyn Transport>,
75}
76
77impl RemoteRunner {
78    /// A runner sending every fit/forward through `transport`.
79    pub fn new(transport: impl Transport + 'static) -> Self {
80        Self {
81            transport: Box::new(transport),
82        }
83    }
84
85    /// Access the underlying transport (for strategy methods).
86    pub fn transport(&self) -> &dyn Transport {
87        self.transport.as_ref()
88    }
89}
90
91impl Runner for RemoteRunner {
92    fn fit(
93        &self,
94        plan: &ExecutionPlan,
95        ctx: &RunContext<'_>,
96        input: &Value,
97        y: Option<&Value>,
98    ) -> Result<(Value, HashMap<String, Value>)> {
99        self.transport.execute(
100            plan,
101            ctx.catalog,
102            input,
103            &RunMode::Fit { y: y.cloned() },
104            ctx.seed,
105        )
106    }
107
108    fn forward(&self, plan: &ExecutionPlan, ctx: &RunContext<'_>, input: &Value) -> Result<Value> {
109        let (output, _states) =
110            self.transport
111                .execute(plan, ctx.catalog, input, &RunMode::Forward, ctx.seed)?;
112        Ok(output)
113    }
114}
somatize_runtime/runner/remote.rs

somatize_runtime/runner/
remote.rs
