feat: creating and removing function executors

Author: seriousben

Diff for: indexify/tests/cli/test_server_task_distribution.py

@@ -30,9 +30,9 @@ def success_func(sleep_secs: float) -> str:
 class TestServerTaskDistribution(unittest.TestCase):
     def test_server_distributes_invocations_fairly_between_two_executors(self):
         print(
-            "Waiting for 10 seconds for Server to notice that any previously existing Executors exited."
+            "Waiting for 30 seconds for Server to notice that any previously existing Executors exited."
         )
-        time.sleep(10)
+        time.sleep(30)
 
         with ExecutorProcessContextManager(
             [
@@ -81,9 +81,9 @@ def test_server_distributes_invocations_fairly_between_two_executors(self):
 
     def test_server_redistributes_invocations_when_new_executor_joins(self):
         print(
-            "Waiting for 10 seconds for Server to notice that any previously existing Executors exited."
+            "Waiting for 30 seconds for Server to notice that any previously existing Executors exited."
         )
-        time.sleep(10)
+        time.sleep(30)
 
         graph = Graph(
             name=test_graph_name(self),

Diff for: server/data_model/src/lib.rs

@@ -1221,6 +1221,30 @@ pub struct FunctionURI {
     pub version: Option<GraphVersion>,
 }
 
+impl FunctionURI {
+    pub fn matches(&self, other: &FunctionURI) -> bool {
+        self.namespace == other.namespace &&
+            self.compute_graph_name == other.compute_graph_name &&
+            self.compute_fn_name == other.compute_fn_name &&
+            (self.version.is_none() || self.version == other.version)
+    }
+}
+
+impl Display for FunctionURI {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}|{}|{}|{}",
+            self.namespace,
+            self.compute_graph_name,
+            self.compute_fn_name,
+            self.version
+                .as_ref()
+                .map_or("None".to_string(), |v| v.to_string())
+        )
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct GpuResources {
     pub count: u32,
@@ -1315,13 +1339,13 @@ impl FunctionExecutor {
     ///
     /// A match occurs when:
     /// 1. The namespace, compute_graph_name, and compute_fn_name are equal
-    /// 2. Either the FunctionURI has no version specified (None), OR
-    ///    the versions match
+    /// 2. Either the FunctionURI has no version specified (None), OR the
+    ///    versions match
     pub fn matches(&self, uri: &FunctionURI) -> bool {
         // Check if namespace, compute_graph_name, and compute_fn_name match
-        let basic_match = self.namespace == uri.namespace
-            && self.compute_graph_name == uri.compute_graph_name
-            && self.compute_fn_name == uri.compute_fn_name;
+        let basic_match = self.namespace == uri.namespace &&
+            self.compute_graph_name == uri.compute_graph_name &&
+            self.compute_fn_name == uri.compute_fn_name;
 
         // If the basic fields match, then check the version
         if basic_match {
@@ -1331,6 +1355,13 @@ impl FunctionExecutor {
             false
         }
     }
+
+    pub fn matches_fn(&self, other: &FunctionExecutor) -> bool {
+        self.namespace == other.namespace &&
+            self.compute_graph_name == other.compute_graph_name &&
+            self.compute_fn_name == other.compute_fn_name &&
+            self.version == other.version
+    }
 }
 
 impl FunctionExecutorBuilder {
@@ -1617,12 +1648,18 @@ pub struct Namespace {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use std::collections::HashMap;
 
+    use super::*;
     use crate::{
-        test_objects::tests::test_compute_fn, ComputeGraph, ComputeGraphCode, ComputeGraphVersion,
-        DynamicEdgeRouter, GraphVersion, Node, RuntimeInformation,
+        test_objects::tests::test_compute_fn,
+        ComputeGraph,
+        ComputeGraphCode,
+        ComputeGraphVersion,
+        DynamicEdgeRouter,
+        GraphVersion,
+        Node,
+        RuntimeInformation,
     };
 
     #[test]

Diff for: server/data_model/src/test_objects.rs

@@ -4,13 +4,28 @@ pub mod tests {
     use rand::{distributions::Alphanumeric, Rng};
 
     use super::super::{
-        ComputeFn, ComputeGraph, ComputeGraphCode, Node, Node::Compute, NodeOutput,
+        ComputeFn,
+        ComputeGraph,
+        ComputeGraphCode,
+        Node,
+        Node::Compute,
+        NodeOutput,
         RuntimeInformation,
     };
     use crate::{
-        DataPayload, DynamicEdgeRouter, ExecutorId, ExecutorMetadata, GraphInvocationCtx,
-        GraphInvocationCtxBuilder, GraphVersion, ImageInformation, InvocationPayload,
-        InvocationPayloadBuilder, NodeOutputBuilder, Task, TaskBuilder,
+        DataPayload,
+        DynamicEdgeRouter,
+        ExecutorId,
+        ExecutorMetadata,
+        GraphInvocationCtx,
+        GraphInvocationCtxBuilder,
+        GraphVersion,
+        ImageInformation,
+        InvocationPayload,
+        InvocationPayloadBuilder,
+        NodeOutputBuilder,
+        Task,
+        TaskBuilder,
     };
 
     pub const TEST_NAMESPACE: &str = "test_ns";

Diff for: server/processor/src/graph_processor.rs

@@ -7,7 +7,9 @@ use metrics::{low_latency_boundaries, Timer};
 use opentelemetry::{metrics::Histogram, KeyValue};
 use state_store::{
     requests::{
-        DeleteComputeGraphRequest, DeleteInvocationRequest, RequestPayload,
+        DeleteComputeGraphRequest,
+        DeleteInvocationRequest,
+        RequestPayload,
         StateMachineUpdateRequest,
     },
     IndexifyState,
@@ -207,67 +209,45 @@ impl GraphProcessor {
     ) -> Result<StateMachineUpdateRequest> {
         debug!("processing state change: {}", state_change);
         let mut indexes = self.indexify_state.in_memory_state.read().await.clone();
-        match &state_change.change_type {
+        let req = match &state_change.change_type {
             ChangeType::InvokeComputeGraph(_) | ChangeType::TaskOutputsIngested(_) => {
-                let scheduler_update = self
+                let mut scheduler_update = self
                     .task_creator
                     .invoke(&state_change.change_type, &mut indexes)
-                    .await;
-                if let Ok(mut result) = scheduler_update {
-                    let placement_result = self.task_allocator.allocate(&mut indexes)?;
-                    result
-                        .new_allocations
-                        .extend(placement_result.new_allocations);
-                    result.updated_tasks.extend(placement_result.updated_tasks);
-                    Ok(StateMachineUpdateRequest {
-                        payload: RequestPayload::SchedulerUpdate(Box::new(result)),
-                        processed_state_changes: vec![state_change.clone()],
-                    })
-                } else {
-                    error!("error invoking task creator: {:?}", scheduler_update.err());
-                    Ok(StateMachineUpdateRequest {
-                        payload: RequestPayload::Noop,
-                        processed_state_changes: vec![state_change.clone()],
-                    })
+                    .await?;
+                scheduler_update.extend(self.task_allocator.allocate(&mut indexes)?);
+                StateMachineUpdateRequest {
+                    payload: RequestPayload::SchedulerUpdate(Box::new(scheduler_update)),
+                    processed_state_changes: vec![state_change.clone()],
                 }
             }
-            ChangeType::ExecutorUpserted(_)
-            | ChangeType::ExecutorRemoved(_)
-            | ChangeType::TombStoneExecutor(_) => {
+            ChangeType::ExecutorUpserted(_) |
+            ChangeType::ExecutorRemoved(_) |
+            ChangeType::TombStoneExecutor(_) => {
                 let scheduler_update = self
                     .task_allocator
-                    .invoke(&state_change.change_type, &mut indexes);
-                if let Ok(result) = scheduler_update {
-                    Ok(StateMachineUpdateRequest {
-                        payload: RequestPayload::SchedulerUpdate(Box::new(result)),
-                        processed_state_changes: vec![state_change.clone()],
-                    })
-                } else {
-                    error!(
-                        "error invoking task allocator: {:?}",
-                        scheduler_update.err()
-                    );
-                    Ok(StateMachineUpdateRequest {
-                        payload: RequestPayload::Noop,
-                        processed_state_changes: vec![state_change.clone()],
-                    })
+                    .invoke(&state_change.change_type, &mut indexes)?;
+                StateMachineUpdateRequest {
+                    payload: RequestPayload::SchedulerUpdate(Box::new(scheduler_update)),
+                    processed_state_changes: vec![state_change.clone()],
                 }
             }
-            ChangeType::TombstoneComputeGraph(request) => Ok(StateMachineUpdateRequest {
+            ChangeType::TombstoneComputeGraph(request) => StateMachineUpdateRequest {
                 payload: RequestPayload::DeleteComputeGraphRequest(DeleteComputeGraphRequest {
                     namespace: request.namespace.clone(),
                     name: request.compute_graph.clone(),
                 }),
                 processed_state_changes: vec![state_change.clone()],
-            }),
-            ChangeType::TombstoneInvocation(request) => Ok(StateMachineUpdateRequest {
+            },
+            ChangeType::TombstoneInvocation(request) => StateMachineUpdateRequest {
                 payload: RequestPayload::DeleteInvocationRequest(DeleteInvocationRequest {
                     namespace: request.namespace.clone(),
                     compute_graph: request.compute_graph.clone(),
                     invocation_id: request.invocation_id.clone(),
                 }),
                 processed_state_changes: vec![state_change.clone()],
-            }),
-        }
+            },
+        };
+        Ok(req)
     }
 }