Cluster state refactor Part 2 (#658)

* Customize session builder

* Add setter for executor slots policy

* Construct Executor with functions

* Add queued and completed timestamps to successful job status

* Add public methods to SchedulerServer

* Public method for getting execution graph

* Public method for stage metrics

* Use node-level local limit (#20)

* Use node-level local limit

* serialize limit in shuffle writer

* Revert "Merge pull request #19 from coralogix/sc-5792"

This reverts commit 08140ef, reversing
changes made to a7f1384.

* add log

* make sure we don't forget limit for shuffle writer

* update accum correctly and try to break early

* Check local limit accumulator before polling for more data

* fix build

Co-authored-by: Martins Purins <[email protected]>

* configure_me_codegen retroactively reserved on our `bind_host` parame… (#520)

* configure_me_codegen retroactively reserved on our `bind_host` parameter name

* Add label and pray

* Add more labels why not

* Add ClusterState trait

* Refactor slightly for clarity

* Revert "Use node-level local limit (#20)"

This reverts commit ff96bcd.

* Revert "Public method for stage metrics"

This reverts commit a802315.

* Revert "Public method for getting execution graph"

This reverts commit 490bda5.

* Revert "Add public methods to SchedulerServer"

This reverts commit 5ad27c0.

* Revert "Add queued and completed timestamps to successful job status"

This reverts commit c615fce.

* Revert "Construct Executor with functions"

This reverts commit 24d4830.

* Always forget the apache header

* WIP

* Implement JobState

* Tests and fixes

* do not hold ref across await point

* Fix clippy warnings

* Fix tomlfmt github action

* uncomment test

---------

Co-authored-by: Martins Purins <[email protected]>
Co-authored-by: Brent Gardner <[email protected]>

Author: 3 people

.github/workflows/rust.yml

@@ -402,7 +402,10 @@ jobs:
           #
           # ignore ./Cargo.toml because putting workspaces in multi-line lists make it easy to read
           ci/scripts/rust_toml_fmt.sh
-          git diff --exit-code
+          if test -f "./Cargo.toml.bak"; then
+              echo "cargo tomlfmt found format violations"
+              exit 1
+          fi
         env:
           CARGO_HOME: "/github/home/.cargo"
           CARGO_TARGET_DIR: "/github/home/target"

ballista/core/proto/ballista.proto

@@ -289,6 +289,7 @@ message ExecutorMetadata {
   ExecutorSpecification specification = 5;
 }
 
+
 // Used by grpc
 message ExecutorRegistration {
   string id = 1;
@@ -336,6 +337,15 @@ message ExecutorResource {
   }
 }
 
+message AvailableTaskSlots {
+  string executor_id = 1;
+  uint32 slots = 2;
+    }
+
+message ExecutorTaskSlots {
+  repeated AvailableTaskSlots task_slots = 1;
+}
+
 message ExecutorData {
   string executor_id = 1;
   repeated ExecutorResourcePair resources = 2;
@@ -544,18 +554,33 @@ message GetJobStatusParams {
 
 message SuccessfulJob {
   repeated PartitionLocation partition_location = 1;
+  uint64 queued_at = 2;
+  uint64 started_at = 3;
+  uint64 ended_at = 4;
 }
 
-message QueuedJob {}
+message QueuedJob {
+  uint64 queued_at = 1;
+}
 
 // TODO: add progress report
-message RunningJob {}
+message RunningJob {
+  uint64 queued_at = 1;
+  uint64 started_at = 2;
+  string scheduler = 3;
+}
 
 message FailedJob {
   string error = 1;
+  uint64 queued_at = 2;
+  uint64 started_at = 3;
+  uint64 ended_at = 4;
 }
 
 message JobStatus {
+  string job_id = 5;
+  string job_name = 6;
+
   oneof status {
     QueuedJob queued = 1;
     RunningJob running = 2;

ballista/core/src/serde/generated/ballista.rs

@@ -581,6 +581,20 @@ pub mod executor_resource {
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
+pub struct AvailableTaskSlots {
+    #[prost(string, tag = "1")]
+    pub executor_id: ::prost::alloc::string::String,
+    #[prost(uint32, tag = "2")]
+    pub slots: u32,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct ExecutorTaskSlots {
+    #[prost(message, repeated, tag = "1")]
+    pub task_slots: ::prost::alloc::vec::Vec<AvailableTaskSlots>,
+}
+#[allow(clippy::derive_partial_eq_without_eq)]
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ExecutorData {
     #[prost(string, tag = "1")]
     pub executor_id: ::prost::alloc::string::String,
@@ -933,23 +947,49 @@ pub struct GetJobStatusParams {
 pub struct SuccessfulJob {
     #[prost(message, repeated, tag = "1")]
     pub partition_location: ::prost::alloc::vec::Vec<PartitionLocation>,
+    #[prost(uint64, tag = "2")]
+    pub queued_at: u64,
+    #[prost(uint64, tag = "3")]
+    pub started_at: u64,
+    #[prost(uint64, tag = "4")]
+    pub ended_at: u64,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct QueuedJob {}
+pub struct QueuedJob {
+    #[prost(uint64, tag = "1")]
+    pub queued_at: u64,
+}
 /// TODO: add progress report
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
-pub struct RunningJob {}
+pub struct RunningJob {
+    #[prost(uint64, tag = "1")]
+    pub queued_at: u64,
+    #[prost(uint64, tag = "2")]
+    pub started_at: u64,
+    #[prost(string, tag = "3")]
+    pub scheduler: ::prost::alloc::string::String,
+}
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct FailedJob {
     #[prost(string, tag = "1")]
     pub error: ::prost::alloc::string::String,
+    #[prost(uint64, tag = "2")]
+    pub queued_at: u64,
+    #[prost(uint64, tag = "3")]
+    pub started_at: u64,
+    #[prost(uint64, tag = "4")]
+    pub ended_at: u64,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct JobStatus {
+    #[prost(string, tag = "5")]
+    pub job_id: ::prost::alloc::string::String,
+    #[prost(string, tag = "6")]
+    pub job_name: ::prost::alloc::string::String,
     #[prost(oneof = "job_status::Status", tags = "1, 2, 3, 4")]
     pub status: ::core::option::Option<job_status::Status>,
 }

ballista/core/src/serde/mod.rs

@@ -26,11 +26,13 @@ use datafusion::execution::FunctionRegistry;
 use datafusion::physical_plan::{ExecutionPlan, Partitioning};
 use datafusion_proto::common::proto_error;
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning;
+use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
 use datafusion_proto::{
     convert_required,
     logical_plan::{AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec},
     physical_plan::{AsExecutionPlan, PhysicalExtensionCodec},
 };
+
 use prost::Message;
 use std::fmt::Debug;
 use std::marker::PhantomData;
@@ -69,16 +71,17 @@ pub fn decode_protobuf(bytes: &[u8]) -> Result<BallistaAction, BallistaError> {
 }
 
 #[derive(Clone, Debug)]
-pub struct BallistaCodec<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {
+pub struct BallistaCodec<
+    T: 'static + AsLogicalPlan = LogicalPlanNode,
+    U: 'static + AsExecutionPlan = PhysicalPlanNode,
+> {
     logical_extension_codec: Arc<dyn LogicalExtensionCodec>,
     physical_extension_codec: Arc<dyn PhysicalExtensionCodec>,
     logical_plan_repr: PhantomData<T>,
     physical_plan_repr: PhantomData<U>,
 }
 
-impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> Default
-    for BallistaCodec<T, U>
-{
+impl Default for BallistaCodec {
     fn default() -> Self {
         Self {
             logical_extension_codec: Arc::new(DefaultLogicalExtensionCodec {}),

ballista/scheduler/scheduler_config_spec.toml

@@ -31,17 +31,10 @@ doc = "Route for proxying flight results via scheduler. Should be of the form 'I
 
 [[param]]
 abbr = "b"
-name = "config_backend"
-type = "ballista_scheduler::state::backend::StateBackend"
-doc = "The configuration backend for the scheduler, possible values: etcd, memory, sled. Default: sled"
-default = "ballista_scheduler::state::backend::StateBackend::Sled"
-
-[[param]]
-abbr = "c"
 name = "cluster_backend"
-type = "ballista_scheduler::state::backend::StateBackend"
+type = "ballista_scheduler::cluster::ClusterStorage"
 doc = "The configuration backend for the scheduler cluster state, possible values: etcd, memory, sled. Default: sled"
-default = "ballista_scheduler::state::backend::StateBackend::Sled"
+default = "ballista_scheduler::cluster::ClusterStorage::Sled"
 
 [[param]]
 abbr = "n"

ballista/scheduler/src/api/handlers.rs

@@ -209,6 +209,7 @@ pub(crate) async fn get_query_stages<T: AsLogicalPlan, U: AsExecutionPlan>(
     {
         Ok(warp::reply::json(&QueryStagesResponse {
             stages: graph
+                .as_ref()
                 .stages()
                 .iter()
                 .map(|(id, stage)| {
@@ -303,7 +304,7 @@ pub(crate) async fn get_job_dot_graph<T: AsLogicalPlan, U: AsExecutionPlan>(
         .await
         .map_err(|_| warp::reject())?
     {
-        ExecutionGraphDot::generate(graph).map_err(|_| warp::reject())
+        ExecutionGraphDot::generate(graph.as_ref()).map_err(|_| warp::reject())
     } else {
         Ok("Not Found".to_string())
     }
@@ -322,7 +323,7 @@ pub(crate) async fn get_query_stage_dot_graph<T: AsLogicalPlan, U: AsExecutionPl
         .await
         .map_err(|_| warp::reject())?
     {
-        ExecutionGraphDot::generate_for_query_stage(graph, stage_id)
+        ExecutionGraphDot::generate_for_query_stage(graph.as_ref(), stage_id)
             .map_err(|_| warp::reject())
     } else {
         Ok("Not Found".to_string())

ballista/scheduler/src/bin/main.rs

@@ -17,18 +17,18 @@
 
 //! Ballista Rust scheduler binary.
 
-use std::{env, io, sync::Arc};
+use std::{env, io};
 
-use anyhow::{Context, Result};
+use anyhow::Result;
 
 use ballista_core::print_version;
 use ballista_scheduler::scheduler_process::start_server;
-#[cfg(feature = "etcd")]
-use ballista_scheduler::state::backend::etcd::EtcdClient;
-use ballista_scheduler::state::backend::memory::MemoryBackendClient;
-#[cfg(feature = "sled")]
-use ballista_scheduler::state::backend::sled::SledClient;
-use ballista_scheduler::state::backend::{StateBackend, StateBackendClient};
+
+use crate::config::{Config, ResultExt};
+use ballista_core::config::LogRotationPolicy;
+use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::config::{ClusterStorageConfig, SchedulerConfig};
+use tracing_subscriber::EnvFilter;
 
 #[macro_use]
 extern crate configure_me;
@@ -43,12 +43,6 @@ mod config {
     ));
 }
 
-use ballista_core::config::LogRotationPolicy;
-use ballista_scheduler::config::SchedulerConfig;
-use ballista_scheduler::state::backend::cluster::DefaultClusterState;
-use config::prelude::*;
-use tracing_subscriber::EnvFilter;
-
 #[tokio::main]
 async fn main() -> Result<()> {
     // parse options
@@ -61,25 +55,14 @@ async fn main() -> Result<()> {
         std::process::exit(0);
     }
 
-    let config_backend = init_kv_backend(&opt.config_backend, &opt).await?;
-
-    let cluster_state = if opt.cluster_backend == opt.config_backend {
-        Arc::new(DefaultClusterState::new(config_backend.clone()))
-    } else {
-        let cluster_kv_store = init_kv_backend(&opt.cluster_backend, &opt).await?;
-
-        Arc::new(DefaultClusterState::new(cluster_kv_store))
-    };
-
     let special_mod_log_level = opt.log_level_setting;
-    let namespace = opt.namespace;
-    let external_host = opt.external_host;
-    let bind_host = opt.bind_host;
-    let port = opt.bind_port;
     let log_dir = opt.log_dir;
     let print_thread_info = opt.print_thread_info;
-    let log_file_name_prefix = format!("scheduler_{namespace}_{external_host}_{port}");
-    let scheduler_name = format!("{external_host}:{port}");
+
+    let log_file_name_prefix = format!(
+        "scheduler_{}_{}_{}",
+        opt.namespace, opt.external_host, opt.bind_port
+    );
 
     let rust_log = env::var(EnvFilter::DEFAULT_ENV);
     let log_filter = EnvFilter::new(rust_log.unwrap_or(special_mod_log_level));
@@ -117,10 +100,13 @@ async fn main() -> Result<()> {
             .init();
     }
 
-    let addr = format!("{bind_host}:{port}");
+    let addr = format!("{}:{}", opt.bind_host, opt.bind_port);
     let addr = addr.parse()?;
 
     let config = SchedulerConfig {
+        namespace: opt.namespace,
+        external_host: opt.external_host,
+        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,
@@ -129,54 +115,13 @@ async fn main() -> Result<()> {
         finished_job_state_clean_up_interval_seconds: opt
             .finished_job_state_clean_up_interval_seconds,
         advertise_flight_sql_endpoint: opt.advertise_flight_sql_endpoint,
+        cluster_storage: ClusterStorageConfig::Memory,
         job_resubmit_interval_ms: (opt.job_resubmit_interval_ms > 0)
             .then_some(opt.job_resubmit_interval_ms),
     };
-    start_server(scheduler_name, config_backend, cluster_state, addr, config).await?;
-    Ok(())
-}
 
-async fn init_kv_backend(
-    backend: &StateBackend,
-    opt: &Config,
-) -> Result<Arc<dyn StateBackendClient>> {
-    let cluster_backend: Arc<dyn StateBackendClient> = match backend {
-        #[cfg(feature = "etcd")]
-        StateBackend::Etcd => {
-            let etcd = etcd_client::Client::connect(&[opt.etcd_urls.clone()], None)
-                .await
-                .context("Could not connect to etcd")?;
-            Arc::new(EtcdClient::new(opt.namespace.clone(), etcd))
-        }
-        #[cfg(not(feature = "etcd"))]
-        StateBackend::Etcd => {
-            unimplemented!(
-                "build the scheduler with the `etcd` feature to use the etcd config backend"
-            )
-        }
-        #[cfg(feature = "sled")]
-        StateBackend::Sled => {
-            if opt.sled_dir.is_empty() {
-                Arc::new(
-                    SledClient::try_new_temporary()
-                        .context("Could not create sled config backend")?,
-                )
-            } else {
-                println!("{}", opt.sled_dir);
-                Arc::new(
-                    SledClient::try_new(opt.sled_dir.clone())
-                        .context("Could not create sled config backend")?,
-                )
-            }
-        }
-        #[cfg(not(feature = "sled"))]
-        StateBackend::Sled => {
-            unimplemented!(
-                "build the scheduler with the `sled` feature to use the sled config backend"
-            )
-        }
-        StateBackend::Memory => Arc::new(MemoryBackendClient::new()),
-    };
+    let cluster = BallistaCluster::new_from_config(&config).await?;
 
-    Ok(cluster_backend)
+    start_server(cluster, addr, config).await?;
+    Ok(())
 }