Rename SlotsPolicy to TaskDistribution

Author: kyotoYaho

ballista/scheduler/scheduler_config_spec.toml

@@ -96,10 +96,10 @@ default = "3600"
 doc = "Delayed interval for cleaning up finished job state. Default: 3600"
 
 [[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"
+name = "task_distribution"
+type = "ballista_scheduler::config::TaskDistribution"
+doc = "The policy of distributing tasks to available executor slots, possible values: bias, round-robin. Default: bias"
+default = "ballista_scheduler::config::TaskDistribution::Bias"
 
 [[param]]
 name = "plugin_dir"

ballista/scheduler/src/bin/main.rs

@@ -109,7 +109,7 @@ async fn main() -> Result<()> {
         bind_port: opt.bind_port,
         scheduling_policy: opt.scheduler_policy,
         event_loop_buffer_size: opt.event_loop_buffer_size,
-        executor_slots_policy: opt.executor_slots_policy,
+        task_distribution: opt.task_distribution,
         finished_job_data_clean_up_interval_seconds: opt
             .finished_job_data_clean_up_interval_seconds,
         finished_job_state_clean_up_interval_seconds: opt

ballista/scheduler/src/cluster/mod.rs

@@ -29,7 +29,7 @@ use crate::cluster::memory::{InMemoryClusterState, InMemoryJobState};
 use crate::cluster::storage::etcd::EtcdClient;
 use crate::cluster::storage::sled::SledClient;
 use crate::cluster::storage::KeyValueStore;
-use crate::config::{ClusterStorageConfig, SchedulerConfig};
+use crate::config::{ClusterStorageConfig, SchedulerConfig, TaskDistribution};
 use crate::scheduler_server::SessionBuilder;
 use crate::state::execution_graph::ExecutionGraph;
 use crate::state::executor_manager::ExecutorReservation;
@@ -195,17 +195,6 @@ impl BallistaCluster {
 /// by any schedulers with a shared `ClusterState`
 pub type ExecutorHeartbeatStream = Pin<Box<dyn Stream<Item = ExecutorHeartbeat> + Send>>;
 
-/// Method of distributing tasks to available executor slots
-#[derive(Debug, Clone, Copy)]
-pub enum TaskDistribution {
-    /// Eagerly assign tasks to executor slots. This will assign as many task slots per executor
-    /// as are currently available
-    Bias,
-    /// Distributed tasks evenely across executors. This will try and iterate through available executors
-    /// and assign one task to each executor until all tasks are assigned.
-    RoundRobin,
-}
-
 /// A trait that contains the necessary method to maintain a globally consistent view of cluster resources
 #[tonic::async_trait]
 pub trait ClusterState: Send + Sync + 'static {

ballista/scheduler/src/config.rs

@@ -26,7 +26,7 @@ use std::fmt;
 #[derive(Debug, Clone)]
 pub struct SchedulerConfig {
     /// Namespace of this scheduler. Schedulers using the same cluster storage and namespace
-    /// will share gloabl cluster state.
+    /// will share global cluster state.
     pub namespace: String,
     /// The external hostname of the scheduler
     pub external_host: String,
@@ -36,8 +36,8 @@ pub struct SchedulerConfig {
     pub scheduling_policy: TaskSchedulingPolicy,
     /// The event loop buffer size. for a system of high throughput, a larger value like 1000000 is recommended
     pub event_loop_buffer_size: u32,
-    /// The executor slots policy for the scheduler. For a cluster with single scheduler, round-robin-local is recommended
-    pub executor_slots_policy: SlotsPolicy,
+    /// Policy of distributing tasks to available executor slots. For a cluster with single scheduler, round-robin is recommended
+    pub task_distribution: TaskDistribution,
     /// The delayed interval for cleaning up finished job data, mainly the shuffle data, 0 means the cleaning up is disabled
     pub finished_job_data_clean_up_interval_seconds: u64,
     /// The delayed interval for cleaning up finished job state stored in the backend, 0 means the cleaning up is disabled.
@@ -62,7 +62,7 @@ impl Default for SchedulerConfig {
             bind_port: 50050,
             scheduling_policy: TaskSchedulingPolicy::PullStaged,
             event_loop_buffer_size: 10000,
-            executor_slots_policy: SlotsPolicy::Bias,
+            task_distribution: TaskDistribution::Bias,
             finished_job_data_clean_up_interval_seconds: 300,
             finished_job_state_clean_up_interval_seconds: 3600,
             advertise_flight_sql_endpoint: None,
@@ -131,8 +131,8 @@ impl SchedulerConfig {
         self
     }
 
-    pub fn with_executor_slots_policy(mut self, policy: SlotsPolicy) -> Self {
-        self.executor_slots_policy = policy;
+    pub fn with_task_distribution(mut self, policy: TaskDistribution) -> Self {
+        self.task_distribution = policy;
         self
     }
 
@@ -161,23 +161,28 @@ pub enum ClusterStorageConfig {
     Sled(Option<String>),
 }
 
-// an enum used to configure the executor slots policy
-// needs to be visible to code generated by configure_me
+/// Policy of distributing tasks to available executor slots
+///
+/// It needs to be visible to code generated by configure_me
 #[derive(Clone, ArgEnum, Copy, Debug, serde::Deserialize)]
-pub enum SlotsPolicy {
+pub enum TaskDistribution {
+    /// Eagerly assign tasks to executor slots. This will assign as many task slots per executor
+    /// as are currently available
     Bias,
+    /// Distributed tasks evenly across executors. This will try and iterate through available executors
+    /// and assign one task to each executor until all tasks are assigned.
     RoundRobin,
 }
 
-impl std::str::FromStr for SlotsPolicy {
+impl std::str::FromStr for TaskDistribution {
     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 {
+impl parse_arg::ParseArgFromStr for TaskDistribution {
     fn describe_type<W: fmt::Write>(mut writer: W) -> fmt::Result {
         write!(writer, "The executor slots policy for the scheduler")
     }

ballista/scheduler/src/state/executor_manager.rs

@@ -17,15 +17,13 @@
 
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
-use crate::cluster::TaskDistribution;
-
 #[cfg(not(test))]
 use ballista_core::error::BallistaError;
 use ballista_core::error::Result;
 use ballista_core::serde::protobuf;
 
 use crate::cluster::ClusterState;
-use crate::config::SlotsPolicy;
+use crate::config::TaskDistribution;
 
 use crate::state::execution_graph::RunningTaskInfo;
 use ballista_core::serde::protobuf::executor_grpc_client::ExecutorGrpcClient;
@@ -101,13 +99,8 @@ pub struct ExecutorManager {
 impl ExecutorManager {
     pub(crate) fn new(
         cluster_state: Arc<dyn ClusterState>,
-        slots_policy: SlotsPolicy,
+        task_distribution: TaskDistribution,
     ) -> Self {
-        let task_distribution = match slots_policy {
-            SlotsPolicy::Bias => TaskDistribution::Bias,
-            SlotsPolicy::RoundRobin => TaskDistribution::RoundRobin,
-        };
-
         Self {
             task_distribution,
             cluster_state,
@@ -511,7 +504,7 @@ impl ExecutorManager {
 #[cfg(test)]
 mod test {
 
-    use crate::config::SlotsPolicy;
+    use crate::config::TaskDistribution;
 
     use crate::scheduler_server::timestamp_secs;
     use crate::state::executor_manager::{ExecutorManager, ExecutorReservation};
@@ -525,17 +518,19 @@ mod test {
 
     #[tokio::test]
     async fn test_reserve_and_cancel() -> Result<()> {
-        test_reserve_and_cancel_inner(SlotsPolicy::Bias).await?;
-        test_reserve_and_cancel_inner(SlotsPolicy::RoundRobin).await?;
+        test_reserve_and_cancel_inner(TaskDistribution::Bias).await?;
+        test_reserve_and_cancel_inner(TaskDistribution::RoundRobin).await?;
 
         Ok(())
     }
 
-    async fn test_reserve_and_cancel_inner(slots_policy: SlotsPolicy) -> Result<()> {
+    async fn test_reserve_and_cancel_inner(
+        task_distribution: TaskDistribution,
+    ) -> Result<()> {
         let cluster = test_cluster_context();
 
         let executor_manager =
-            ExecutorManager::new(cluster.cluster_state(), slots_policy);
+            ExecutorManager::new(cluster.cluster_state(), task_distribution);
 
         let executors = test_executors(10, 4);
 
@@ -551,7 +546,7 @@ mod test {
         assert_eq!(
             reservations.len(),
             40,
-            "Expected 40 reservations for policy {slots_policy:?}"
+            "Expected 40 reservations for policy {task_distribution:?}"
         );
 
         // Now cancel them
@@ -563,25 +558,27 @@ mod test {
         assert_eq!(
             reservations.len(),
             40,
-            "Expected 40 reservations for policy {slots_policy:?}"
+            "Expected 40 reservations for policy {task_distribution:?}"
         );
 
         Ok(())
     }
 
     #[tokio::test]
     async fn test_reserve_partial() -> Result<()> {
-        test_reserve_partial_inner(SlotsPolicy::Bias).await?;
-        test_reserve_partial_inner(SlotsPolicy::RoundRobin).await?;
+        test_reserve_partial_inner(TaskDistribution::Bias).await?;
+        test_reserve_partial_inner(TaskDistribution::RoundRobin).await?;
 
         Ok(())
     }
 
-    async fn test_reserve_partial_inner(slots_policy: SlotsPolicy) -> Result<()> {
+    async fn test_reserve_partial_inner(
+        task_distribution: TaskDistribution,
+    ) -> Result<()> {
         let cluster = test_cluster_context();
 
         let executor_manager =
-            ExecutorManager::new(cluster.cluster_state(), slots_policy);
+            ExecutorManager::new(cluster.cluster_state(), task_distribution);
 
         let executors = test_executors(10, 4);
 
@@ -621,21 +618,23 @@ mod test {
 
     #[tokio::test]
     async fn test_reserve_concurrent() -> Result<()> {
-        test_reserve_concurrent_inner(SlotsPolicy::Bias).await?;
-        test_reserve_concurrent_inner(SlotsPolicy::RoundRobin).await?;
+        test_reserve_concurrent_inner(TaskDistribution::Bias).await?;
+        test_reserve_concurrent_inner(TaskDistribution::RoundRobin).await?;
 
         Ok(())
     }
 
-    async fn test_reserve_concurrent_inner(slots_policy: SlotsPolicy) -> Result<()> {
+    async fn test_reserve_concurrent_inner(
+        task_distribution: TaskDistribution,
+    ) -> Result<()> {
         let (sender, mut receiver) =
             tokio::sync::mpsc::channel::<Result<Vec<ExecutorReservation>>>(1000);
 
         let executors = test_executors(10, 4);
 
         let cluster = test_cluster_context();
         let executor_manager =
-            ExecutorManager::new(cluster.cluster_state(), slots_policy);
+            ExecutorManager::new(cluster.cluster_state(), task_distribution);
 
         for (executor_metadata, executor_data) in executors {
             executor_manager
@@ -670,17 +669,19 @@ mod test {
 
     #[tokio::test]
     async fn test_register_reserve() -> Result<()> {
-        test_register_reserve_inner(SlotsPolicy::Bias).await?;
-        test_register_reserve_inner(SlotsPolicy::RoundRobin).await?;
+        test_register_reserve_inner(TaskDistribution::Bias).await?;
+        test_register_reserve_inner(TaskDistribution::RoundRobin).await?;
 
         Ok(())
     }
 
-    async fn test_register_reserve_inner(slots_policy: SlotsPolicy) -> Result<()> {
+    async fn test_register_reserve_inner(
+        task_distribution: TaskDistribution,
+    ) -> Result<()> {
         let cluster = test_cluster_context();
 
         let executor_manager =
-            ExecutorManager::new(cluster.cluster_state(), slots_policy);
+            ExecutorManager::new(cluster.cluster_state(), task_distribution);
 
         let executors = test_executors(10, 4);
 
@@ -702,17 +703,19 @@ mod test {
 
     #[tokio::test]
     async fn test_ignore_fenced_executors() -> Result<()> {
-        test_ignore_fenced_executors_inner(SlotsPolicy::Bias).await?;
-        test_ignore_fenced_executors_inner(SlotsPolicy::RoundRobin).await?;
+        test_ignore_fenced_executors_inner(TaskDistribution::Bias).await?;
+        test_ignore_fenced_executors_inner(TaskDistribution::RoundRobin).await?;
 
         Ok(())
     }
 
-    async fn test_ignore_fenced_executors_inner(slots_policy: SlotsPolicy) -> Result<()> {
+    async fn test_ignore_fenced_executors_inner(
+        task_distribution: TaskDistribution,
+    ) -> Result<()> {
         let cluster = test_cluster_context();
 
         let executor_manager =
-            ExecutorManager::new(cluster.cluster_state(), slots_policy);
+            ExecutorManager::new(cluster.cluster_state(), task_distribution);
 
         // Setup two executors initially
         let executors = test_executors(2, 4);

ballista/scheduler/src/state/mod.rs

@@ -116,7 +116,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
         Self {
             executor_manager: ExecutorManager::new(
                 cluster.cluster_state(),
-                config.executor_slots_policy,
+                config.task_distribution,
             ),
             task_manager: TaskManager::new(
                 cluster.job_state(),
@@ -140,7 +140,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerState<T,
         Self {
             executor_manager: ExecutorManager::new(
                 cluster.cluster_state(),
-                config.executor_slots_policy,
+                config.task_distribution,
             ),
             task_manager: TaskManager::with_launcher(
                 cluster.job_state(),