1use crate::cache::MemoryCache;
8use crate::event_bus::EventBus;
9use crate::executor::{self, Context, GraphInfo};
10use crate::node_catalog::NodeCatalog;
11use crate::runner::Runner;
12use crate::runner::Transport;
13use somatize_compiler::{CompileMode, CompileResult, compile};
14use somatize_core::cache::{CacheKey, CacheStore};
15use somatize_core::error::{Result, SomaError};
16use somatize_core::event::Event;
17use somatize_core::fingerprint::ArchitectureFingerprint;
18use somatize_core::graph::Graph;
19use somatize_core::store::{DataRef, DataStore};
20use somatize_core::util::timestamp_id;
21use somatize_core::value::Value;
22use std::collections::HashMap;
23use std::sync::Arc;
24
25pub struct GraphSession {
37 graph: Graph,
38 catalog: NodeCatalog,
39 cache: Arc<dyn CacheStore>,
40 event_bus: Arc<EventBus>,
41 data_store: Option<Arc<dyn DataStore>>,
42 transport: Option<Arc<dyn Transport>>,
43 driver: Option<crate::effects::EffectDriver>,
47 fitted: bool,
48}
49
50impl GraphSession {
51 pub fn new(graph: Graph, catalog: NodeCatalog) -> Self {
54 Self {
55 graph,
56 catalog,
57 cache: Arc::new(MemoryCache::default()),
58 event_bus: Arc::new(EventBus::new(256)),
59 data_store: None,
60 transport: None,
61 driver: None,
62 fitted: false,
63 }
64 }
65
66 pub fn with_cache(mut self, cache: Arc<dyn CacheStore>) -> Self {
69 self.cache = cache;
70 self
71 }
72
73 pub fn with_event_bus(mut self, bus: Arc<EventBus>) -> Self {
76 self.event_bus = bus;
77 self
78 }
79
80 pub fn with_data_store(mut self, store: Arc<dyn DataStore>) -> Self {
82 self.data_store = Some(store);
83 self
84 }
85
86 pub fn with_transport(mut self, transport: Arc<dyn Transport>) -> Self {
88 self.transport = Some(transport);
89 self
90 }
91
92 pub fn with_driver(mut self, driver: crate::effects::EffectDriver) -> Self {
99 self.driver = Some(driver);
100 self
101 }
102
103 fn run_driver(&self) -> Option<crate::effects::EffectDriver> {
105 self.driver
106 .as_ref()
107 .map(|d| d.clone().with_catalog(Arc::new(self.catalog.clone())))
108 }
109
110 pub fn compile(&self, mode: CompileMode) -> Result<CompileResult> {
114 compile(&self.graph, &self.catalog, mode, Some(self.cache.as_ref()))
115 }
116
117 pub fn run(&mut self, mode: CompileMode) -> Result<HashMap<String, Value>> {
123 let CompileResult { plan, diagnostics } =
124 compile(&self.graph, &self.catalog, mode, Some(self.cache.as_ref()))?;
125
126 for diag in &diagnostics {
127 tracing::warn!("compile diagnostic: {:?}", diag);
128 }
129
130 let graph_info = GraphInfo::from_graph(&self.graph);
131 let run_id = timestamp_id("graph_run");
132 let mut ctx =
133 Context::new(self.event_bus.clone(), run_id.clone()).with_graph_info(graph_info);
134
135 if let Some(store) = &self.data_store {
136 ctx = ctx.with_data_store(store.clone());
137 }
138 if let Some(transport) = &self.transport {
139 ctx = ctx.with_transport(transport.clone());
140 }
141 if let Some(driver) = self.run_driver() {
142 ctx = ctx.with_driver(driver);
143 }
144
145 self.event_bus.emit(Event::RunStarted {
146 run_id: run_id.clone(),
147 plan_summary: plan.summary(),
148 });
149 let start = std::time::Instant::now();
150 if let Err(e) = executor::execute(&plan, &mut ctx, &self.catalog, self.cache.as_ref()) {
151 self.event_bus.emit(Event::RunFailed {
152 run_id,
153 error: e.to_string(),
154 });
155 return Err(e);
156 }
157 self.event_bus.emit(Event::RunCompleted {
158 run_id,
159 duration: start.elapsed(),
160 });
161
162 Ok(ctx.into_outputs())
163 }
164
165 pub fn fit(&mut self, x: &Value, y: Option<&Value>) -> Result<HashMap<String, Value>> {
171 self.graph.validate()?;
172
173 let CompileResult { plan, .. } = compile(
174 &self.graph,
175 &self.catalog,
176 CompileMode::NoCache,
177 Some(self.cache.as_ref()),
178 )?;
179
180 let run_id = timestamp_id("fit");
181 self.event_bus.emit(Event::RunStarted {
182 run_id: run_id.clone(),
183 plan_summary: plan.summary(),
184 });
185 let start = std::time::Instant::now();
186
187 let runner = crate::runner::LocalRunner;
188 let mut ctx = crate::runner::RunContext::new(
189 &self.catalog,
190 self.cache.as_ref(),
191 &self.event_bus,
192 &run_id,
193 GraphInfo::from_graph(&self.graph),
194 );
195 if let Some(driver) = self.run_driver() {
196 ctx = ctx.with_driver(driver);
197 }
198 let result = runner.fit(&plan, &ctx, x, y);
199 let (_last_output, mut all_outputs) = match result {
200 Ok(out) => {
201 self.event_bus.emit(Event::RunCompleted {
202 run_id,
203 duration: start.elapsed(),
204 });
205 out
206 }
207 Err(e) => {
208 self.event_bus.emit(Event::RunFailed {
209 run_id,
210 error: e.to_string(),
211 });
212 return Err(e);
213 }
214 };
215
216 for (key, value) in &all_outputs {
218 if let Some(node_id) = somatize_core::keys::node_of_state_key(key) {
219 self.catalog.try_set_state(node_id, value.clone())?;
220 }
221 }
222
223 all_outputs.retain(|k, _| somatize_core::keys::node_of_state_key(k).is_none());
225
226 self.fitted = true;
227 Ok(all_outputs)
228 }
229
230 pub fn forward_with(
237 &self,
238 x: &Value,
239 strategy: &dyn crate::forward::ForwardStrategy,
240 ) -> Result<Value> {
241 let driver = self.run_driver();
242 strategy.forward(
243 &self.graph,
244 &crate::forward::ForwardEnv {
245 catalog: &self.catalog,
246 cache: self.cache.as_ref(),
247 event_bus: &self.event_bus,
248 data_store: self.data_store.as_ref(),
249 driver: driver.as_ref(),
250 },
251 x,
252 )
253 }
254
255 pub fn forward(&self, x: &Value) -> Result<Value> {
257 self.forward_with(x, &crate::forward::Standard)
258 }
259
260 pub fn persist_states(&self) -> Result<DataRef> {
264 let store = self
265 .data_store
266 .as_ref()
267 .ok_or_else(|| SomaError::Execution {
268 node_id: "session".into(),
269 message: "persist_states requires a data store".into(),
270 })?;
271
272 let sorted = self.graph.topological_sort()?;
273 let mut states_map = serde_json::Map::new();
274 for node_id in &sorted {
275 if let Some(state) = self.catalog.get_state(node_id) {
276 let json = serde_json::to_value(&*state)
277 .map_err(|e| SomaError::Other(format!("state serialize: {e}")))?;
278 states_map.insert(node_id.to_string(), json);
279 }
280 }
281
282 let states_value = Value::json(serde_json::Value::Object(states_map));
283 let fingerprint = self.graph_config_hash()?;
284 let key = CacheKey::from_parts(&[b"graph_states", fingerprint.as_bytes()]);
285 store.put(&key, &states_value)
286 }
287
288 pub fn load_states(&mut self, data_ref: &DataRef) -> Result<()> {
290 let store = self
291 .data_store
292 .as_ref()
293 .ok_or_else(|| SomaError::Execution {
294 node_id: "session".into(),
295 message: "load_states requires a data store".into(),
296 })?;
297
298 let states_value = store.get(data_ref)?;
299 let states_json = states_value
300 .as_json()
301 .ok_or_else(|| SomaError::Other("persisted states must be JSON".into()))?;
302 let obj = states_json
303 .as_object()
304 .ok_or_else(|| SomaError::Other("persisted states must be a JSON object".into()))?;
305
306 for (node_id, json_val) in obj {
307 let value: Value = serde_json::from_value(json_val.clone())
308 .map_err(|e| SomaError::Other(format!("state deserialize: {e}")))?;
309 self.catalog.try_set_state(node_id.clone(), value)?;
310 }
311
312 self.fitted = true;
313 Ok(())
314 }
315
316 pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<Event> {
320 self.event_bus.subscribe()
321 }
322
323 pub fn event_bus(&self) -> &Arc<EventBus> {
325 &self.event_bus
326 }
327
328 pub fn is_fitted(&self) -> bool {
330 self.fitted
331 }
332
333 pub fn graph(&self) -> &Graph {
335 &self.graph
336 }
337
338 pub fn catalog(&self) -> &NodeCatalog {
340 &self.catalog
341 }
342
343 pub fn catalog_mut(&mut self) -> &mut NodeCatalog {
345 &mut self.catalog
346 }
347
348 fn graph_config_hash(&self) -> Result<String> {
359 Ok(ArchitectureFingerprint::of(&self.graph)?.digest)
360 }
361}
362
363pub fn graph_run(
374 graph: &Graph,
375 catalog: &NodeCatalog,
376 mode: CompileMode,
377 cache: Arc<dyn CacheStore>,
378) -> Result<HashMap<String, Value>> {
379 GraphSession::new(graph.clone(), catalog.clone())
380 .with_cache(cache)
381 .run(mode)
382}
383
384pub fn graph_fit(
386 graph: &Graph,
387 catalog: &NodeCatalog,
388 x: &Value,
389 y: Option<&Value>,
390 cache: Arc<dyn CacheStore>,
391) -> Result<HashMap<String, Value>> {
392 GraphSession::new(graph.clone(), catalog.clone())
393 .with_cache(cache)
394 .fit(x, y)
395}
396
397pub fn graph_predict(
399 graph: &Graph,
400 catalog: &NodeCatalog,
401 x: &Value,
402 cache: Arc<dyn CacheStore>,
403) -> Result<Value> {
404 GraphSession::new(graph.clone(), catalog.clone())
405 .with_cache(cache)
406 .forward(x)
407}
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412 use crate::cache::MemoryCache;
413 use somatize_compiler::NodeRegistry;
414 use somatize_core::cache::CacheKey;
415 use somatize_core::error::Result;
416 use somatize_core::filter::{FilterKind, FilterMeta, StreamMode};
417 use somatize_core::graph::{Edge, Node};
418
419 struct DoublerFilter;
422 impl somatize_core::filter::Filter for DoublerFilter {
423 fn config_hash(&self) -> CacheKey {
424 CacheKey::from_parts(&[b"Doubler"])
425 }
426 fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
427 Ok(Value::Empty)
428 }
429 fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
430 let (data, shape) = x
431 .as_tensor()
432 .ok_or(SomaError::Other("need tensor".into()))?;
433 Ok(Value::tensor(
434 data.iter().map(|v| v * 2.0).collect(),
435 shape.to_vec(),
436 ))
437 }
438 fn meta(&self) -> FilterMeta {
439 FilterMeta {
440 name: "Doubler".into(),
441 kind: FilterKind::Stateless,
442 cacheable: true,
443 differentiable: true,
444 deterministic: true,
445 stream_mode: StreamMode::FixedState,
446 distribution: somatize_core::filter::Distribution::Local,
447 input_schema: None,
448 output_schema: None,
449 }
450 }
451 }
452
453 struct AdderFilter(f64);
454 impl somatize_core::filter::Filter for AdderFilter {
455 fn config_hash(&self) -> CacheKey {
456 CacheKey::from_parts(&[b"Adder", &self.0.to_le_bytes()])
457 }
458 fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
459 Ok(Value::Empty)
460 }
461 fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
462 let (data, shape) = x
463 .as_tensor()
464 .ok_or(SomaError::Other("need tensor".into()))?;
465 Ok(Value::tensor(
466 data.iter().map(|v| v + self.0).collect(),
467 shape.to_vec(),
468 ))
469 }
470 fn meta(&self) -> FilterMeta {
471 FilterMeta {
472 name: "Adder".into(),
473 kind: FilterKind::Stateless,
474 cacheable: true,
475 differentiable: true,
476 deterministic: true,
477 stream_mode: StreamMode::FixedState,
478 distribution: somatize_core::filter::Distribution::Local,
479 input_schema: None,
480 output_schema: None,
481 }
482 }
483 }
484
485 struct MeanFilter;
486 impl somatize_core::filter::Filter for MeanFilter {
487 fn config_hash(&self) -> CacheKey {
488 CacheKey::from_parts(&[b"Mean"])
489 }
490 fn fit(&self, x: &Value, _y: Option<&Value>) -> Result<Value> {
491 let (data, _) = x
492 .as_tensor()
493 .ok_or(SomaError::Other("need tensor".into()))?;
494 let mean = data.iter().sum::<f64>() / data.len() as f64;
495 Ok(Value::json(serde_json::json!({ "mean": mean })))
496 }
497 fn forward(&self, x: &Value, state: &Value) -> Result<Value> {
498 let (data, shape) = x
499 .as_tensor()
500 .ok_or(SomaError::Other("need tensor".into()))?;
501 let mean = state
502 .as_json()
503 .and_then(|j| j["mean"].as_f64())
504 .unwrap_or(0.0);
505 Ok(Value::tensor(
506 data.iter().map(|v| v - mean).collect(),
507 shape.to_vec(),
508 ))
509 }
510 fn meta(&self) -> FilterMeta {
511 FilterMeta {
512 name: "Mean".into(),
513 kind: FilterKind::Trainable,
514 cacheable: true,
515 differentiable: true,
516 deterministic: true,
517 stream_mode: StreamMode::FixedState,
518 distribution: somatize_core::filter::Distribution::Local,
519 input_schema: None,
520 output_schema: None,
521 }
522 }
523 }
524
525 fn linear_graph(ids: &[&str]) -> Graph {
526 let mut g = Graph::new();
527 for &id in ids {
528 g.nodes.push(Node::new(id, id, id));
529 }
530 for (i, pair) in ids.windows(2).enumerate() {
531 g.edges.push(Edge::data(format!("e{i}"), pair[0], pair[1]));
532 }
533 g
534 }
535
536 #[test]
539 fn session_run_linear() {
540 let graph = linear_graph(&["double", "add"]);
541 let mut lib = NodeCatalog::new();
542 lib.register("double", Box::new(DoublerFilter));
543 lib.register("add", Box::new(AdderFilter(10.0)));
544
545 let mut session = GraphSession::new(graph, lib);
546
547 let cache = MemoryCache::default();
548 session = session.with_cache(Arc::new(cache));
549
550 let CompileResult { plan, .. } = session.compile(CompileMode::NoCache).unwrap();
552 let bus = Arc::new(EventBus::new(64));
553 let mut ctx =
554 Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(session.graph()));
555 ctx.set(
556 somatize_core::keys::GRAPH_INPUT,
557 Value::tensor(vec![1.0, 2.0, 3.0], vec![3]),
558 );
559 executor::execute(&plan, &mut ctx, session.catalog(), &MemoryCache::default()).unwrap();
560
561 let outputs: HashMap<String, Value> = ctx.into_outputs();
562
563 let result = outputs.get("add").unwrap();
564 let (data, _) = result.as_tensor().unwrap();
565 assert_eq!(data, &[12.0, 14.0, 16.0]);
566 }
567
568 #[test]
569 fn session_fit_and_forward() {
570 let graph = linear_graph(&["mean", "double"]);
571 let mut lib = NodeCatalog::new();
572 lib.register("mean", Box::new(MeanFilter));
573 lib.register("double", Box::new(DoublerFilter));
574
575 let mut session = GraphSession::new(graph, lib);
576
577 let x = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
578 let outputs = session.fit(&x, None).unwrap();
579
580 let result = outputs.get("double").unwrap();
583 let (data, _) = result.as_tensor().unwrap();
584 assert_eq!(data, &[-20.0, 0.0, 20.0]);
585
586 assert!(session.is_fitted());
587 }
588
589 #[test]
590 fn session_compile_diagnostics() {
591 let graph = linear_graph(&["double"]);
592 let mut lib = NodeCatalog::new();
593 lib.register("double", Box::new(DoublerFilter));
594
595 let session = GraphSession::new(graph, lib);
596 let result = session.compile(CompileMode::NoCache).unwrap();
597 assert!(result.plan.node_count() > 0);
598 }
599
600 #[test]
603 fn graph_run_linear() {
604 let graph = linear_graph(&["double", "add"]);
605 let mut lib = NodeCatalog::new();
606 lib.register("double", Box::new(DoublerFilter));
607 lib.register("add", Box::new(AdderFilter(10.0)));
608
609 let cache = MemoryCache::default();
610
611 let outputs = {
612 let CompileResult { plan, .. } =
613 compile(&graph, &lib, CompileMode::NoCache, None).unwrap();
614 let bus = Arc::new(EventBus::new(64));
615 let mut ctx = Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(&graph));
616 ctx.set(
617 somatize_core::keys::GRAPH_INPUT,
618 Value::tensor(vec![1.0, 2.0, 3.0], vec![3]),
619 );
620 executor::execute(&plan, &mut ctx, &lib, &cache).unwrap();
621 ctx.into_outputs()
622 };
623
624 let result = outputs.get("add").unwrap();
625 let (data, _) = result.as_tensor().unwrap();
626 assert_eq!(data, &[12.0, 14.0, 16.0]);
627 }
628
629 #[test]
630 fn graph_run_diamond() {
631 let mut graph = Graph::new();
632 graph.nodes.push(Node::new("double", "Double", "double"));
633 graph.nodes.push(Node::new("add", "Add", "add"));
634 graph.nodes.push(Node::new("merge", "Merge", "merge"));
635 graph.edges.push(Edge::data("e1", "double", "merge"));
636 graph.edges.push(Edge::data("e2", "add", "merge"));
637
638 let mut lib = NodeCatalog::new();
639 lib.register("double", Box::new(DoublerFilter));
640 lib.register("add", Box::new(AdderFilter(100.0)));
641
642 struct MergeFilter;
643 impl somatize_core::filter::Filter for MergeFilter {
644 fn config_hash(&self) -> CacheKey {
645 CacheKey::from_parts(&[b"Merge"])
646 }
647 fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
648 Ok(Value::Empty)
649 }
650 fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
651 Ok(x.clone())
652 }
653 fn meta(&self) -> FilterMeta {
654 FilterMeta {
655 name: "Merge".into(),
656 kind: FilterKind::Stateless,
657 cacheable: true,
658 differentiable: false,
659 deterministic: true,
660 stream_mode: StreamMode::FixedState,
661 distribution: somatize_core::filter::Distribution::Local,
662 input_schema: None,
663 output_schema: None,
664 }
665 }
666 }
667 lib.register("merge", Box::new(MergeFilter));
668
669 let cache = MemoryCache::default();
670 let CompileResult { plan, .. } = compile(&graph, &lib, CompileMode::NoCache, None).unwrap();
671
672 let bus = Arc::new(EventBus::new(64));
673 let mut ctx = Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(&graph));
674 ctx.set(
675 somatize_core::keys::GRAPH_INPUT,
676 Value::tensor(vec![5.0], vec![1]),
677 );
678 executor::execute(&plan, &mut ctx, &lib, &cache).unwrap();
679
680 let merge_output = ctx.get("merge").unwrap();
681 assert!(
682 merge_output.as_json().is_some(),
683 "merge should receive JSON from multiple predecessors"
684 );
685 }
686
687 #[test]
688 fn graph_fit_trainable() {
689 let graph = linear_graph(&["mean", "double"]);
690 let mut lib = NodeCatalog::new();
691 lib.register("mean", Box::new(MeanFilter));
692 lib.register("double", Box::new(DoublerFilter));
693
694 let cache = Arc::new(MemoryCache::default());
695 let x = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
696
697 let outputs = graph_fit(&graph, &lib, &x, None, cache.clone()).unwrap();
698
699 let result = outputs.get("double").unwrap();
700 let (data, _) = result.as_tensor().unwrap();
701 assert_eq!(data, &[-20.0, 0.0, 20.0]);
702
703 assert!(!cache.is_empty());
704 }
705
706 #[test]
707 fn the_catalog_is_the_compiler_registry() {
708 let mut lib = NodeCatalog::new();
709 lib.register("a", Box::new(DoublerFilter));
710
711 let registry: &dyn NodeRegistry = &lib;
712 assert!(registry.meta("a").is_some());
713 assert_eq!(registry.meta("a").unwrap().name, "Doubler");
714 assert!(registry.config_hash("a").is_some());
715 assert!(registry.meta("b").is_none());
716 }
717
718 fn session_of(graph: Graph) -> GraphSession {
719 let mut lib = NodeCatalog::new();
720 for node in &graph.nodes {
721 lib.register(&node.id, Box::new(DoublerFilter));
722 }
723 GraphSession::new(graph, lib)
724 }
725
726 #[test]
730 fn state_address_separates_graphs_that_share_node_ids() {
731 let chain = session_of(linear_graph(&["a", "b", "c"]));
732
733 let mut fan = Graph::new();
735 for id in ["a", "b", "c"] {
736 fan.nodes.push(Node::new(id, id, id));
737 }
738 fan.edges.push(Edge::data("e0", "a", "b"));
739 fan.edges.push(Edge::data("e1", "a", "c"));
740 let fan = session_of(fan);
741
742 assert_ne!(
743 chain.graph_config_hash().unwrap(),
744 fan.graph_config_hash().unwrap(),
745 "two differently wired graphs must not persist states to one address"
746 );
747 }
748
749 #[test]
752 fn state_address_is_stable_for_the_same_graph() {
753 assert_eq!(
754 session_of(linear_graph(&["a", "b"]))
755 .graph_config_hash()
756 .unwrap(),
757 session_of(linear_graph(&["a", "b"]))
758 .graph_config_hash()
759 .unwrap()
760 );
761 }
762}