Add round robin executor slots reservation policy for the scheduler to evenly assign tasks to executors

ballista/scheduler/scheduler_config_spec.toml

@@ -72,6 +72,12 @@ type = "ballista_core::config::TaskSchedulingPolicy"
 doc = "The scheduing policy for the scheduler, possible values: pull-staged, push-staged. Default: pull-staged"
 default = "ballista_core::config::TaskSchedulingPolicy::PullStaged"
 
+[[param]]
+name = "executor_slots_policy"
+type = "ballista_scheduler::config::SlotsPolicy"
+doc = "The executor slots policy for the scheduler, possible values: bias, round-robin. Default: bias"
+default = "ballista_scheduler::config::SlotsPolicy::Bias"
+
 [[param]]
 name = "plugin_dir"
 type = "String"

ballista/scheduler/src/config.rs

@@ -0,0 +1,44 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+//
+
+//! Ballista scheduler specific configuration
+
+use clap::ArgEnum;
+use std::fmt;
+
+// an enum used to configure the executor slots policy
+// needs to be visible to code generated by configure_me
+#[derive(Clone, ArgEnum, Copy, Debug, serde::Deserialize)]
+pub enum SlotsPolicy {
+    Bias,
+    RoundRobin,
+}
+
+impl std::str::FromStr for SlotsPolicy {
+    type Err = String;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        ArgEnum::from_str(s, true)
+    }
+}
+
+impl parse_arg::ParseArgFromStr for SlotsPolicy {
+    fn describe_type<W: fmt::Write>(mut writer: W) -> fmt::Result {
+        write!(writer, "The executor slots policy for the scheduler")
+    }
+}

ballista/scheduler/src/lib.rs

@@ -18,6 +18,7 @@
 #![doc = include_str ! ("../README.md")]
 
 pub mod api;
+pub mod config;
 pub mod display;
 pub mod planner;
 pub mod scheduler_server;

ballista/scheduler/src/main.rs

@@ -62,6 +62,7 @@ mod config {
 }
 
 use ballista_core::utils::create_grpc_server;
+use ballista_scheduler::config::SlotsPolicy;
 #[cfg(feature = "flight-sql")]
 use ballista_scheduler::flight_sql::FlightSqlServiceImpl;
 use config::prelude::*;
@@ -71,7 +72,8 @@ async fn start_server(
     scheduler_name: String,
     config_backend: Arc<dyn StateBackendClient>,
     addr: SocketAddr,
-    policy: TaskSchedulingPolicy,
+    scheduling_policy: TaskSchedulingPolicy,
+    slots_policy: SlotsPolicy,
 ) -> Result<()> {
     info!(
         "Ballista v{} Scheduler listening on {:?}",
@@ -80,14 +82,15 @@ async fn start_server(
     // Should only call SchedulerServer::new() once in the process
     info!(
         "Starting Scheduler grpc server with task scheduling policy of {:?}",
-        policy
+        scheduling_policy
     );
     let mut scheduler_server: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
-        match policy {
+        match scheduling_policy {
             TaskSchedulingPolicy::PushStaged => SchedulerServer::new_with_policy(
                 scheduler_name,
                 config_backend.clone(),
-                policy,
+                scheduling_policy,
+                slots_policy,
                 BallistaCodec::default(),
                 default_session_builder,
             ),
@@ -239,7 +242,15 @@ async fn main() -> Result<()> {
         }
     };
 
-    let policy: TaskSchedulingPolicy = opt.scheduler_policy;
-    start_server(scheduler_name, client, addr, policy).await?;
+    let scheduling_policy: TaskSchedulingPolicy = opt.scheduler_policy;
+    let slots_policy: SlotsPolicy = opt.executor_slots_policy;
+    start_server(
+        scheduler_name,
+        client,
+        addr,
+        scheduling_policy,
+        slots_policy,
+    )
+    .await?;
     Ok(())
 }

ballista/scheduler/src/scheduler_server/mod.rs

@@ -30,6 +30,7 @@ use datafusion::logical_plan::LogicalPlan;
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_proto::logical_plan::AsLogicalPlan;
 
+use crate::config::SlotsPolicy;
 use log::{error, warn};
 
 use crate::scheduler_server::event::QueryStageSchedulerEvent;
@@ -72,6 +73,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             scheduler_name,
             config,
             TaskSchedulingPolicy::PullStaged,
+            SlotsPolicy::Bias,
             codec,
             default_session_builder,
         )
@@ -87,6 +89,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             scheduler_name,
             config,
             TaskSchedulingPolicy::PullStaged,
+            SlotsPolicy::Bias,
             codec,
             session_builder,
         )
@@ -95,7 +98,8 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     pub fn new_with_policy(
         scheduler_name: String,
         config: Arc<dyn StateBackendClient>,
-        policy: TaskSchedulingPolicy,
+        scheduling_policy: TaskSchedulingPolicy,
+        slots_policy: SlotsPolicy,
         codec: BallistaCodec<T, U>,
         session_builder: SessionBuilder,
     ) -> Self {
@@ -104,9 +108,10 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             session_builder,
             codec,
             scheduler_name.clone(),
+            slots_policy,
         ));
 
-        SchedulerServer::new_with_state(scheduler_name, policy, state)
+        SchedulerServer::new_with_state(scheduler_name, scheduling_policy, state)
     }
 
     pub(crate) fn new_with_state(
@@ -294,6 +299,7 @@ mod test {
     };
     use ballista_core::error::Result;
 
+    use crate::config::SlotsPolicy;
     use ballista_core::serde::protobuf::{
         failed_task, job_status, task_status, ExecutionError, FailedTask, JobStatus,
         PhysicalPlanNode, ShuffleWritePartition, SuccessfulTask, TaskStatus,
@@ -753,14 +759,15 @@ mod test {
     }
 
     async fn test_scheduler(
-        policy: TaskSchedulingPolicy,
+        scheduling_policy: TaskSchedulingPolicy,
     ) -> Result<SchedulerServer<LogicalPlanNode, PhysicalPlanNode>> {
         let state_storage = Arc::new(StandaloneClient::try_new_temporary()?);
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new_with_policy(
                 "localhost:50050".to_owned(),
                 state_storage.clone(),
-                policy,
+                scheduling_policy,
+                SlotsPolicy::Bias,
                 BallistaCodec::default(),
                 default_session_builder,
             );