In memory state

server/Cargo.toml

@@ -49,7 +49,7 @@ object_store = {git = "https://github.com/tensorlakeai/arrow-rs", branch="remove
 futures = "0.3.31"
 bytes = "1.9.0"
 pin-project-lite = "0.2.16"
-async-trait = "0.1.85"
+async-trait = "0.1.86"
 tokio-stream = "0.1.17"
 slatedb = { git = "https://github.com/diptanu/slatedb" }
 rust-embed = { version = "8.5.0", features = ["mime-guess"] }
@@ -64,6 +64,7 @@ tower-http = { version = "0.6.2", default-features = false, features = [
     "cors",
     "trace",
 ] }
+im = {version = "15.1.0", features = ["serde"]}
 pin-project = "1.1.8"
 ciborium = "0.2.2"
 uuid = { version = "1.12.1", features = ["v4"] }
@@ -128,6 +129,7 @@ axum-tracing-opentelemetry = { version = "0.25.0", features = [
     "tracing_level_info",
 ] }
 tower-otel-http-metrics = { workspace = true }
+im = { workspace = true }
 
 [dev-dependencies]
 tempfile = { workspace = true }

server/processor/src/graph_processor.rs

@@ -1,7 +1,13 @@
 use std::{sync::Arc, vec};
 
 use anyhow::Result;
-use data_model::{ChangeType, StateChange};
+use data_model::{
+    ChangeType,
+    GraphInvocationCtx,
+    StateChange,
+    TaskOutcome,
+    TaskOutputsIngestedEvent,
+};
 use state_store::{
     requests::{
         DeleteComputeGraphRequest,
@@ -10,11 +16,10 @@
         MutateClusterTopologyRequest,
         NamespaceProcessorUpdateRequest,
         ReductionTasks,
         RequestPayload,
         StateMachineUpdateRequest,
         TaskAllocationUpdateRequest,
-    },
-    IndexifyState,
+    }, IndexifyState
 };
 use tokio::sync::Notify;
 use tracing::{error, info};
@@ -173,7 +178,10 @@
                 info!("invoking compute graph: {:?}", event);
                 let task_creation_result = self
                     .task_creator
-                    .handle_invoke_compute_graph(event.clone())
+                    .handle_invoke_compute_graph(
+                        event.clone(),
+                        &self.indexify_state.in_memory_state().await,
+                    )
                     .await?;
                 Ok(task_creation_result_to_sm_update(
                     &event.namespace,
@@ -183,24 +191,33 @@
                     &state_change,
                 ))
             }
-            ChangeType::TaskOutputsIngested(event) => Ok(StateMachineUpdateRequest {
-                payload: RequestPayload::FinalizeTask(FinalizeTaskRequest {
-                    namespace: event.namespace.clone(),
-                    compute_graph: event.compute_graph.clone(),
-                    compute_fn: event.compute_fn.clone(),
-                    invocation_id: event.invocation_id.clone(),
-                    task_id: event.task_id.clone(),
-                    task_outcome: event.outcome.clone(),
-                    executor_id: event.executor_id.clone(),
-                    diagnostics: event.diagnostic.clone(),
-                }),
-                processed_state_changes: vec![state_change.clone()],
-            }),
+            ChangeType::TaskOutputsIngested(event) => {
+                let graph_ctx =
+                    update_graph_ctx_finalize_task(event, self.indexify_state.clone()).await;
+                if let None = graph_ctx {
+                    return Ok(StateMachineUpdateRequest {
+                        payload: RequestPayload::Noop,
+                        processed_state_changes: vec![state_change.clone()],
+                    });
+                }
+                Ok(StateMachineUpdateRequest {
+                    payload: RequestPayload::FinalizeTask(FinalizeTaskRequest {
+                        namespace: event.namespace.clone(),
+                        compute_graph: event.compute_graph.clone(),
+                        compute_fn: event.compute_fn.clone(),
+                        invocation_id: event.invocation_id.clone(),
+                        task_id: event.task_id.clone(),
+                        task_outcome: event.outcome.clone(),
+                        executor_id: event.executor_id.clone(),
+                        diagnostics: event.diagnostic.clone(),
+                        invocation_ctx: graph_ctx,
+                    }),
+                    processed_state_changes: vec![state_change.clone()],
+                })
+            }
             ChangeType::TaskFinalized(event) => {
-                let task_creation_result = self
-                    .task_creator
-                    .handle_task_finished_inner(self.indexify_state.clone(), event)
-                    .await?;
+                let task_creation_result =
+                    self.task_creator.handle_task_finished_inner(event).await?;
                 Ok(task_creation_result_to_sm_update(
                     &event.namespace,
                     &event.compute_graph,
@@ -283,7 +300,7 @@
     state_change: &StateChange,
 ) -> StateMachineUpdateRequest {
     StateMachineUpdateRequest {
-        payload: RequestPayload::NamespaceProcessorUpdate(NamespaceProcessorUpdateRequest {
+        payload: RequestPayload::TaskCreatorUpdate(NamespaceProcessorUpdateRequest {
             namespace: ns.to_string(),
             compute_graph: compute_graph.to_string(),
             invocation_id: invocation_id.to_string(),
@@ -292,6 +309,7 @@
                 new_reduction_tasks: task_creation_result.new_reduction_tasks,
                 processed_reduction_tasks: task_creation_result.processed_reduction_tasks,
             },
+            invocation_ctx: task_creation_result.invocation_ctx,
         }),
         processed_state_changes: vec![state_change.clone()],
     }
@@ -310,3 +328,35 @@
         processed_state_changes: vec![state_change.clone()],
     }
 }
+
+async fn update_graph_ctx_finalize_task(
+    event: &TaskOutputsIngestedEvent,
+    indexify_state: Arc<IndexifyState>,
+) -> Option<GraphInvocationCtx> {
+    let in_memory_state = indexify_state.in_memory_state().await;
+    let invocation_ctx_map= in_memory_state
+        .invocation_ctx
+        .clone();
+
+    let may_be_invocation_ctx = invocation_ctx_map
+        .get(&format!(
+            "{}|{}|{}",
+            event.namespace, event.compute_graph, event.invocation_id
+        ));
+    if may_be_invocation_ctx.is_none() {
+        error!("namespace: {}, invocation: {}, cg: {}, fn: {}, task: {} no invocation ctx found for task outputs ingested event",
+    event.namespace, event.invocation_id, event.compute_graph, event.compute_fn, event.task_id);
+        return None;
+    }
+    let mut invocation_ctx = may_be_invocation_ctx.unwrap().clone();
+    invocation_ctx.outstanding_tasks -= 1;
+    invocation_ctx
+        .fn_task_analytics
+        .entry(event.compute_fn.clone())
+        .and_modify(|e| match event.outcome {
+            TaskOutcome::Success => e.success(),
+            TaskOutcome::Failure => e.fail(),
+            _ => {}
+        });
+    Some(invocation_ctx)
+}