use crate::cache::MemoryCache;
use crate::event_bus::EventBus;
use crate::executor::{self, Context, GraphInfo};
use crate::node_catalog::NodeCatalog;
use crate::runner::Runner;
use crate::runner::Transport;
use somatize_compiler::{CompileMode, CompileResult, compile};
use somatize_core::cache::{CacheKey, CacheStore};
use somatize_core::error::{Result, SomaError};
use somatize_core::event::Event;
use somatize_core::fingerprint::ArchitectureFingerprint;
use somatize_core::graph::Graph;
use somatize_core::store::{DataRef, DataStore};
use somatize_core::util::timestamp_id;
use somatize_core::value::Value;
use std::collections::HashMap;
use std::sync::Arc;
pub struct GraphSession {
graph: Graph,
catalog: NodeCatalog,
cache: Arc<dyn CacheStore>,
event_bus: Arc<EventBus>,
data_store: Option<Arc<dyn DataStore>>,
transport: Option<Arc<dyn Transport>>,
driver: Option<crate::effects::EffectDriver>,
fitted: bool,
}
impl GraphSession {
pub fn new(graph: Graph, catalog: NodeCatalog) -> Self {
Self {
graph,
catalog,
cache: Arc::new(MemoryCache::default()),
event_bus: Arc::new(EventBus::new(256)),
data_store: None,
transport: None,
driver: None,
fitted: false,
}
}
pub fn with_cache(mut self, cache: Arc<dyn CacheStore>) -> Self {
self.cache = cache;
self
}
pub fn with_event_bus(mut self, bus: Arc<EventBus>) -> Self {
self.event_bus = bus;
self
}
pub fn with_data_store(mut self, store: Arc<dyn DataStore>) -> Self {
self.data_store = Some(store);
self
}
pub fn with_transport(mut self, transport: Arc<dyn Transport>) -> Self {
self.transport = Some(transport);
self
}
pub fn with_driver(mut self, driver: crate::effects::EffectDriver) -> Self {
self.driver = Some(driver);
self
}
fn run_driver(&self) -> Option<crate::effects::EffectDriver> {
self.driver
.as_ref()
.map(|d| d.clone().with_catalog(Arc::new(self.catalog.clone())))
}
pub fn compile(&self, mode: CompileMode) -> Result<CompileResult> {
compile(&self.graph, &self.catalog, mode, Some(self.cache.as_ref()))
}
pub fn run(&mut self, mode: CompileMode) -> Result<HashMap<String, Value>> {
let CompileResult { plan, diagnostics } =
compile(&self.graph, &self.catalog, mode, Some(self.cache.as_ref()))?;
for diag in &diagnostics {
tracing::warn!("compile diagnostic: {:?}", diag);
}
let graph_info = GraphInfo::from_graph(&self.graph);
let run_id = timestamp_id("graph_run");
let mut ctx =
Context::new(self.event_bus.clone(), run_id.clone()).with_graph_info(graph_info);
if let Some(store) = &self.data_store {
ctx = ctx.with_data_store(store.clone());
}
if let Some(transport) = &self.transport {
ctx = ctx.with_transport(transport.clone());
}
if let Some(driver) = self.run_driver() {
ctx = ctx.with_driver(driver);
}
self.event_bus.emit(Event::RunStarted {
run_id: run_id.clone(),
plan_summary: plan.summary(),
});
let start = std::time::Instant::now();
if let Err(e) = executor::execute(&plan, &mut ctx, &self.catalog, self.cache.as_ref()) {
self.event_bus.emit(Event::RunFailed {
run_id,
error: e.to_string(),
});
return Err(e);
}
self.event_bus.emit(Event::RunCompleted {
run_id,
duration: start.elapsed(),
});
Ok(ctx.into_outputs())
}
pub fn fit(&mut self, x: &Value, y: Option<&Value>) -> Result<HashMap<String, Value>> {
self.graph.validate()?;
let CompileResult { plan, .. } = compile(
&self.graph,
&self.catalog,
CompileMode::NoCache,
Some(self.cache.as_ref()),
)?;
let run_id = timestamp_id("fit");
self.event_bus.emit(Event::RunStarted {
run_id: run_id.clone(),
plan_summary: plan.summary(),
});
let start = std::time::Instant::now();
let runner = crate::runner::LocalRunner;
let mut ctx = crate::runner::RunContext::new(
&self.catalog,
self.cache.as_ref(),
&self.event_bus,
&run_id,
GraphInfo::from_graph(&self.graph),
);
if let Some(driver) = self.run_driver() {
ctx = ctx.with_driver(driver);
}
let result = runner.fit(&plan, &ctx, x, y);
let (_last_output, mut all_outputs) = match result {
Ok(out) => {
self.event_bus.emit(Event::RunCompleted {
run_id,
duration: start.elapsed(),
});
out
}
Err(e) => {
self.event_bus.emit(Event::RunFailed {
run_id,
error: e.to_string(),
});
return Err(e);
}
};
for (key, value) in &all_outputs {
if let Some(node_id) = somatize_core::keys::node_of_state_key(key) {
self.catalog.try_set_state(node_id, value.clone())?;
}
}
all_outputs.retain(|k, _| somatize_core::keys::node_of_state_key(k).is_none());
self.fitted = true;
Ok(all_outputs)
}
pub fn forward_with(
&self,
x: &Value,
strategy: &dyn crate::forward::ForwardStrategy,
) -> Result<Value> {
let driver = self.run_driver();
strategy.forward(
&self.graph,
&crate::forward::ForwardEnv {
catalog: &self.catalog,
cache: self.cache.as_ref(),
event_bus: &self.event_bus,
data_store: self.data_store.as_ref(),
driver: driver.as_ref(),
},
x,
)
}
pub fn forward(&self, x: &Value) -> Result<Value> {
self.forward_with(x, &crate::forward::Standard)
}
pub fn persist_states(&self) -> Result<DataRef> {
let store = self
.data_store
.as_ref()
.ok_or_else(|| SomaError::Execution {
node_id: "session".into(),
message: "persist_states requires a data store".into(),
})?;
let sorted = self.graph.topological_sort()?;
let mut states_map = serde_json::Map::new();
for node_id in &sorted {
if let Some(state) = self.catalog.get_state(node_id) {
let json = serde_json::to_value(&*state)
.map_err(|e| SomaError::Other(format!("state serialize: {e}")))?;
states_map.insert(node_id.to_string(), json);
}
}
let states_value = Value::json(serde_json::Value::Object(states_map));
let fingerprint = self.graph_config_hash()?;
let key = CacheKey::from_parts(&[b"graph_states", fingerprint.as_bytes()]);
store.put(&key, &states_value)
}
pub fn load_states(&mut self, data_ref: &DataRef) -> Result<()> {
let store = self
.data_store
.as_ref()
.ok_or_else(|| SomaError::Execution {
node_id: "session".into(),
message: "load_states requires a data store".into(),
})?;
let states_value = store.get(data_ref)?;
let states_json = states_value
.as_json()
.ok_or_else(|| SomaError::Other("persisted states must be JSON".into()))?;
let obj = states_json
.as_object()
.ok_or_else(|| SomaError::Other("persisted states must be a JSON object".into()))?;
for (node_id, json_val) in obj {
let value: Value = serde_json::from_value(json_val.clone())
.map_err(|e| SomaError::Other(format!("state deserialize: {e}")))?;
self.catalog.try_set_state(node_id.clone(), value)?;
}
self.fitted = true;
Ok(())
}
pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<Event> {
self.event_bus.subscribe()
}
pub fn event_bus(&self) -> &Arc<EventBus> {
&self.event_bus
}
pub fn is_fitted(&self) -> bool {
self.fitted
}
pub fn graph(&self) -> &Graph {
&self.graph
}
pub fn catalog(&self) -> &NodeCatalog {
&self.catalog
}
pub fn catalog_mut(&mut self) -> &mut NodeCatalog {
&mut self.catalog
}
fn graph_config_hash(&self) -> Result<String> {
Ok(ArchitectureFingerprint::of(&self.graph)?.digest)
}
}
pub fn graph_run(
graph: &Graph,
catalog: &NodeCatalog,
mode: CompileMode,
cache: Arc<dyn CacheStore>,
) -> Result<HashMap<String, Value>> {
GraphSession::new(graph.clone(), catalog.clone())
.with_cache(cache)
.run(mode)
}
pub fn graph_fit(
graph: &Graph,
catalog: &NodeCatalog,
x: &Value,
y: Option<&Value>,
cache: Arc<dyn CacheStore>,
) -> Result<HashMap<String, Value>> {
GraphSession::new(graph.clone(), catalog.clone())
.with_cache(cache)
.fit(x, y)
}
pub fn graph_predict(
graph: &Graph,
catalog: &NodeCatalog,
x: &Value,
cache: Arc<dyn CacheStore>,
) -> Result<Value> {
GraphSession::new(graph.clone(), catalog.clone())
.with_cache(cache)
.forward(x)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cache::MemoryCache;
use somatize_compiler::NodeRegistry;
use somatize_core::cache::CacheKey;
use somatize_core::error::Result;
use somatize_core::filter::{FilterKind, FilterMeta, StreamMode};
use somatize_core::graph::{Edge, Node};
struct DoublerFilter;
impl somatize_core::filter::Filter for DoublerFilter {
fn config_hash(&self) -> CacheKey {
CacheKey::from_parts(&[b"Doubler"])
}
fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
Ok(Value::Empty)
}
fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
let (data, shape) = x
.as_tensor()
.ok_or(SomaError::Other("need tensor".into()))?;
Ok(Value::tensor(
data.iter().map(|v| v * 2.0).collect(),
shape.to_vec(),
))
}
fn meta(&self) -> FilterMeta {
FilterMeta {
name: "Doubler".into(),
kind: FilterKind::Stateless,
cacheable: true,
differentiable: true,
deterministic: true,
stream_mode: StreamMode::FixedState,
distribution: somatize_core::filter::Distribution::Local,
input_schema: None,
output_schema: None,
}
}
}
struct AdderFilter(f64);
impl somatize_core::filter::Filter for AdderFilter {
fn config_hash(&self) -> CacheKey {
CacheKey::from_parts(&[b"Adder", &self.0.to_le_bytes()])
}
fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
Ok(Value::Empty)
}
fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
let (data, shape) = x
.as_tensor()
.ok_or(SomaError::Other("need tensor".into()))?;
Ok(Value::tensor(
data.iter().map(|v| v + self.0).collect(),
shape.to_vec(),
))
}
fn meta(&self) -> FilterMeta {
FilterMeta {
name: "Adder".into(),
kind: FilterKind::Stateless,
cacheable: true,
differentiable: true,
deterministic: true,
stream_mode: StreamMode::FixedState,
distribution: somatize_core::filter::Distribution::Local,
input_schema: None,
output_schema: None,
}
}
}
struct MeanFilter;
impl somatize_core::filter::Filter for MeanFilter {
fn config_hash(&self) -> CacheKey {
CacheKey::from_parts(&[b"Mean"])
}
fn fit(&self, x: &Value, _y: Option<&Value>) -> Result<Value> {
let (data, _) = x
.as_tensor()
.ok_or(SomaError::Other("need tensor".into()))?;
let mean = data.iter().sum::<f64>() / data.len() as f64;
Ok(Value::json(serde_json::json!({ "mean": mean })))
}
fn forward(&self, x: &Value, state: &Value) -> Result<Value> {
let (data, shape) = x
.as_tensor()
.ok_or(SomaError::Other("need tensor".into()))?;
let mean = state
.as_json()
.and_then(|j| j["mean"].as_f64())
.unwrap_or(0.0);
Ok(Value::tensor(
data.iter().map(|v| v - mean).collect(),
shape.to_vec(),
))
}
fn meta(&self) -> FilterMeta {
FilterMeta {
name: "Mean".into(),
kind: FilterKind::Trainable,
cacheable: true,
differentiable: true,
deterministic: true,
stream_mode: StreamMode::FixedState,
distribution: somatize_core::filter::Distribution::Local,
input_schema: None,
output_schema: None,
}
}
}
fn linear_graph(ids: &[&str]) -> Graph {
let mut g = Graph::new();
for &id in ids {
g.nodes.push(Node::new(id, id, id));
}
for (i, pair) in ids.windows(2).enumerate() {
g.edges.push(Edge::data(format!("e{i}"), pair[0], pair[1]));
}
g
}
#[test]
fn session_run_linear() {
let graph = linear_graph(&["double", "add"]);
let mut lib = NodeCatalog::new();
lib.register("double", Box::new(DoublerFilter));
lib.register("add", Box::new(AdderFilter(10.0)));
let mut session = GraphSession::new(graph, lib);
let cache = MemoryCache::default();
session = session.with_cache(Arc::new(cache));
let CompileResult { plan, .. } = session.compile(CompileMode::NoCache).unwrap();
let bus = Arc::new(EventBus::new(64));
let mut ctx =
Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(session.graph()));
ctx.set(
somatize_core::keys::GRAPH_INPUT,
Value::tensor(vec![1.0, 2.0, 3.0], vec![3]),
);
executor::execute(&plan, &mut ctx, session.catalog(), &MemoryCache::default()).unwrap();
let outputs: HashMap<String, Value> = ctx.into_outputs();
let result = outputs.get("add").unwrap();
let (data, _) = result.as_tensor().unwrap();
assert_eq!(data, &[12.0, 14.0, 16.0]);
}
#[test]
fn session_fit_and_forward() {
let graph = linear_graph(&["mean", "double"]);
let mut lib = NodeCatalog::new();
lib.register("mean", Box::new(MeanFilter));
lib.register("double", Box::new(DoublerFilter));
let mut session = GraphSession::new(graph, lib);
let x = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
let outputs = session.fit(&x, None).unwrap();
let result = outputs.get("double").unwrap();
let (data, _) = result.as_tensor().unwrap();
assert_eq!(data, &[-20.0, 0.0, 20.0]);
assert!(session.is_fitted());
}
#[test]
fn session_compile_diagnostics() {
let graph = linear_graph(&["double"]);
let mut lib = NodeCatalog::new();
lib.register("double", Box::new(DoublerFilter));
let session = GraphSession::new(graph, lib);
let result = session.compile(CompileMode::NoCache).unwrap();
assert!(result.plan.node_count() > 0);
}
#[test]
fn graph_run_linear() {
let graph = linear_graph(&["double", "add"]);
let mut lib = NodeCatalog::new();
lib.register("double", Box::new(DoublerFilter));
lib.register("add", Box::new(AdderFilter(10.0)));
let cache = MemoryCache::default();
let outputs = {
let CompileResult { plan, .. } =
compile(&graph, &lib, CompileMode::NoCache, None).unwrap();
let bus = Arc::new(EventBus::new(64));
let mut ctx = Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(&graph));
ctx.set(
somatize_core::keys::GRAPH_INPUT,
Value::tensor(vec![1.0, 2.0, 3.0], vec![3]),
);
executor::execute(&plan, &mut ctx, &lib, &cache).unwrap();
ctx.into_outputs()
};
let result = outputs.get("add").unwrap();
let (data, _) = result.as_tensor().unwrap();
assert_eq!(data, &[12.0, 14.0, 16.0]);
}
#[test]
fn graph_run_diamond() {
let mut graph = Graph::new();
graph.nodes.push(Node::new("double", "Double", "double"));
graph.nodes.push(Node::new("add", "Add", "add"));
graph.nodes.push(Node::new("merge", "Merge", "merge"));
graph.edges.push(Edge::data("e1", "double", "merge"));
graph.edges.push(Edge::data("e2", "add", "merge"));
let mut lib = NodeCatalog::new();
lib.register("double", Box::new(DoublerFilter));
lib.register("add", Box::new(AdderFilter(100.0)));
struct MergeFilter;
impl somatize_core::filter::Filter for MergeFilter {
fn config_hash(&self) -> CacheKey {
CacheKey::from_parts(&[b"Merge"])
}
fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
Ok(Value::Empty)
}
fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
Ok(x.clone())
}
fn meta(&self) -> FilterMeta {
FilterMeta {
name: "Merge".into(),
kind: FilterKind::Stateless,
cacheable: true,
differentiable: false,
deterministic: true,
stream_mode: StreamMode::FixedState,
distribution: somatize_core::filter::Distribution::Local,
input_schema: None,
output_schema: None,
}
}
}
lib.register("merge", Box::new(MergeFilter));
let cache = MemoryCache::default();
let CompileResult { plan, .. } = compile(&graph, &lib, CompileMode::NoCache, None).unwrap();
let bus = Arc::new(EventBus::new(64));
let mut ctx = Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(&graph));
ctx.set(
somatize_core::keys::GRAPH_INPUT,
Value::tensor(vec![5.0], vec![1]),
);
executor::execute(&plan, &mut ctx, &lib, &cache).unwrap();
let merge_output = ctx.get("merge").unwrap();
assert!(
merge_output.as_json().is_some(),
"merge should receive JSON from multiple predecessors"
);
}
#[test]
fn graph_fit_trainable() {
let graph = linear_graph(&["mean", "double"]);
let mut lib = NodeCatalog::new();
lib.register("mean", Box::new(MeanFilter));
lib.register("double", Box::new(DoublerFilter));
let cache = Arc::new(MemoryCache::default());
let x = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
let outputs = graph_fit(&graph, &lib, &x, None, cache.clone()).unwrap();
let result = outputs.get("double").unwrap();
let (data, _) = result.as_tensor().unwrap();
assert_eq!(data, &[-20.0, 0.0, 20.0]);
assert!(!cache.is_empty());
}
#[test]
fn the_catalog_is_the_compiler_registry() {
let mut lib = NodeCatalog::new();
lib.register("a", Box::new(DoublerFilter));
let registry: &dyn NodeRegistry = &lib;
assert!(registry.meta("a").is_some());
assert_eq!(registry.meta("a").unwrap().name, "Doubler");
assert!(registry.config_hash("a").is_some());
assert!(registry.meta("b").is_none());
}
fn session_of(graph: Graph) -> GraphSession {
let mut lib = NodeCatalog::new();
for node in &graph.nodes {
lib.register(&node.id, Box::new(DoublerFilter));
}
GraphSession::new(graph, lib)
}
#[test]
fn state_address_separates_graphs_that_share_node_ids() {
let chain = session_of(linear_graph(&["a", "b", "c"]));
let mut fan = Graph::new();
for id in ["a", "b", "c"] {
fan.nodes.push(Node::new(id, id, id));
}
fan.edges.push(Edge::data("e0", "a", "b"));
fan.edges.push(Edge::data("e1", "a", "c"));
let fan = session_of(fan);
assert_ne!(
chain.graph_config_hash().unwrap(),
fan.graph_config_hash().unwrap(),
"two differently wired graphs must not persist states to one address"
);
}
#[test]
fn state_address_is_stable_for_the_same_graph() {
assert_eq!(
session_of(linear_graph(&["a", "b"]))
.graph_config_hash()
.unwrap(),
session_of(linear_graph(&["a", "b"]))
.graph_config_hash()
.unwrap()
);
}
}