1use crate::executor::{Context, compute_node, output_key, store_output};
34use crate::node_catalog::{NodeCatalog, NodeImpl};
35use somatize_core::cache::{CacheKey, CacheStore};
36use somatize_core::error::{Result, SomaError};
37use somatize_core::event::Event;
38use somatize_core::filter::StreamMode;
39use somatize_core::node::{NodeMeta, NodeOutcome};
40use somatize_core::value::Value;
41use std::sync::Arc;
42use std::time::{Duration, Instant};
43
44struct StreamNode {
46 id: String,
47 node: NodeImpl,
48 meta: NodeMeta,
49 stream_mode: StreamMode,
52 base_state: Arc<Value>,
55 barrier: Vec<Value>,
57 evolving: Option<Value>,
59 started: bool,
60 chunks: u64,
61 cache_hits: u64,
62 cache_misses: u64,
63 compute: Duration,
64}
65
66pub struct StreamRun {
74 nodes: Vec<StreamNode>,
75 chunk_count: usize,
76}
77
78impl StreamRun {
79 pub fn new(node_ids: &[String], catalog: &NodeCatalog) -> Result<Self> {
84 let nodes = node_ids
85 .iter()
86 .map(|id| {
87 let node = catalog
88 .node(id)
89 .ok_or_else(|| SomaError::NodeNotFound(id.clone()))?
90 .clone();
91 let stream_mode = match &node {
94 NodeImpl::Filter(f) => f.meta().stream_mode,
95 NodeImpl::Step(_) => {
96 return Err(SomaError::Execution {
97 node_id: id.clone(),
98 message: "a step cannot run inside a stream plan".into(),
99 });
100 }
101 };
102 let meta = node.meta();
103 let base_state = catalog
104 .get_state(id)
105 .unwrap_or_else(|| Arc::new(Value::Empty));
106 Ok(StreamNode {
107 id: id.clone(),
108 node,
109 meta,
110 stream_mode,
111 base_state,
112 barrier: Vec::new(),
113 evolving: None,
114 started: false,
115 chunks: 0,
116 cache_hits: 0,
117 cache_misses: 0,
118 compute: Duration::ZERO,
119 })
120 })
121 .collect::<Result<Vec<_>>>()?;
122 Ok(Self {
123 nodes,
124 chunk_count: 0,
125 })
126 }
127
128 pub fn process_chunk(
131 &mut self,
132 chunk: Value,
133 ctx: &mut Context,
134 cache: &dyn CacheStore,
135 ) -> Result<Option<Value>> {
136 let stage = format!("chunk {}", self.chunk_count);
137 self.chunk_count += 1;
138 let mut current = chunk;
139 for i in 0..self.nodes.len() {
140 if matches!(self.nodes[i].stream_mode, StreamMode::Barrier) {
141 self.nodes[i].barrier.push(current);
142 return Ok(None);
143 }
144 current = self.run_compute(i, current, &stage, ctx, cache)?;
145 }
146 Ok(Some(current))
147 }
148
149 pub fn flush(&mut self, ctx: &mut Context, cache: &dyn CacheStore) -> Result<Option<Value>> {
153 let mut current: Option<Value> = None;
154 for i in 0..self.nodes.len() {
155 if !self.nodes[i].barrier.is_empty() {
156 let buffer = std::mem::take(&mut self.nodes[i].barrier);
157 let materialized = materialize_buffer(&buffer)?;
158 current = Some(self.run_compute(i, materialized, "flush", ctx, cache)?);
159 } else if let Some(v) = current.take() {
160 current = Some(self.run_compute(i, v, "flush", ctx, cache)?);
161 }
162 }
163 Ok(current)
164 }
165
166 pub fn finish(&mut self, ctx: &Context) {
168 for node in &mut self.nodes {
169 if !node.started {
170 continue;
171 }
172 node.started = false;
173 ctx.event_bus.emit(Event::NodeCompleted {
174 run_id: ctx.run_id.clone(),
175 node_id: node.id.clone(),
176 duration: node.compute,
177 output_summary: format!(
178 "stream: {} chunks, {} hits, {} misses",
179 node.chunks, node.cache_hits, node.cache_misses
180 ),
181 });
182 }
183 }
184
185 pub fn chunks_processed(&self) -> usize {
188 self.chunk_count
189 }
190
191 fn run_compute(
195 &mut self,
196 i: usize,
197 input: Value,
198 stage: &str,
199 ctx: &Context,
200 cache: &dyn CacheStore,
201 ) -> Result<Value> {
202 let node = &mut self.nodes[i];
203 if !node.started {
204 node.started = true;
205 ctx.event_bus.emit(Event::NodeStarted {
206 run_id: ctx.run_id.clone(),
207 node_id: node.id.clone(),
208 kind: node.meta.kind,
209 effectful: node.meta.effectful,
210 });
211 }
212
213 let state_ref: &Value = match &node.evolving {
214 Some(v) => v,
215 None => node.base_state.as_ref(),
216 };
217 let input_key = CacheKey::for_value(&input);
218 let key = output_key(&node.node, &node.meta, state_ref, &input_key, ctx.seed);
219
220 if let Some(k) = &key {
221 if let Ok(Some((cached, _tier))) = cache.get_located(k) {
222 node.chunks += 1;
223 node.cache_hits += 1;
224 if matches!(node.stream_mode, StreamMode::Evolving) {
225 node.evolving = Some(cached.clone());
226 }
227 return Ok(cached);
228 }
229 node.cache_misses += 1;
230 }
231
232 let started_at = Instant::now();
233 match compute_node(&node.node, &node.id, ctx, &input, state_ref) {
234 Ok(NodeOutcome::Produced(out)) => {
235 let duration = started_at.elapsed();
236 node.compute += duration;
237 node.chunks += 1;
238 if let Some(k) = &key {
239 store_output(
240 cache,
241 k,
242 &out,
243 &node.id,
244 &ctx.run_id,
245 duration,
246 node.meta.deterministic,
247 );
248 }
249 if matches!(node.stream_mode, StreamMode::Evolving) {
250 node.evolving = Some(out.clone());
251 }
252 Ok(out)
253 }
254 Ok(NodeOutcome::HandOff { .. } | NodeOutcome::Paused { .. }) => {
258 Err(SomaError::Execution {
259 node_id: node.id.clone(),
260 message: "a step cannot run inside a stream plan".into(),
261 })
262 }
263 Err(e) => {
264 ctx.event_bus.emit(Event::NodeFailed {
265 run_id: ctx.run_id.clone(),
266 node_id: node.id.clone(),
267 error: format!("{stage}: {e}"),
268 });
269 Err(e)
270 }
271 }
272 }
273}
274
275#[derive(Default)]
283pub struct StreamOutput {
284 all_data: Vec<f64>,
285 result_shape: Option<Vec<usize>>,
286 non_tensor: Option<Value>,
287}
288
289impl StreamOutput {
290 pub fn new() -> Self {
292 Self::default()
293 }
294
295 pub fn push(&mut self, output: Value) {
297 match output {
298 Value::Tensor { values, shape } => {
299 if self.result_shape.is_none() {
300 self.result_shape = Some(shape);
301 }
302 self.all_data.extend_from_slice(values.as_slice());
303 }
304 other => self.non_tensor = Some(other),
305 }
306 }
307
308 pub fn finish(self) -> Value {
311 if let Some(mut shape) = self.result_shape {
312 let row_size: usize = shape.iter().skip(1).product::<usize>().max(1);
313 shape[0] = self.all_data.len() / row_size;
314 return Value::tensor(self.all_data, shape);
315 }
316 self.non_tensor.unwrap_or(Value::Empty)
317 }
318}
319
320pub fn materialize_buffer(buffer: &[Value]) -> Result<Value> {
322 if buffer.is_empty() {
323 return Ok(Value::Empty);
324 }
325 let mut all_data = Vec::new();
326 let mut total_rows = 0;
327 let mut cols = 0;
328
329 for chunk in buffer {
330 match chunk {
331 Value::Tensor { values, shape } => {
332 all_data.extend(values.iter());
333 if shape.len() == 1 {
334 total_rows += shape[0];
335 cols = 1;
336 } else if shape.len() >= 2 {
337 total_rows += shape[0];
338 cols = shape[1];
339 }
340 }
341 _ => {
342 return Err(SomaError::Other(
343 "barrier buffer contains non-tensor values".into(),
344 ));
345 }
346 }
347 }
348
349 if cols <= 1 {
350 Ok(Value::tensor(all_data, vec![total_rows]))
351 } else {
352 Ok(Value::tensor(all_data, vec![total_rows, cols]))
353 }
354}
355
356#[cfg(test)]
357mod tests {
358 use super::*;
359 use crate::cache::memory::MemoryCache;
360 use crate::event_bus::EventBus;
361 use somatize_core::error::Result as SomaResult;
362 use somatize_core::filter::{Distribution, Filter, FilterKind, FilterMeta};
363
364 fn meta(name: &str, stream_mode: StreamMode, cacheable: bool) -> FilterMeta {
365 FilterMeta {
366 name: name.into(),
367 kind: FilterKind::Stateless,
368 cacheable,
369 differentiable: false,
370 deterministic: true,
371 stream_mode,
372 distribution: Distribution::Local,
373 input_schema: None,
374 output_schema: None,
375 }
376 }
377
378 struct DoubleChunk;
379 impl Filter for DoubleChunk {
380 fn config_hash(&self) -> CacheKey {
381 CacheKey::from_parts(&[b"DoubleChunk"])
382 }
383 fn fit(&self, _x: &Value, _y: Option<&Value>) -> SomaResult<Value> {
384 Ok(Value::Empty)
385 }
386 fn forward(&self, x: &Value, _state: &Value) -> SomaResult<Value> {
387 if let Value::Tensor { values, shape } = x {
388 Ok(Value::tensor(
389 values.iter().map(|v| v * 2.0).collect(),
390 shape.clone(),
391 ))
392 } else {
393 Ok(x.clone())
394 }
395 }
396 fn meta(&self) -> FilterMeta {
397 meta("DoubleChunk", StreamMode::FixedState, true)
398 }
399 }
400
401 struct UncachedDouble;
403 impl Filter for UncachedDouble {
404 fn config_hash(&self) -> CacheKey {
405 CacheKey::from_parts(&[b"UncachedDouble"])
406 }
407 fn fit(&self, _x: &Value, _y: Option<&Value>) -> SomaResult<Value> {
408 Ok(Value::Empty)
409 }
410 fn forward(&self, x: &Value, _state: &Value) -> SomaResult<Value> {
411 DoubleChunk.forward(x, &Value::Empty)
412 }
413 fn meta(&self) -> FilterMeta {
414 meta("UncachedDouble", StreamMode::FixedState, false)
415 }
416 }
417
418 struct Accumulator;
420 impl Filter for Accumulator {
421 fn config_hash(&self) -> CacheKey {
422 CacheKey::from_parts(&[b"Accumulator"])
423 }
424 fn fit(&self, _x: &Value, _y: Option<&Value>) -> SomaResult<Value> {
425 Ok(Value::Empty)
426 }
427 fn forward(&self, x: &Value, _state: &Value) -> SomaResult<Value> {
428 Ok(x.clone())
429 }
430 fn meta(&self) -> FilterMeta {
431 meta("Accumulator", StreamMode::Barrier, true)
432 }
433 }
434
435 struct RunningSum;
438 impl Filter for RunningSum {
439 fn config_hash(&self) -> CacheKey {
440 CacheKey::from_parts(&[b"RunningSum"])
441 }
442 fn fit(&self, _x: &Value, _y: Option<&Value>) -> SomaResult<Value> {
443 Ok(Value::tensor(vec![0.0], vec![1]))
444 }
445 fn forward(&self, x: &Value, state: &Value) -> SomaResult<Value> {
446 let x_sum: f64 = match x {
447 Value::Tensor { values, .. } => values.iter().sum(),
448 _ => 0.0,
449 };
450 let state_sum: f64 = match state {
451 Value::Tensor { values, .. } => values.first().copied().unwrap_or(0.0),
452 _ => 0.0,
453 };
454 Ok(Value::tensor(vec![x_sum + state_sum], vec![1]))
455 }
456 fn meta(&self) -> FilterMeta {
457 let mut m = meta("RunningSum", StreamMode::Evolving, false);
458 m.kind = FilterKind::Trainable;
459 m
460 }
461 }
462
463 struct Panicker;
464 impl Filter for Panicker {
465 fn config_hash(&self) -> CacheKey {
466 CacheKey::from_parts(&[b"Panicker"])
467 }
468 fn fit(&self, _x: &Value, _y: Option<&Value>) -> SomaResult<Value> {
469 Ok(Value::Empty)
470 }
471 fn forward(&self, _x: &Value, _state: &Value) -> SomaResult<Value> {
472 panic!("chunk went sideways")
473 }
474 fn meta(&self) -> FilterMeta {
475 meta("Panicker", StreamMode::FixedState, true)
476 }
477 }
478
479 fn harness(nodes: Vec<(&str, Box<dyn Filter>)>) -> (StreamRun, Context, MemoryCache) {
480 let mut catalog = NodeCatalog::new();
481 let mut ids = Vec::new();
482 for (id, filter) in nodes {
483 catalog.register(id, filter);
484 ids.push(id.to_string());
485 }
486 let run = StreamRun::new(&ids, &catalog).unwrap();
487 let ctx = Context::new(Arc::new(EventBus::new(64)), "stream-test");
488 (run, ctx, MemoryCache::default())
489 }
490
491 fn tensor(vals: &[f64]) -> Value {
492 Value::tensor(vals.to_vec(), vec![vals.len()])
493 }
494
495 #[test]
496 fn fixed_state_processes_each_chunk() {
497 let (mut run, mut ctx, cache) = harness(vec![("double", Box::new(DoubleChunk))]);
498 let out = run
499 .process_chunk(tensor(&[1.0, 2.0]), &mut ctx, &cache)
500 .unwrap();
501 assert_eq!(out, Some(tensor(&[2.0, 4.0])));
502 let out = run.process_chunk(tensor(&[3.0]), &mut ctx, &cache).unwrap();
503 assert_eq!(out, Some(tensor(&[6.0])));
504 }
505
506 #[test]
507 fn barrier_accumulates_then_flushes() {
508 let (mut run, mut ctx, cache) = harness(vec![("acc", Box::new(Accumulator))]);
509 assert_eq!(
510 run.process_chunk(tensor(&[1.0, 2.0]), &mut ctx, &cache)
511 .unwrap(),
512 None
513 );
514 assert_eq!(
515 run.process_chunk(tensor(&[3.0, 4.0]), &mut ctx, &cache)
516 .unwrap(),
517 None
518 );
519 let flushed = run.flush(&mut ctx, &cache).unwrap().unwrap();
520 assert_eq!(flushed, tensor(&[1.0, 2.0, 3.0, 4.0]));
521 }
522
523 #[test]
524 fn evolving_state_accumulates() {
525 let (mut run, mut ctx, cache) = harness(vec![("sum", Box::new(RunningSum))]);
526 let r1 = run
527 .process_chunk(tensor(&[10.0]), &mut ctx, &cache)
528 .unwrap()
529 .unwrap();
530 assert_eq!(r1, tensor(&[10.0]));
531 let r2 = run
532 .process_chunk(tensor(&[5.0]), &mut ctx, &cache)
533 .unwrap()
534 .unwrap();
535 assert_eq!(r2, tensor(&[15.0]), "10 + 5: the output was the state");
536 }
537
538 #[test]
539 fn mixed_pipeline_fixed_then_barrier() {
540 let (mut run, mut ctx, cache) = harness(vec![
541 ("double", Box::new(DoubleChunk)),
542 ("acc", Box::new(Accumulator)),
543 ]);
544 assert_eq!(
545 run.process_chunk(tensor(&[1.0]), &mut ctx, &cache).unwrap(),
546 None
547 );
548 assert_eq!(
549 run.process_chunk(tensor(&[2.0]), &mut ctx, &cache).unwrap(),
550 None
551 );
552 let flushed = run.flush(&mut ctx, &cache).unwrap().unwrap();
553 assert_eq!(flushed, tensor(&[2.0, 4.0]), "doubled then accumulated");
554 }
555
556 #[test]
559 fn uncacheable_chunks_are_not_cached() {
560 let (mut run, mut ctx, cache) = harness(vec![("raw", Box::new(UncachedDouble))]);
561 run.process_chunk(tensor(&[1.0]), &mut ctx, &cache).unwrap();
562 run.process_chunk(tensor(&[2.0]), &mut ctx, &cache).unwrap();
563 assert!(
564 cache.is_empty(),
565 "an uncacheable filter's chunks reached the store"
566 );
567 }
568
569 #[test]
572 fn cached_chunks_are_served_and_counted() {
573 let mut catalog = NodeCatalog::new();
574 catalog.register("double", Box::new(DoubleChunk));
575 let ids = vec!["double".to_string()];
576 let cache = MemoryCache::default();
577 let mut ctx = Context::new(Arc::new(EventBus::new(64)), "stream-test");
578
579 let mut first = StreamRun::new(&ids, &catalog).unwrap();
580 let a = first
581 .process_chunk(tensor(&[5.0]), &mut ctx, &cache)
582 .unwrap();
583 assert!(!cache.is_empty(), "the chunk should have been cached");
584
585 let mut second = StreamRun::new(&ids, &catalog).unwrap();
586 let b = second
587 .process_chunk(tensor(&[5.0]), &mut ctx, &cache)
588 .unwrap();
589 assert_eq!(a, b);
590 assert_eq!(second.nodes[0].cache_hits, 1);
591 assert_eq!(second.nodes[0].cache_misses, 0);
592 }
593
594 #[test]
598 fn a_chunk_cache_key_follows_the_run_seed() {
599 let mut catalog = NodeCatalog::new();
600 catalog.register("double", Box::new(DoubleChunk));
601 let ids = vec!["double".to_string()];
602 let cache = MemoryCache::default();
603 let bus = Arc::new(EventBus::new(64));
604
605 for seed in [Some(1), Some(2), None] {
606 let mut ctx = Context::new(bus.clone(), "stream-test").with_seed(seed);
607 let mut run = StreamRun::new(&ids, &catalog).unwrap();
608 run.process_chunk(tensor(&[1.0, 2.0]), &mut ctx, &cache)
609 .unwrap();
610 }
611 assert_eq!(
612 cache.len(),
613 3,
614 "each seed must own its own cache line for the same chunk"
615 );
616 }
617
618 #[test]
621 fn non_finite_chunks_do_not_share_a_cache_key() {
622 let nan = tensor(&[f64::NAN]);
623 let inf = tensor(&[f64::INFINITY]);
624 assert_eq!(
625 serde_json::to_vec(&nan).unwrap(),
626 serde_json::to_vec(&inf).unwrap(),
627 "if this ever stops being true the bug is gone by other means"
628 );
629
630 let (mut run, mut ctx, cache) = harness(vec![("double", Box::new(DoubleChunk))]);
631 let out_nan = run.process_chunk(nan, &mut ctx, &cache).unwrap().unwrap();
632 let out_inf = run.process_chunk(inf, &mut ctx, &cache).unwrap().unwrap();
633
634 let first = |v: &Value| match v {
635 Value::Tensor { values, .. } => values[0],
636 other => panic!("expected a tensor, got {other:?}"),
637 };
638 assert!(first(&out_nan).is_nan(), "NaN doubled is still NaN");
639 assert_eq!(
640 first(&out_inf),
641 f64::INFINITY,
642 "the infinite chunk was served the NaN chunk's cached output"
643 );
644 }
645
646 #[test]
649 fn a_panicking_chunk_is_contained() {
650 let (mut run, mut ctx, cache) = harness(vec![("boom", Box::new(Panicker))]);
651 let err = run
652 .process_chunk(tensor(&[1.0]), &mut ctx, &cache)
653 .unwrap_err();
654 assert!(err.to_string().contains("panicked"), "{err}");
655 }
656
657 #[test]
661 fn barrier_flush_goes_through_the_cache() {
662 let mut catalog = NodeCatalog::new();
663 catalog.register("acc", Box::new(Accumulator));
664 let ids = vec!["acc".to_string()];
665 let cache = MemoryCache::default();
666 let mut ctx = Context::new(Arc::new(EventBus::new(64)), "stream-test");
667
668 let mut first = StreamRun::new(&ids, &catalog).unwrap();
669 first
670 .process_chunk(tensor(&[1.0]), &mut ctx, &cache)
671 .unwrap();
672 first
673 .process_chunk(tensor(&[2.0]), &mut ctx, &cache)
674 .unwrap();
675 first.flush(&mut ctx, &cache).unwrap();
676 assert!(!cache.is_empty(), "the flush output should be cached");
677
678 let mut second = StreamRun::new(&ids, &catalog).unwrap();
679 second
680 .process_chunk(tensor(&[1.0]), &mut ctx, &cache)
681 .unwrap();
682 second
683 .process_chunk(tensor(&[2.0]), &mut ctx, &cache)
684 .unwrap();
685 second.flush(&mut ctx, &cache).unwrap();
686 assert_eq!(second.nodes[0].cache_hits, 1, "the flush should be a hit");
687 }
688
689 #[test]
692 fn an_unknown_node_is_an_error_not_a_skip() {
693 let catalog = NodeCatalog::new();
694 let Err(err) = StreamRun::new(&["ghost".to_string()], &catalog) else {
695 panic!("an unknown node must not stream");
696 };
697 assert!(matches!(err, SomaError::NodeNotFound(id) if id == "ghost"));
698 }
699}