added queue size config (#1497)

* added queue size config

* Move comments about queue size to struct field.

---------

Co-authored-by: Julio Martinez <julio@tensorflow.ai>

Author: diptanu

server/processor/src/graph_processor.rs

@@ -24,10 +24,15 @@ pub struct GraphProcessor {
     pub task_cache: Arc<task_cache::TaskCache>,
     pub state_transition_latency: Histogram<f64>,
     pub processor_processing_latency: Histogram<f64>,
+    pub queue_size: u32,
 }
 
 impl GraphProcessor {
-    pub fn new(indexify_state: Arc<IndexifyState>, task_cache: Arc<task_cache::TaskCache>) -> Self {
+    pub fn new(
+        indexify_state: Arc<IndexifyState>,
+        task_cache: Arc<task_cache::TaskCache>,
+        queue_size: u32,
+    ) -> Self {
         let meter = opentelemetry::global::meter("processor_metrics");
 
         let processor_processing_latency = meter
@@ -49,6 +54,7 @@ impl GraphProcessor {
             task_cache,
             state_transition_latency,
             processor_processing_latency,
+            queue_size,
         }
     }
 
@@ -214,7 +220,7 @@ impl GraphProcessor {
 
                 scheduler_update.extend(self.task_cache.try_allocate(indexes.clone()));
 
-                scheduler_update.extend(task_allocator::allocate(indexes, clock)?);
+                scheduler_update.extend(task_allocator::allocate(indexes, clock, self.queue_size)?);
                 StateMachineUpdateRequest {
                     payload: RequestPayload::SchedulerUpdate(Box::new(scheduler_update)),
                     processed_state_changes: vec![state_change.clone()],
@@ -223,8 +229,12 @@ impl GraphProcessor {
             ChangeType::ExecutorUpserted(_) |
             ChangeType::ExecutorRemoved(_) |
             ChangeType::TombStoneExecutor(_) => {
-                let scheduler_update =
-                    task_allocator::invoke(indexes, clock, &state_change.change_type)?;
+                let scheduler_update = task_allocator::invoke(
+                    indexes,
+                    clock,
+                    &state_change.change_type,
+                    self.queue_size,
+                )?;
                 StateMachineUpdateRequest {
                     payload: RequestPayload::SchedulerUpdate(Box::new(scheduler_update)),
                     processed_state_changes: vec![state_change.clone()],

server/processor/src/task_allocator.rs

@@ -29,30 +29,32 @@ use state_store::{
 };
 use tracing::{debug, error, info, info_span, warn};
 
-// Maximum number of allocations per executor.
-//
-// In the future, this should be a dynamic value based on:
-// - function concurrency configuration
-// - function batching configuration
-// - function timeout configuration
-const MAX_ALLOCATIONS_PER_FN_EXECUTOR: usize = 20;
-
 struct TaskAllocationProcessor<'a> {
     in_memory_state: &'a mut InMemoryState,
     clock: u64,
+
+    // Maximum number of allocations per executor.
+    //
+    // In the future, this should be a dynamic value based on:
+    // - function concurrency configuration
+    // - function batching configuration
+    // - function timeout configuration
+    queue_size: u32,
 }
 
 #[tracing::instrument(skip(in_memory_state, clock, change))]
 pub fn invoke(
     in_memory_state: Arc<RwLock<InMemoryState>>,
     clock: u64,
     change: &ChangeType,
+    queue_size: u32,
 ) -> Result<SchedulerUpdateRequest> {
     let mut in_memory_state = in_memory_state.write().unwrap();
 
     let mut task_allocator = TaskAllocationProcessor {
         in_memory_state: &mut in_memory_state.deref_mut(),
         clock,
+        queue_size,
     };
 
     task_allocator.invoke(change)
@@ -64,12 +66,14 @@ pub fn invoke(
 pub fn allocate(
     in_memory_state: Arc<RwLock<InMemoryState>>,
     clock: u64,
+    queue_size: u32,
 ) -> Result<SchedulerUpdateRequest> {
     let mut in_memory_state = in_memory_state.write().unwrap();
 
     let mut task_allocator = TaskAllocationProcessor {
         in_memory_state: &mut in_memory_state.deref_mut(),
         clock,
+        queue_size,
     };
 
     task_allocator.allocate()
@@ -255,7 +259,7 @@ impl<'a> TaskAllocationProcessor<'a> {
         let mut update = SchedulerUpdateRequest::default();
         let mut function_executors = self
             .in_memory_state
-            .candidate_function_executors(task, MAX_ALLOCATIONS_PER_FN_EXECUTOR)?;
+            .candidate_function_executors(task, self.queue_size)?;
         if function_executors.function_executors.is_empty() &&
             function_executors.num_pending_function_executors == 0
         {
@@ -269,7 +273,7 @@ impl<'a> TaskAllocationProcessor<'a> {
             )?;
             function_executors = self
                 .in_memory_state
-                .candidate_function_executors(task, MAX_ALLOCATIONS_PER_FN_EXECUTOR)?;
+                .candidate_function_executors(task, self.queue_size)?;
         }
         info!(
             num_function_executors = function_executors.function_executors.len(),

server/src/config.rs

@@ -17,6 +17,7 @@ pub struct ServerConfig {
     pub blob_storage: BlobStorageConfig,
     pub tracing: TracingConfig,
     pub executor: ExecutorConfig,
+    pub queue_size: u32,
 }
 
 impl Default for ServerConfig {
@@ -30,6 +31,7 @@ impl Default for ServerConfig {
             blob_storage: Default::default(),
             tracing: TracingConfig::default(),
             executor: ExecutorConfig::default(),
+            queue_size: 20,
         }
     }
 }

server/src/service.rs

@@ -80,6 +80,7 @@ impl Service {
         let graph_processor = Arc::new(GraphProcessor::new(
             indexify_state.clone(),
             task_cache.clone(),
+            config.queue_size,
         ));
         Ok(Self {
             config,

server/state_store/src/in_memory_state.rs

@@ -1017,7 +1017,7 @@ impl InMemoryState {
     pub fn candidate_function_executors(
         &self,
         task: &Task,
-        capacity_threshold: usize,
+        capacity_threshold: u32,
     ) -> Result<CandidateFunctionExecutors> {
         let mut candidates = Vec::new();
         let fn_uri = FunctionURI::from(task);
@@ -1043,7 +1043,7 @@ impl InMemoryState {
                     .and_then(|alloc_map| alloc_map.get(&function_executor.function_executor.id))
                     .map(|allocs| allocs.len())
                     .unwrap_or(0);
-                if allocation_count < capacity_threshold {
+                if (allocation_count as u32) < capacity_threshold {
                     candidates.push(function_executor.clone());
                 }
             }