From b1e58fb5c7ddb094252a5f5601767e0c10e2142d Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 5 Aug 2022 10:39:06 +0800 Subject: [PATCH 01/10] Add timeout and keep-alive settings for Grpc Client --- ballista/rust/client/src/context.rs | 7 +++++-- ballista/rust/core/src/client.rs | 6 ++++-- .../src/execution_plans/distributed_query.rs | 6 ++++-- .../src/execution_plans/shuffle_reader.rs | 2 ++ ballista/rust/core/src/utils.rs | 20 +++++++++++++++++++ ballista/rust/executor/src/main.rs | 5 ++++- .../rust/scheduler/src/state/task_manager.rs | 3 ++- 7 files changed, 41 insertions(+), 8 deletions(-) diff --git a/ballista/rust/client/src/context.rs b/ballista/rust/client/src/context.rs index 03cd0799b2..32e1f66175 100644 --- a/ballista/rust/client/src/context.rs +++ b/ballista/rust/client/src/context.rs @@ -28,7 +28,9 @@ use std::sync::Arc; use ballista_core::config::BallistaConfig; use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient; use ballista_core::serde::protobuf::{ExecuteQueryParams, KeyValuePair}; -use ballista_core::utils::create_df_ctx_with_ballista_query_planner; +use ballista_core::utils::{ + create_df_ctx_with_ballista_query_planner, create_grpc_client_connection, +}; use datafusion_proto::protobuf::LogicalPlanNode; use datafusion::catalog::TableReference; @@ -93,9 +95,10 @@ impl BallistaContext { "Connecting to Ballista scheduler at {}", scheduler_url.clone() ); - let mut scheduler = SchedulerGrpcClient::connect(scheduler_url.clone()) + let connection = create_grpc_client_connection(scheduler_url.clone()) .await .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; + let mut scheduler = SchedulerGrpcClient::new(connection); let remote_session_id = scheduler .execute_query(ExecuteQueryParams { diff --git a/ballista/rust/core/src/client.rs b/ballista/rust/core/src/client.rs index a5c4a062be..dfe2003fb2 100644 --- a/ballista/rust/core/src/client.rs +++ b/ballista/rust/core/src/client.rs @@ -39,6 +39,7 @@ use datafusion::arrow::{ record_batch::RecordBatch, }; +use crate::utils::create_grpc_client_connection; use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream}; use futures::{Stream, StreamExt}; use log::debug; @@ -57,8 +58,8 @@ impl BallistaClient { pub async fn try_new(host: &str, port: u16) -> Result { let addr = format!("http://{}:{}", host, port); debug!("BallistaClient connecting to {}", addr); - let flight_client = - FlightServiceClient::connect(addr.clone()) + let connection = + create_grpc_client_connection(addr.clone()) .await .map_err(|e| { BallistaError::General(format!( @@ -66,6 +67,7 @@ impl BallistaClient { addr, e )) })?; + let flight_client = FlightServiceClient::new(connection); debug!("BallistaClient connected OK"); Ok(Self { flight_client }) diff --git a/ballista/rust/core/src/execution_plans/distributed_query.rs b/ballista/rust/core/src/execution_plans/distributed_query.rs index 62e7ff0202..11666a8e2b 100644 --- a/ballista/rust/core/src/execution_plans/distributed_query.rs +++ b/ballista/rust/core/src/execution_plans/distributed_query.rs @@ -23,6 +23,7 @@ use crate::serde::protobuf::{ ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult, KeyValuePair, PartitionLocation, }; +use crate::utils::create_grpc_client_connection; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::error::{ArrowError, Result as ArrowResult}; use datafusion::arrow::record_batch::RecordBatch; @@ -235,11 +236,12 @@ async fn execute_query( ) -> Result> + Send> { info!("Connecting to Ballista scheduler at {}", scheduler_url); // TODO reuse the scheduler to avoid connecting to the Ballista scheduler again and again - - let mut scheduler = SchedulerGrpcClient::connect(scheduler_url.clone()) + let connection = create_grpc_client_connection(scheduler_url) .await .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; + let mut scheduler = SchedulerGrpcClient::new(connection); + let query_result = scheduler .execute_query(query) .await diff --git a/ballista/rust/core/src/execution_plans/shuffle_reader.rs b/ballista/rust/core/src/execution_plans/shuffle_reader.rs index 3046a2276f..c69d120b1c 100644 --- a/ballista/rust/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/rust/core/src/execution_plans/shuffle_reader.rs @@ -205,6 +205,8 @@ async fn fetch_partition( ) -> Result { let metadata = &location.executor_meta; let partition_id = &location.partition_id; + // TODO for shuffle client connections, we should avoid creating new connections again and again. + // And we should also avoid to keep alive too many connections for long time. let mut ballista_client = BallistaClient::try_new(metadata.host.as_str(), metadata.port as u16) .await diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs index 7493b1ee93..c2f465c763 100644 --- a/ballista/rust/core/src/utils.rs +++ b/ballista/rust/core/src/utils.rs @@ -49,7 +49,10 @@ use std::io::{BufWriter, Write}; use std::marker::PhantomData; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; +use std::time::Duration; use std::{fs::File, pin::Pin}; +use tonic::codegen::StdError; +use tonic::transport::{Channel, Error}; /// Stream data to disk in Arrow IPC format @@ -310,3 +313,20 @@ impl QueryPlanner for BallistaQueryPlanner { } } } + +pub async fn create_grpc_client_connection( + dst: D, +) -> std::result::Result +where + D: std::convert::TryInto, + D::Error: Into, +{ + let endpoint = tonic::transport::Endpoint::new(dst)? + .connect_timeout(Duration::from_secs(20)) + .timeout(Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Duration::from_secs(3600))) + .http2_keep_alive_interval(Duration::from_secs(300)) + .keep_alive_timeout(Duration::from_secs(20)) + .keep_alive_while_idle(true); + endpoint.connect().await +} diff --git a/ballista/rust/executor/src/main.rs b/ballista/rust/executor/src/main.rs index 6a87be7766..01048936a5 100644 --- a/ballista/rust/executor/src/main.rs +++ b/ballista/rust/executor/src/main.rs @@ -39,6 +39,7 @@ use ballista_core::serde::protobuf::{ }; use ballista_core::serde::scheduler::ExecutorSpecification; use ballista_core::serde::BallistaCodec; +use ballista_core::utils::create_grpc_client_connection; use ballista_core::{print_version, BALLISTA_VERSION}; use ballista_executor::executor::Executor; use ballista_executor::flight_service::BallistaFlightService; @@ -134,10 +135,12 @@ async fn main() -> Result<()> { opt.concurrent_tasks, )); - let scheduler = SchedulerGrpcClient::connect(scheduler_url) + let connection = create_grpc_client_connection(scheduler_url) .await .context("Could not connect to scheduler")?; + let scheduler = SchedulerGrpcClient::new(connection); + let default_codec: BallistaCodec = BallistaCodec::default(); diff --git a/ballista/rust/scheduler/src/state/task_manager.rs b/ballista/rust/scheduler/src/state/task_manager.rs index e3ceb610dc..17abeebe7a 100644 --- a/ballista/rust/scheduler/src/state/task_manager.rs +++ b/ballista/rust/scheduler/src/state/task_manager.rs @@ -440,7 +440,8 @@ impl TaskManager } else { let executor_id = executor.id.clone(); let executor_url = format!("http://{}:{}", executor.host, executor.grpc_port); - let mut client = ExecutorGrpcClient::connect(executor_url).await?; + let connection = ballista_core::utils::create_grpc_client_connection(executor_url).await?; + let mut client = ExecutorGrpcClient::new(connection); clients.insert(executor_id, client.clone()); client .launch_task(protobuf::LaunchTaskParams { From 387ce59b8c962b2bdc549c9e360955c76a5abcf4 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 5 Aug 2022 11:13:03 +0800 Subject: [PATCH 02/10] Add timeout and keep-alive settings for Grpc Server --- ballista/rust/executor/src/executor_server.rs | 8 +++++++- ballista/rust/executor/src/main.rs | 11 +++++++++-- ballista/rust/executor/src/standalone.rs | 13 ++++++++++--- ballista/rust/scheduler/src/main.rs | 5 +++++ ballista/rust/scheduler/src/standalone.rs | 13 ++++++++++--- ballista/rust/scheduler/src/state/task_manager.rs | 3 ++- 6 files changed, 43 insertions(+), 10 deletions(-) diff --git a/ballista/rust/executor/src/executor_server.rs b/ballista/rust/executor/src/executor_server.rs index 7fa4893964..c16ad2132e 100644 --- a/ballista/rust/executor/src/executor_server.rs +++ b/ballista/rust/executor/src/executor_server.rs @@ -84,7 +84,13 @@ pub async fn startup( info!("Setup executor grpc service for {:?}", addr); let server = ExecutorGrpcServer::new(executor_server.clone()); - let grpc_server_future = Server::builder().add_service(server).serve(addr); + let grpc_server_future = Server::builder() + .timeout(Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Duration::from_secs(3600))) + .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) + .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) + .add_service(server) + .serve(addr); tokio::spawn(async move { grpc_server_future.await }); } diff --git a/ballista/rust/executor/src/main.rs b/ballista/rust/executor/src/main.rs index 01048936a5..acceabc860 100644 --- a/ballista/rust/executor/src/main.rs +++ b/ballista/rust/executor/src/main.rs @@ -187,8 +187,15 @@ async fn main() -> Result<()> { "Ballista v{} Rust Executor listening on {:?}", BALLISTA_VERSION, addr ); - let server_future = - tokio::spawn(Server::builder().add_service(server).serve(addr)); + let server_future = tokio::spawn( + Server::builder() + .timeout(Core_Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Core_Duration::from_secs(3600))) + .http2_keepalive_interval(Option::Some(Core_Duration::from_secs(300))) + .http2_keepalive_timeout(Option::Some(Core_Duration::from_secs(20))) + .add_service(server) + .serve(addr), + ); server_future .await .context("Tokio error")? diff --git a/ballista/rust/executor/src/standalone.rs b/ballista/rust/executor/src/standalone.rs index a2685afa44..7b73f8cda5 100644 --- a/ballista/rust/executor/src/standalone.rs +++ b/ballista/rust/executor/src/standalone.rs @@ -30,6 +30,7 @@ use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_proto::logical_plan::AsLogicalPlan; use log::info; use std::sync::Arc; +use std::time::Duration; use tempfile::TempDir; use tokio::net::TcpListener; use tonic::transport::{Channel, Server}; @@ -84,9 +85,15 @@ pub async fn new_standalone_executor< let service = BallistaFlightService::new(executor.clone()); let server = FlightServiceServer::new(service); tokio::spawn( - Server::builder().add_service(server).serve_with_incoming( - tokio_stream::wrappers::TcpListenerStream::new(listener), - ), + Server::builder() + .timeout(Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Duration::from_secs(3600))) + .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) + .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) + .add_service(server) + .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new( + listener, + )), ); tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec)); diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs index 61e4ae3601..afe64cba46 100644 --- a/ballista/rust/scheduler/src/main.rs +++ b/ballista/rust/scheduler/src/main.rs @@ -23,6 +23,7 @@ use ballista_scheduler::scheduler_server::externalscaler::external_scaler_server use futures::future::{self, Either, TryFutureExt}; use hyper::{server::conn::AddrStream, service::make_service_fn, Server}; use std::convert::Infallible; +use std::time::Duration; use std::{net::SocketAddr, sync::Arc}; use tonic::transport::server::Connected; use tonic::transport::Server as TonicServer; @@ -109,6 +110,10 @@ async fn start_server( let keda_scaler = ExternalScalerServer::new(scheduler_server.clone()); let mut tonic = TonicServer::builder() + .timeout(Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Duration::from_secs(3600))) + .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) + .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) .add_service(scheduler_grpc_server) .add_service(flight_sql_server) .add_service(keda_scaler) diff --git a/ballista/rust/scheduler/src/standalone.rs b/ballista/rust/scheduler/src/standalone.rs index 4f358a4644..da756c666d 100644 --- a/ballista/rust/scheduler/src/standalone.rs +++ b/ballista/rust/scheduler/src/standalone.rs @@ -23,6 +23,7 @@ use ballista_core::{ }; use datafusion_proto::protobuf::LogicalPlanNode; use log::info; +use std::time::Duration; use std::{net::SocketAddr, sync::Arc}; use tokio::net::TcpListener; use tonic::transport::Server; @@ -50,9 +51,15 @@ pub async fn new_standalone_scheduler() -> Result { BALLISTA_VERSION, addr ); tokio::spawn( - Server::builder().add_service(server).serve_with_incoming( - tokio_stream::wrappers::TcpListenerStream::new(listener), - ), + Server::builder() + .timeout(Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Duration::from_secs(3600))) + .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) + .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) + .add_service(server) + .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new( + listener, + )), ); Ok(addr) diff --git a/ballista/rust/scheduler/src/state/task_manager.rs b/ballista/rust/scheduler/src/state/task_manager.rs index 17abeebe7a..c4c7fe53ad 100644 --- a/ballista/rust/scheduler/src/state/task_manager.rs +++ b/ballista/rust/scheduler/src/state/task_manager.rs @@ -440,7 +440,8 @@ impl TaskManager } else { let executor_id = executor.id.clone(); let executor_url = format!("http://{}:{}", executor.host, executor.grpc_port); - let connection = ballista_core::utils::create_grpc_client_connection(executor_url).await?; + let connection = + ballista_core::utils::create_grpc_client_connection(executor_url).await?; let mut client = ExecutorGrpcClient::new(connection); clients.insert(executor_id, client.clone()); client From b0976aa222437817edb0d01f26d34bfbc9c1f562 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 5 Aug 2022 11:30:31 +0800 Subject: [PATCH 03/10] move server settings to utils --- ballista/rust/core/src/utils.rs | 10 +++++++++- ballista/rust/executor/src/executor_server.rs | 11 +++-------- ballista/rust/executor/src/main.rs | 14 +++----------- ballista/rust/executor/src/standalone.rs | 10 +++------- ballista/rust/scheduler/src/main.rs | 9 ++------- ballista/rust/scheduler/src/standalone.rs | 10 ++-------- 6 files changed, 22 insertions(+), 42 deletions(-) diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs index c2f465c763..5e50f923d8 100644 --- a/ballista/rust/core/src/utils.rs +++ b/ballista/rust/core/src/utils.rs @@ -52,7 +52,7 @@ use std::sync::Arc; use std::time::Duration; use std::{fs::File, pin::Pin}; use tonic::codegen::StdError; -use tonic::transport::{Channel, Error}; +use tonic::transport::{Channel, Error, Server}; /// Stream data to disk in Arrow IPC format @@ -330,3 +330,11 @@ where .keep_alive_while_idle(true); endpoint.connect().await } + +pub fn create_grpc_server() -> Server { + Server::builder() + .timeout(Duration::from_secs(20)) + .tcp_keepalive(Option::Some(Duration::from_secs(3600))) + .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) + .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) +} diff --git a/ballista/rust/executor/src/executor_server.rs b/ballista/rust/executor/src/executor_server.rs index c16ad2132e..e484e05a20 100644 --- a/ballista/rust/executor/src/executor_server.rs +++ b/ballista/rust/executor/src/executor_server.rs @@ -22,7 +22,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::mpsc; use log::{debug, error, info}; -use tonic::transport::{Channel, Server}; +use tonic::transport::Channel; use tonic::{Request, Response, Status}; use ballista_core::error::BallistaError; @@ -39,6 +39,7 @@ use ballista_core::serde::protobuf::{ }; use ballista_core::serde::scheduler::ExecutorState; use ballista_core::serde::{AsExecutionPlan, BallistaCodec}; +use ballista_core::utils::create_grpc_server; use datafusion::execution::context::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::logical_plan::AsLogicalPlan; @@ -84,13 +85,7 @@ pub async fn startup( info!("Setup executor grpc service for {:?}", addr); let server = ExecutorGrpcServer::new(executor_server.clone()); - let grpc_server_future = Server::builder() - .timeout(Duration::from_secs(20)) - .tcp_keepalive(Option::Some(Duration::from_secs(3600))) - .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) - .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) - .add_service(server) - .serve(addr); + let grpc_server_future = create_grpc_server().add_service(server).serve(addr); tokio::spawn(async move { grpc_server_future.await }); } diff --git a/ballista/rust/executor/src/main.rs b/ballista/rust/executor/src/main.rs index acceabc860..cc7f41e373 100644 --- a/ballista/rust/executor/src/main.rs +++ b/ballista/rust/executor/src/main.rs @@ -28,7 +28,6 @@ use log::{error, info}; use tempfile::TempDir; use tokio::fs::ReadDir; use tokio::{fs, time}; -use tonic::transport::Server; use uuid::Uuid; use ballista_core::config::TaskSchedulingPolicy; @@ -39,7 +38,7 @@ use ballista_core::serde::protobuf::{ }; use ballista_core::serde::scheduler::ExecutorSpecification; use ballista_core::serde::BallistaCodec; -use ballista_core::utils::create_grpc_client_connection; +use ballista_core::utils::{create_grpc_client_connection, create_grpc_server}; use ballista_core::{print_version, BALLISTA_VERSION}; use ballista_executor::executor::Executor; use ballista_executor::flight_service::BallistaFlightService; @@ -187,15 +186,8 @@ async fn main() -> Result<()> { "Ballista v{} Rust Executor listening on {:?}", BALLISTA_VERSION, addr ); - let server_future = tokio::spawn( - Server::builder() - .timeout(Core_Duration::from_secs(20)) - .tcp_keepalive(Option::Some(Core_Duration::from_secs(3600))) - .http2_keepalive_interval(Option::Some(Core_Duration::from_secs(300))) - .http2_keepalive_timeout(Option::Some(Core_Duration::from_secs(20))) - .add_service(server) - .serve(addr), - ); + let server_future = + tokio::spawn(create_grpc_server().add_service(server).serve(addr)); server_future .await .context("Tokio error")? diff --git a/ballista/rust/executor/src/standalone.rs b/ballista/rust/executor/src/standalone.rs index 7b73f8cda5..ca5513fae2 100644 --- a/ballista/rust/executor/src/standalone.rs +++ b/ballista/rust/executor/src/standalone.rs @@ -20,6 +20,7 @@ use crate::{execution_loop, executor::Executor, flight_service::BallistaFlightSe use arrow_flight::flight_service_server::FlightServiceServer; use ballista_core::serde::scheduler::ExecutorSpecification; use ballista_core::serde::{AsExecutionPlan, BallistaCodec}; +use ballista_core::utils::create_grpc_server; use ballista_core::{ error::Result, serde::protobuf::executor_registration::OptionalHost, @@ -30,10 +31,9 @@ use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; use datafusion_proto::logical_plan::AsLogicalPlan; use log::info; use std::sync::Arc; -use std::time::Duration; use tempfile::TempDir; use tokio::net::TcpListener; -use tonic::transport::{Channel, Server}; +use tonic::transport::Channel; use uuid::Uuid; pub async fn new_standalone_executor< @@ -85,11 +85,7 @@ pub async fn new_standalone_executor< let service = BallistaFlightService::new(executor.clone()); let server = FlightServiceServer::new(service); tokio::spawn( - Server::builder() - .timeout(Duration::from_secs(20)) - .tcp_keepalive(Option::Some(Duration::from_secs(3600))) - .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) - .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) + create_grpc_server() .add_service(server) .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new( listener, diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs index afe64cba46..417e18559e 100644 --- a/ballista/rust/scheduler/src/main.rs +++ b/ballista/rust/scheduler/src/main.rs @@ -23,10 +23,8 @@ use ballista_scheduler::scheduler_server::externalscaler::external_scaler_server use futures::future::{self, Either, TryFutureExt}; use hyper::{server::conn::AddrStream, service::make_service_fn, Server}; use std::convert::Infallible; -use std::time::Duration; use std::{net::SocketAddr, sync::Arc}; use tonic::transport::server::Connected; -use tonic::transport::Server as TonicServer; use tower::Service; use ballista_core::BALLISTA_VERSION; @@ -61,6 +59,7 @@ mod config { )); } +use ballista_core::utils::create_grpc_server; use ballista_scheduler::flight_sql::FlightSqlServiceImpl; use config::prelude::*; use datafusion::execution::context::default_session_builder; @@ -109,11 +108,7 @@ async fn start_server( let keda_scaler = ExternalScalerServer::new(scheduler_server.clone()); - let mut tonic = TonicServer::builder() - .timeout(Duration::from_secs(20)) - .tcp_keepalive(Option::Some(Duration::from_secs(3600))) - .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) - .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) + let mut tonic = create_grpc_server() .add_service(scheduler_grpc_server) .add_service(flight_sql_server) .add_service(keda_scaler) diff --git a/ballista/rust/scheduler/src/standalone.rs b/ballista/rust/scheduler/src/standalone.rs index da756c666d..4c2f5f0e66 100644 --- a/ballista/rust/scheduler/src/standalone.rs +++ b/ballista/rust/scheduler/src/standalone.rs @@ -17,17 +17,15 @@ use ballista_core::serde::protobuf::PhysicalPlanNode; use ballista_core::serde::BallistaCodec; +use ballista_core::utils::create_grpc_server; use ballista_core::{ error::Result, serde::protobuf::scheduler_grpc_server::SchedulerGrpcServer, BALLISTA_VERSION, }; use datafusion_proto::protobuf::LogicalPlanNode; use log::info; -use std::time::Duration; use std::{net::SocketAddr, sync::Arc}; use tokio::net::TcpListener; -use tonic::transport::Server; - use crate::{ scheduler_server::SchedulerServer, state::backend::standalone::StandaloneClient, }; @@ -51,11 +49,7 @@ pub async fn new_standalone_scheduler() -> Result { BALLISTA_VERSION, addr ); tokio::spawn( - Server::builder() - .timeout(Duration::from_secs(20)) - .tcp_keepalive(Option::Some(Duration::from_secs(3600))) - .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) - .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) + create_grpc_server() .add_service(server) .serve_with_incoming(tokio_stream::wrappers::TcpListenerStream::new( listener, From cc1b0d41b52c6a6dfa2b219040eaf324565fe19c Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 5 Aug 2022 11:34:45 +0800 Subject: [PATCH 04/10] fix fmt --- ballista/rust/scheduler/src/standalone.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ballista/rust/scheduler/src/standalone.rs b/ballista/rust/scheduler/src/standalone.rs index 4c2f5f0e66..0de81b85cd 100644 --- a/ballista/rust/scheduler/src/standalone.rs +++ b/ballista/rust/scheduler/src/standalone.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +use crate::{ + scheduler_server::SchedulerServer, state::backend::standalone::StandaloneClient, +}; use ballista_core::serde::protobuf::PhysicalPlanNode; use ballista_core::serde::BallistaCodec; use ballista_core::utils::create_grpc_server; @@ -26,9 +29,6 @@ use datafusion_proto::protobuf::LogicalPlanNode; use log::info; use std::{net::SocketAddr, sync::Arc}; use tokio::net::TcpListener; -use crate::{ - scheduler_server::SchedulerServer, state::backend::standalone::StandaloneClient, -}; pub async fn new_standalone_scheduler() -> Result { let client = StandaloneClient::try_new_temporary()?; From 008366a2240b15e58590800ef3edd1b9b18da2f1 Mon Sep 17 00:00:00 2001 From: Wang Date: Fri, 5 Aug 2022 11:46:13 +0800 Subject: [PATCH 05/10] set tcp_nodelay to true explicitly --- ballista/rust/core/src/utils.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs index 5e50f923d8..d356703a03 100644 --- a/ballista/rust/core/src/utils.rs +++ b/ballista/rust/core/src/utils.rs @@ -324,6 +324,8 @@ where let endpoint = tonic::transport::Endpoint::new(dst)? .connect_timeout(Duration::from_secs(20)) .timeout(Duration::from_secs(20)) + // Disable Nagle's Algorithm since we don't want packets to wait + .tcp_nodelay(true) .tcp_keepalive(Option::Some(Duration::from_secs(3600))) .http2_keep_alive_interval(Duration::from_secs(300)) .keep_alive_timeout(Duration::from_secs(20)) @@ -334,6 +336,8 @@ where pub fn create_grpc_server() -> Server { Server::builder() .timeout(Duration::from_secs(20)) + // Disable Nagle's Algorithm since we don't want packets to wait + .tcp_nodelay(true) .tcp_keepalive(Option::Some(Duration::from_secs(3600))) .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) From d9bd029eaa7821669695cdbd70eadc0c4f3f6041 Mon Sep 17 00:00:00 2001 From: Wang Date: Wed, 10 Aug 2022 19:31:37 +0800 Subject: [PATCH 06/10] Ballista Executor report plan/operators metrics to Ballista Scheduler --- ballista/rust/core/proto/ballista.proto | 37 +++++++ .../core/src/serde/scheduler/from_proto.rs | 98 +++++++++++++++++ .../rust/core/src/serde/scheduler/to_proto.rs | 68 ++++++++++++ ballista/rust/core/src/utils.rs | 14 +++ ballista/rust/executor/src/execution_loop.rs | 12 +- ballista/rust/executor/src/executor_server.rs | 17 ++- ballista/rust/executor/src/lib.rs | 15 ++- ballista/rust/scheduler/src/planner.rs | 2 +- .../src/scheduler_server/event_loop.rs | 1 + .../scheduler/src/scheduler_server/mod.rs | 3 + .../scheduler/src/state/execution_graph.rs | 104 +++++++++++++++++- .../rust/scheduler/src/state/task_manager.rs | 20 ++++ 12 files changed, 378 insertions(+), 13 deletions(-) diff --git a/ballista/rust/core/proto/ballista.proto b/ballista/rust/core/proto/ballista.proto index 4e2c55f69b..04ada4bcff 100644 --- a/ballista/rust/core/proto/ballista.proto +++ b/ballista/rust/core/proto/ballista.proto @@ -429,6 +429,7 @@ message ExecutionGraphStage { repeated TaskStatus task_statuses = 6; uint32 output_link = 7; bool resolved = 8; + repeated OperatorMetricsSet stage_metrics = 9; } message ExecutionGraph { @@ -503,6 +504,41 @@ message ColumnStats { uint32 distinct_count = 4; } +message OperatorMetricsSet { + repeated OperatorMetric metrics = 1; +} + + +message NamedCount { + string name = 1; + uint64 value = 2; +} + +message NamedGauge { + string name = 1; + uint64 value = 2; +} + +message NamedTime { + string name = 1; + uint64 value = 2; +} + +message OperatorMetric { + oneof metric { + uint64 output_rows = 1; + uint64 elapse_time = 2; + uint64 spill_count = 3; + uint64 spilled_bytes = 4; + uint64 current_memory_usage = 5; + NamedCount count = 6; + NamedGauge gauge = 7; + NamedTime time = 8; + int64 start_timestamp = 9; + int64 end_timestamp = 10; + } +} + // Used by scheduler message ExecutorMetadata { string id = 1; @@ -594,6 +630,7 @@ message TaskStatus { FailedTask failed = 3; CompletedTask completed = 4; } + repeated OperatorMetricsSet metrics = 5; } message PollWorkParams { diff --git a/ballista/rust/core/src/serde/scheduler/from_proto.rs b/ballista/rust/core/src/serde/scheduler/from_proto.rs index b401f1fdf7..970a5878b2 100644 --- a/ballista/rust/core/src/serde/scheduler/from_proto.rs +++ b/ballista/rust/core/src/serde/scheduler/from_proto.rs @@ -15,11 +15,19 @@ // specific language governing permissions and limitations // under the License. +use chrono::{TimeZone, Utc}; +use datafusion::physical_plan::metrics::{ + Count, Gauge, MetricValue, MetricsSet, Time, Timestamp, +}; +use datafusion::physical_plan::Metric; use std::convert::TryInto; +use std::sync::Arc; +use std::time::Duration; use crate::error::BallistaError; use crate::serde::protobuf; use crate::serde::protobuf::action::ActionType; +use crate::serde::protobuf::{operator_metric, NamedCount, NamedGauge, NamedTime}; use crate::serde::scheduler::{Action, PartitionId, PartitionLocation, PartitionStats}; impl TryInto for protobuf::Action { @@ -104,3 +112,93 @@ impl TryInto for protobuf::PartitionLocation { }) } } + +impl TryInto for protobuf::OperatorMetric { + type Error = BallistaError; + + fn try_into(self) -> Result { + match self.metric { + Some(operator_metric::Metric::OutputRows(value)) => { + let count = Count::new(); + count.add(value as usize); + Ok(MetricValue::OutputRows(count)) + } + Some(operator_metric::Metric::ElapseTime(value)) => { + let time = Time::new(); + time.add_duration(Duration::from_nanos(value)); + Ok(MetricValue::ElapsedCompute(time)) + } + Some(operator_metric::Metric::SpillCount(value)) => { + let count = Count::new(); + count.add(value as usize); + Ok(MetricValue::SpillCount(count)) + } + Some(operator_metric::Metric::SpilledBytes(value)) => { + let count = Count::new(); + count.add(value as usize); + Ok(MetricValue::SpilledBytes(count)) + } + Some(operator_metric::Metric::CurrentMemoryUsage(value)) => { + let gauge = Gauge::new(); + gauge.add(value as usize); + Ok(MetricValue::CurrentMemoryUsage(gauge)) + } + Some(operator_metric::Metric::Count(NamedCount { name, value })) => { + let count = Count::new(); + count.add(value as usize); + Ok(MetricValue::Count { + name: name.into(), + count, + }) + } + Some(operator_metric::Metric::Gauge(NamedGauge { name, value })) => { + let gauge = Gauge::new(); + gauge.add(value as usize); + Ok(MetricValue::Gauge { + name: name.into(), + gauge, + }) + } + Some(operator_metric::Metric::Time(NamedTime { name, value })) => { + let time = Time::new(); + time.add_duration(Duration::from_nanos(value)); + Ok(MetricValue::Time { + name: name.into(), + time, + }) + } + Some(operator_metric::Metric::StartTimestamp(value)) => { + let timestamp = Timestamp::new(); + timestamp.set(Utc.timestamp_nanos(value)); + Ok(MetricValue::StartTimestamp(timestamp)) + } + Some(operator_metric::Metric::EndTimestamp(value)) => { + let timestamp = Timestamp::new(); + timestamp.set(Utc.timestamp_nanos(value)); + Ok(MetricValue::EndTimestamp(timestamp)) + } + None => Err(BallistaError::General( + "scheduler::from_proto(OperatorMetric) metric is None.".to_owned(), + )), + } + } +} + +impl TryInto for protobuf::OperatorMetricsSet { + type Error = BallistaError; + + fn try_into(self) -> Result { + let mut ms = MetricsSet::new(); + let metrics = self + .metrics + .into_iter() + .map(|m| m.try_into()) + .collect::, BallistaError>>()?; + + for value in metrics { + let new_metric = Arc::new(Metric::new(value, None)); + ms.push(new_metric) + } + Ok(ms) + } +} diff --git a/ballista/rust/core/src/serde/scheduler/to_proto.rs b/ballista/rust/core/src/serde/scheduler/to_proto.rs index 4c1c5d1618..4dbf5d9678 100644 --- a/ballista/rust/core/src/serde/scheduler/to_proto.rs +++ b/ballista/rust/core/src/serde/scheduler/to_proto.rs @@ -15,11 +15,13 @@ // specific language governing permissions and limitations // under the License. +use datafusion::physical_plan::metrics::{MetricValue, MetricsSet}; use std::convert::TryInto; use crate::error::BallistaError; use crate::serde::protobuf; use crate::serde::protobuf::action::ActionType; +use crate::serde::protobuf::{operator_metric, NamedCount, NamedGauge, NamedTime}; use crate::serde::scheduler::{Action, PartitionId, PartitionLocation, PartitionStats}; use datafusion::physical_plan::Partitioning; @@ -105,3 +107,69 @@ pub fn hash_partitioning_to_proto( } } } + +impl TryInto for &MetricValue { + type Error = BallistaError; + + fn try_into(self) -> Result { + match self { + MetricValue::OutputRows(count) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::OutputRows(count.value() as u64)), + }), + MetricValue::ElapsedCompute(time) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::ElapseTime(time.value() as u64)), + }), + MetricValue::SpillCount(count) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::SpillCount(count.value() as u64)), + }), + MetricValue::SpilledBytes(count) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::SpilledBytes(count.value() as u64)), + }), + MetricValue::CurrentMemoryUsage(gauge) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::CurrentMemoryUsage( + gauge.value() as u64 + )), + }), + MetricValue::Count { name, count } => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::Count(NamedCount { + name: name.to_string(), + value: count.value() as u64, + })), + }), + MetricValue::Gauge { name, gauge } => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::Gauge(NamedGauge { + name: name.to_string(), + value: gauge.value() as u64, + })), + }), + MetricValue::Time { name, time } => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::Time(NamedTime { + name: name.to_string(), + value: time.value() as u64, + })), + }), + MetricValue::StartTimestamp(timestamp) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::StartTimestamp( + timestamp.value().map(|m| m.timestamp_nanos()).unwrap_or(0), + )), + }), + MetricValue::EndTimestamp(timestamp) => Ok(protobuf::OperatorMetric { + metric: Some(operator_metric::Metric::EndTimestamp( + timestamp.value().map(|m| m.timestamp_nanos()).unwrap_or(0), + )), + }), + } + } +} + +impl TryInto for MetricsSet { + type Error = BallistaError; + + fn try_into(self) -> Result { + let metrics = self + .iter() + .map(|m| m.value().try_into()) + .collect::, BallistaError>>()?; + Ok(protobuf::OperatorMetricsSet { metrics }) + } +} diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs index d356703a03..5fe280e65c 100644 --- a/ballista/rust/core/src/utils.rs +++ b/ballista/rust/core/src/utils.rs @@ -38,6 +38,7 @@ use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::file_format::{CsvExec, ParquetExec}; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::hash_join::HashJoinExec; +use datafusion::physical_plan::metrics::MetricsSet; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::{metrics, ExecutionPlan, RecordBatchStream}; @@ -342,3 +343,16 @@ pub fn create_grpc_server() -> Server { .http2_keepalive_interval(Option::Some(Duration::from_secs(300))) .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) } + +pub fn collect_plan_metrics(plan: &Arc) -> Vec { + let mut metrics_array = Vec::::new(); + if let Some(metrics) = plan.metrics() { + metrics_array.push(metrics); + } + plan.children().iter().for_each(|c| { + collect_plan_metrics(c) + .into_iter() + .for_each(|e| metrics_array.push(e)) + }); + metrics_array +} diff --git a/ballista/rust/executor/src/execution_loop.rs b/ballista/rust/executor/src/execution_loop.rs index c94081244a..63b02fc138 100644 --- a/ballista/rust/executor/src/execution_loop.rs +++ b/ballista/rust/executor/src/execution_loop.rs @@ -28,12 +28,14 @@ use ballista_core::error::BallistaError; use ballista_core::serde::physical_plan::from_proto::parse_protobuf_hash_partitioning; use ballista_core::serde::scheduler::ExecutorSpecification; use ballista_core::serde::{AsExecutionPlan, BallistaCodec}; +use ballista_core::utils::collect_plan_metrics; use datafusion::execution::context::TaskContext; use datafusion_proto::logical_plan::AsLogicalPlan; use futures::FutureExt; use log::{debug, error, info, trace, warn}; use std::any::Any; use std::collections::HashMap; +use std::convert::TryInto; use std::error::Error; use std::ops::Deref; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -190,7 +192,7 @@ async fn run_received_tasks, BallistaError>>() + .ok(); + let _ = task_status_sender.send(as_task_status( execution_result, executor.metadata.id.clone(), task_id, + operator_metrics, )); }); diff --git a/ballista/rust/executor/src/executor_server.rs b/ballista/rust/executor/src/executor_server.rs index e484e05a20..1555e1e197 100644 --- a/ballista/rust/executor/src/executor_server.rs +++ b/ballista/rust/executor/src/executor_server.rs @@ -16,6 +16,7 @@ // under the License. use std::collections::HashMap; +use std::convert::TryInto; use std::ops::Deref; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -39,7 +40,7 @@ use ballista_core::serde::protobuf::{ }; use ballista_core::serde::scheduler::ExecutorState; use ballista_core::serde::{AsExecutionPlan, BallistaCodec}; -use ballista_core::utils::create_grpc_server; +use ballista_core::utils::{collect_plan_metrics, create_grpc_server}; use datafusion::execution::context::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::logical_plan::AsLogicalPlan; @@ -239,7 +240,7 @@ impl ExecutorServer ExecutorServer, BallistaError>>()?; let executor_id = &self.executor.metadata.id; - let task_status = as_task_status(execution_result, executor_id.clone(), task_id); + let task_status = as_task_status( + execution_result, + executor_id.clone(), + task_id, + Some(operator_metrics), + ); let task_status_sender = self.executor_env.tx_task_status.clone(); task_status_sender.send(task_status).await.unwrap(); diff --git a/ballista/rust/executor/src/lib.rs b/ballista/rust/executor/src/lib.rs index 4d145b269e..e93993b671 100644 --- a/ballista/rust/executor/src/lib.rs +++ b/ballista/rust/executor/src/lib.rs @@ -32,21 +32,27 @@ pub use standalone::new_standalone_executor; use log::info; use ballista_core::serde::protobuf::{ - task_status, CompletedTask, FailedTask, PartitionId, ShuffleWritePartition, - TaskStatus, + task_status, CompletedTask, FailedTask, OperatorMetricsSet, PartitionId, + ShuffleWritePartition, TaskStatus, }; pub fn as_task_status( execution_result: ballista_core::error::Result>, executor_id: String, task_id: PartitionId, + operator_metrics: Option>, ) -> TaskStatus { + let metrics = operator_metrics.unwrap_or_default(); match execution_result { Ok(partitions) => { - info!("Task {:?} finished", task_id); - + info!( + "Task {:?} finished with operator_metrics array size {}", + task_id, + metrics.len() + ); TaskStatus { task_id: Some(task_id), + metrics, status: Some(task_status::Status::Completed(CompletedTask { executor_id, partitions, @@ -59,6 +65,7 @@ pub fn as_task_status( TaskStatus { task_id: Some(task_id), + metrics, status: Some(task_status::Status::Failed(FailedTask { error: format!("Task failed due to Tokio error: {}", error_msg), })), diff --git a/ballista/rust/scheduler/src/planner.rs b/ballista/rust/scheduler/src/planner.rs index 1c94693728..9c393bace4 100644 --- a/ballista/rust/scheduler/src/planner.rs +++ b/ballista/rust/scheduler/src/planner.rs @@ -63,7 +63,7 @@ impl DistributedPlanner { job_id: &'a str, execution_plan: Arc, ) -> Result>> { - info!("planning query stages"); + info!("planning query stages for job {}", job_id); let (new_plan, mut stages) = self.plan_query_stages_internal(job_id, execution_plan)?; stages.push(create_shuffle_writer( diff --git a/ballista/rust/scheduler/src/scheduler_server/event_loop.rs b/ballista/rust/scheduler/src/scheduler_server/event_loop.rs index b0037ab0cf..d6397ba6b9 100644 --- a/ballista/rust/scheduler/src/scheduler_server/event_loop.rs +++ b/ballista/rust/scheduler/src/scheduler_server/event_loop.rs @@ -322,6 +322,7 @@ mod test { stage_id: 1, partition_id: 0, }), + metrics: vec![], status: Some(task_status::Status::Completed(CompletedTask { executor_id: "executor-1".to_string(), partitions, diff --git a/ballista/rust/scheduler/src/scheduler_server/mod.rs b/ballista/rust/scheduler/src/scheduler_server/mod.rs index a6a26d8064..2fac229b0a 100644 --- a/ballista/rust/scheduler/src/scheduler_server/mod.rs +++ b/ballista/rust/scheduler/src/scheduler_server/mod.rs @@ -364,6 +364,7 @@ mod test { executor_id: "executor-1".to_owned(), partitions, })), + metrics: vec![], task_id: Some(PartitionId { job_id: job_id.to_owned(), stage_id: task.partition.stage_id as u32, @@ -493,6 +494,7 @@ mod test { partitions, }, )), + metrics: vec![], task_id: Some(PartitionId { job_id: job_id.to_owned(), stage_id: task.partition.stage_id as u32, @@ -627,6 +629,7 @@ mod test { error: "".to_string(), }, )), + metrics: vec![], task_id: Some(PartitionId { job_id: job_id.to_owned(), stage_id: task.partition.stage_id as u32, diff --git a/ballista/rust/scheduler/src/state/execution_graph.rs b/ballista/rust/scheduler/src/state/execution_graph.rs index 1412f7e016..59e4d78aec 100644 --- a/ballista/rust/scheduler/src/state/execution_graph.rs +++ b/ballista/rust/scheduler/src/state/execution_graph.rs @@ -20,7 +20,7 @@ use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ShuffleWriterExec, UnresolvedShuffleExec}; use ballista_core::serde::protobuf::{ - self, CompletedJob, JobStatus, QueuedJob, TaskStatus, + self, CompletedJob, JobStatus, OperatorMetricsSet, QueuedJob, TaskStatus, }; use ballista_core::serde::protobuf::{job_status, FailedJob, ShuffleWritePartition}; use ballista_core::serde::protobuf::{task_status, RunningTask}; @@ -28,14 +28,16 @@ use ballista_core::serde::scheduler::{ ExecutorMetadata, PartitionId, PartitionLocation, PartitionStats, }; use datafusion::physical_plan::{ - accept, ExecutionPlan, ExecutionPlanVisitor, Partitioning, + accept, ExecutionPlan, ExecutionPlanVisitor, Metric, Partitioning, }; use log::debug; use std::collections::HashMap; use std::convert::TryInto; use std::fmt::{Debug, Formatter}; +use ballista_core::utils::collect_plan_metrics; use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::metrics::{MetricValue, MetricsSet}; use std::sync::Arc; /// This data structure collects the partition locations for an `ExecutionStage`. @@ -103,6 +105,8 @@ pub struct ExecutionStage { /// Flag indicating whether all input partitions have been resolved and the plan /// has UnresovledShuffleExec operators resolved to ShuffleReadExec operators. pub(crate) resolved: bool, + /// Combined metrics of the already finished tasks in the stage, If it is None, no task is finished yet. + pub(crate) stage_metrics: Option>, } impl Debug for ExecutionStage { @@ -153,6 +157,7 @@ impl ExecutionStage { task_statuses: vec![None; num_tasks], output_link, resolved, + stage_metrics: None, } } @@ -230,6 +235,60 @@ impl ExecutionStage { self.task_statuses[partition] = Some(status); } + /// update and combine the task metrics to the stage metrics + pub fn update_task_metrics( + &mut self, + partition: usize, + metrics: Vec, + ) -> Result<()> { + if let Some(combined_metrics) = &mut self.stage_metrics { + if metrics.len() != combined_metrics.len() { + return Err(BallistaError::Internal(format!("Error updating task metrics to stage {}, task metrics array size {} does not equal \ + with the stage metrics array size {} for task {}", self.stage_id, metrics.len(), combined_metrics.len(), partition))); + } + let metrics_values_array = metrics + .into_iter() + .map(|ms| { + ms.metrics + .into_iter() + .map(|m| m.try_into()) + .collect::>>() + }) + .collect::>>()?; + + let new_metrics_set = combined_metrics + .iter_mut() + .zip(metrics_values_array) + .map(|(first, second)| { + Self::combine_metrics_set(first, second, partition) + }) + .collect(); + self.stage_metrics = Some(new_metrics_set) + } else { + let new_metrics_set = metrics + .into_iter() + .map(|ms| ms.try_into()) + .collect::>>()?; + if !new_metrics_set.is_empty() { + self.stage_metrics = Some(new_metrics_set) + } + } + Ok(()) + } + + pub fn combine_metrics_set( + first: &mut MetricsSet, + second: Vec, + partition: usize, + ) -> MetricsSet { + for metric_value in second { + // TODO recheck the lable logic + let new_metric = Arc::new(Metric::new(metric_value, Some(partition))); + first.push(new_metric); + } + first.aggregate_by_partition() + } + /// Add input partitions published from an input stage. pub fn add_input_partitions( &mut self, @@ -467,8 +526,8 @@ impl ExecutionGraph { self.stages.values().all(|s| s.complete()) } - /// Update task statuses in the graph. This will push shuffle partitions to their - /// respective shuffle read stages. + /// Update task statuses and task metrics in the graph. + /// This will also push shuffle partitions to their respective shuffle read stages. pub fn update_task_status( &mut self, executor: &ExecutorMetadata, @@ -482,6 +541,7 @@ impl ExecutionGraph { stage_id, partition_id, }), + metrics: operator_metrics, status: Some(task_status), } = status { @@ -497,6 +557,7 @@ impl ExecutionGraph { let partition = partition_id as usize; if let Some(stage) = self.stages.get_mut(&stage_id) { stage.update_task_status(partition, task_status.clone()); + let stage_plan = stage.plan.clone(); let stage_complete = stage.complete(); // TODO Should be able to reschedule this task. @@ -513,6 +574,40 @@ impl ExecutionGraph { } else if let task_status::Status::Completed(completed_task) = task_status { + // update task metrics for completed task + stage.update_task_metrics(partition, operator_metrics)?; + + // if this stage is completed, we want to combine the stage metrics to plan's metric set and print out the plan + if stage_complete && stage.stage_metrics.as_ref().is_some() { + let mut plan_metrics = collect_plan_metrics(&stage_plan); + let stage_metrics = stage + .stage_metrics + .as_ref() + .expect("stage metrics should not be None."); + if plan_metrics.len() != stage_metrics.len() { + return Err(BallistaError::Internal(format!("Error combine stage metrics to plan for stage {}, plan metrics array size {} does not equal \ + to the stage metrics array size {}", stage_id, plan_metrics.len(), stage_metrics.len()))); + } + plan_metrics.iter_mut().zip(stage_metrics).for_each( + |(plan_metric, stage_metric)| { + stage_metric + .iter() + .for_each(|s| plan_metric.push(s.clone())); + }, + ); + + // TODO the plan_metrics update above is a snapshot clone from the plan metrics. + // TODO Need to modify DataFusion to return metricset reference, not clone. + + println!( + "=== [{}/{}/{}] Stage finished, physical plan with metrics ===\n{}\n", + job_id, + stage_id, + partition, + DisplayableExecutionPlan::with_full_metrics(stage_plan.as_ref()).indent() + ); + } + let locations = partition_to_location( self.job_id.as_str(), stage_id, @@ -808,6 +903,7 @@ mod test { executor_id: "executor-1".to_owned(), partitions, })), + metrics: vec![], task_id: Some(protobuf::PartitionId { job_id: job_id.clone(), stage_id: task.partition.stage_id as u32, diff --git a/ballista/rust/scheduler/src/state/task_manager.rs b/ballista/rust/scheduler/src/state/task_manager.rs index c4c7fe53ad..ab9ff26e7a 100644 --- a/ballista/rust/scheduler/src/state/task_manager.rs +++ b/ballista/rust/scheduler/src/state/task_manager.rs @@ -610,6 +610,16 @@ impl TaskManager }, ); } + let stage_metrics = if stage.stage_metrics.is_empty() { + None + } else { + let ms = stage + .stage_metrics + .into_iter() + .map(|m| m.try_into()) + .collect::>>()?; + Some(ms) + }; let execution_stage = ExecutionStage { stage_id: stage.stage_id as usize, @@ -620,6 +630,7 @@ impl TaskManager task_statuses, output_link, resolved: stage.resolved, + stage_metrics, }; stages.insert(stage_id, execution_stage); } @@ -701,6 +712,8 @@ impl TaskManager stage_id: stage_id as u32, partition_id: partition as u32, }), + // task metrics should not persist. + metrics: vec![], status: Some(status), }) }) @@ -709,6 +722,12 @@ impl TaskManager let output_partitioning = hash_partitioning_to_proto(stage.output_partitioning.as_ref())?; + let stage_metrics = stage + .stage_metrics + .unwrap_or_default() + .into_iter() + .map(|m| m.try_into()) + .collect::>>()?; Ok(protobuf::ExecutionGraphStage { stage_id: stage_id as u64, partitions: stage.partitions as u32, @@ -718,6 +737,7 @@ impl TaskManager task_statuses, output_link, resolved: stage.resolved, + stage_metrics, }) }) .collect::>>()?; From ceb86565d33e0024bbeaa742e6bd14788b5fd30e Mon Sep 17 00:00:00 2001 From: Wang Date: Mon, 15 Aug 2022 11:11:44 +0800 Subject: [PATCH 07/10] Resolve review comments --- ballista/rust/scheduler/src/state/execution_graph.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ballista/rust/scheduler/src/state/execution_graph.rs b/ballista/rust/scheduler/src/state/execution_graph.rs index 59e4d78aec..e52ae2402b 100644 --- a/ballista/rust/scheduler/src/state/execution_graph.rs +++ b/ballista/rust/scheduler/src/state/execution_graph.rs @@ -30,7 +30,7 @@ use ballista_core::serde::scheduler::{ use datafusion::physical_plan::{ accept, ExecutionPlan, ExecutionPlanVisitor, Metric, Partitioning, }; -use log::debug; +use log::{debug, log}; use std::collections::HashMap; use std::convert::TryInto; use std::fmt::{Debug, Formatter}; @@ -599,12 +599,12 @@ impl ExecutionGraph { // TODO the plan_metrics update above is a snapshot clone from the plan metrics. // TODO Need to modify DataFusion to return metricset reference, not clone. - println!( + log!( "=== [{}/{}/{}] Stage finished, physical plan with metrics ===\n{}\n", job_id, stage_id, partition, - DisplayableExecutionPlan::with_full_metrics(stage_plan.as_ref()).indent() + DisplayableExecutionPlan::with_metrics(stage_plan.as_ref()).indent() ); } From 736d0fdc102733fcb2f89aac70a6ed536bef8458 Mon Sep 17 00:00:00 2001 From: Wang Date: Mon, 15 Aug 2022 12:45:03 +0800 Subject: [PATCH 08/10] Fix plan display with metrics --- ballista/rust/scheduler/src/display.rs | 128 ++++++++++++++++++ ballista/rust/scheduler/src/lib.rs | 1 + .../scheduler/src/state/execution_graph.rs | 12 +- 3 files changed, 135 insertions(+), 6 deletions(-) create mode 100644 ballista/rust/scheduler/src/display.rs diff --git a/ballista/rust/scheduler/src/display.rs b/ballista/rust/scheduler/src/display.rs new file mode 100644 index 0000000000..e4557a6b52 --- /dev/null +++ b/ballista/rust/scheduler/src/display.rs @@ -0,0 +1,128 @@ +// 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. + +//! Implementation of ballista physical plan display with metrics. See +//! [`crate::physical_plan::displayable`] for examples of how to +//! format + +use datafusion::logical_plan::{StringifiedPlan, ToStringifiedPlan}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ + accept, DisplayFormatType, ExecutionPlan, ExecutionPlanVisitor, +}; +use std::fmt; + +/// Wraps an `ExecutionPlan` to display this plan with metrics collected/aggregated. +/// The metrics must be collected in the same order as how we visit and display the plan. +pub struct DisplayableBallistaExecutionPlan<'a> { + inner: &'a dyn ExecutionPlan, + metrics: &'a Vec, +} + +impl<'a> DisplayableBallistaExecutionPlan<'a> { + /// Create a wrapper around an [`'ExecutionPlan'] which can be + /// pretty printed with aggregated metrics. + pub fn new(inner: &'a dyn ExecutionPlan, metrics: &'a Vec) -> Self { + Self { inner, metrics } + } + + /// Return a `format`able structure that produces a single line + /// per node. + /// + /// ```text + /// ProjectionExec: expr=[a] + /// CoalesceBatchesExec: target_batch_size=4096 + /// FilterExec: a < 5 + /// RepartitionExec: partitioning=RoundRobinBatch(16) + /// CsvExec: source=...", + /// ``` + pub fn indent(&self) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + metrics: &'a Vec, + } + impl<'a> fmt::Display for Wrapper<'a> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let t = DisplayFormatType::Default; + let mut visitor = IndentVisitor { + t, + f, + indent: 0, + metrics: self.metrics, + metric_index: 0, + }; + accept(self.plan, &mut visitor) + } + } + Wrapper { + plan: self.inner, + metrics: self.metrics, + } + } +} + +/// Formats plans with a single line per node. +struct IndentVisitor<'a, 'b> { + /// How to format each node + t: DisplayFormatType, + /// Write to this formatter + f: &'a mut fmt::Formatter<'b>, + /// Indent size + indent: usize, + /// The metrics along with the plan + metrics: &'a Vec, + /// The metric index + metric_index: usize, +} + +impl<'a, 'b> ExecutionPlanVisitor for IndentVisitor<'a, 'b> { + type Error = fmt::Error; + fn pre_visit( + &mut self, + plan: &dyn ExecutionPlan, + ) -> std::result::Result { + write!(self.f, "{:indent$}", "", indent = self.indent * 2)?; + plan.fmt_as(self.t, self.f)?; + if let Some(metrics) = self.metrics.get(self.metric_index) { + let metrics = metrics + .aggregate_by_partition() + .sorted_for_display() + .timestamps_removed(); + write!(self.f, ", metrics=[{}]", metrics)?; + } else { + write!(self.f, ", metrics=[]")?; + } + writeln!(self.f)?; + self.indent += 1; + self.metric_index += 1; + Ok(true) + } + + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result { + self.indent -= 1; + Ok(true) + } +} + +impl<'a> ToStringifiedPlan for DisplayableBallistaExecutionPlan<'a> { + fn to_stringified( + &self, + plan_type: datafusion::logical_plan::PlanType, + ) -> StringifiedPlan { + StringifiedPlan::new(plan_type, self.indent().to_string()) + } +} diff --git a/ballista/rust/scheduler/src/lib.rs b/ballista/rust/scheduler/src/lib.rs index 2c2fa4100a..95d2f00ea4 100644 --- a/ballista/rust/scheduler/src/lib.rs +++ b/ballista/rust/scheduler/src/lib.rs @@ -20,6 +20,7 @@ pub mod api; pub mod planner; pub mod scheduler_server; +pub mod display; #[cfg(feature = "sled")] pub mod standalone; pub mod state; diff --git a/ballista/rust/scheduler/src/state/execution_graph.rs b/ballista/rust/scheduler/src/state/execution_graph.rs index e52ae2402b..766d6ec1bc 100644 --- a/ballista/rust/scheduler/src/state/execution_graph.rs +++ b/ballista/rust/scheduler/src/state/execution_graph.rs @@ -16,6 +16,7 @@ // under the License. use crate::planner::DistributedPlanner; +use crate::display::DisplayableBallistaExecutionPlan; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ShuffleWriterExec, UnresolvedShuffleExec}; @@ -30,7 +31,7 @@ use ballista_core::serde::scheduler::{ use datafusion::physical_plan::{ accept, ExecutionPlan, ExecutionPlanVisitor, Metric, Partitioning, }; -use log::{debug, log}; +use log::{debug, info}; use std::collections::HashMap; use std::convert::TryInto; use std::fmt::{Debug, Formatter}; @@ -579,6 +580,8 @@ impl ExecutionGraph { // if this stage is completed, we want to combine the stage metrics to plan's metric set and print out the plan if stage_complete && stage.stage_metrics.as_ref().is_some() { + // The plan_metrics collected here is a snapshot clone from the plan metrics. + // They are all empty now and need to combine with the stage metrics in the ExecutionStages let mut plan_metrics = collect_plan_metrics(&stage_plan); let stage_metrics = stage .stage_metrics @@ -596,15 +599,12 @@ impl ExecutionGraph { }, ); - // TODO the plan_metrics update above is a snapshot clone from the plan metrics. - // TODO Need to modify DataFusion to return metricset reference, not clone. - - log!( + info!( "=== [{}/{}/{}] Stage finished, physical plan with metrics ===\n{}\n", job_id, stage_id, partition, - DisplayableExecutionPlan::with_metrics(stage_plan.as_ref()).indent() + DisplayableBallistaExecutionPlan::new(stage_plan.as_ref(), plan_metrics.as_ref()).indent() ); } From 439316d888906691a79b1183869bb8fccab0c7a2 Mon Sep 17 00:00:00 2001 From: Wang Date: Mon, 15 Aug 2022 13:48:25 +0800 Subject: [PATCH 09/10] fix shuffle writer metrics collection --- ballista/rust/core/src/utils.rs | 4 +-- ballista/rust/executor/src/execution_loop.rs | 10 ++++--- ballista/rust/executor/src/executor.rs | 26 ++++++++++++------- ballista/rust/executor/src/executor_server.rs | 10 +++++-- ballista/rust/executor/src/metrics/mod.rs | 7 ++--- ballista/rust/scheduler/src/lib.rs | 2 +- .../scheduler/src/state/execution_graph.rs | 5 ++-- 7 files changed, 42 insertions(+), 22 deletions(-) diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs index 5fe280e65c..ae8d57731f 100644 --- a/ballista/rust/core/src/utils.rs +++ b/ballista/rust/core/src/utils.rs @@ -344,13 +344,13 @@ pub fn create_grpc_server() -> Server { .http2_keepalive_timeout(Option::Some(Duration::from_secs(20))) } -pub fn collect_plan_metrics(plan: &Arc) -> Vec { +pub fn collect_plan_metrics(plan: &dyn ExecutionPlan) -> Vec { let mut metrics_array = Vec::::new(); if let Some(metrics) = plan.metrics() { metrics_array.push(metrics); } plan.children().iter().for_each(|c| { - collect_plan_metrics(c) + collect_plan_metrics(c.as_ref()) .into_iter() .for_each(|e| metrics_array.push(e)) }); diff --git a/ballista/rust/executor/src/execution_loop.rs b/ballista/rust/executor/src/execution_loop.rs index 63b02fc138..471377e417 100644 --- a/ballista/rust/executor/src/execution_loop.rs +++ b/ballista/rust/executor/src/execution_loop.rs @@ -185,14 +185,18 @@ async fn run_received_tasks, + shuffle_writer: Arc, task_ctx: Arc, _shuffle_output_partitioning: Option, ) -> Result, BallistaError> { + let partitions = shuffle_writer.execute_shuffle_write(part, task_ctx).await?; + self.metrics_collector + .record_stage(&job_id, stage_id, part, shuffle_writer); + + Ok(partitions) + } + + /// Recreate the shuffle writer with the correct working directory. + pub fn new_shuffle_writer( + &self, + job_id: String, + stage_id: usize, + plan: Arc, + ) -> Result, BallistaError> { let exec = if let Some(shuffle_writer) = plan.as_any().downcast_ref::() { // recreate the shuffle writer with the correct working directory ShuffleWriterExec::try_new( - job_id.clone(), + job_id, stage_id, plan.children()[0].clone(), self.work_dir.clone(), @@ -109,13 +123,7 @@ impl Executor { .to_string(), )) }?; - - let partitions = exec.execute_shuffle_write(part, task_ctx).await?; - - self.metrics_collector - .record_stage(&job_id, stage_id, part, exec); - - Ok(partitions) + Ok(Arc::new(exec)) } pub fn work_dir(&self) -> &str { diff --git a/ballista/rust/executor/src/executor_server.rs b/ballista/rust/executor/src/executor_server.rs index 1555e1e197..591c88a1cb 100644 --- a/ballista/rust/executor/src/executor_server.rs +++ b/ballista/rust/executor/src/executor_server.rs @@ -234,13 +234,19 @@ impl ExecutorServer ExecutorServer, ); } @@ -45,14 +46,14 @@ impl ExecutorMetricsCollector for LoggingMetricsCollector { job_id: &str, stage_id: usize, partition: usize, - plan: ShuffleWriterExec, + plan: Arc, ) { println!( "=== [{}/{}/{}] Physical plan with metrics ===\n{}\n", job_id, stage_id, partition, - DisplayableExecutionPlan::with_metrics(&plan).indent() + DisplayableExecutionPlan::with_metrics(plan.as_ref()).indent() ); } } diff --git a/ballista/rust/scheduler/src/lib.rs b/ballista/rust/scheduler/src/lib.rs index 95d2f00ea4..838eb56256 100644 --- a/ballista/rust/scheduler/src/lib.rs +++ b/ballista/rust/scheduler/src/lib.rs @@ -18,9 +18,9 @@ #![doc = include_str ! ("../README.md")] pub mod api; +pub mod display; pub mod planner; pub mod scheduler_server; -pub mod display; #[cfg(feature = "sled")] pub mod standalone; pub mod state; diff --git a/ballista/rust/scheduler/src/state/execution_graph.rs b/ballista/rust/scheduler/src/state/execution_graph.rs index 766d6ec1bc..b13f7c58c0 100644 --- a/ballista/rust/scheduler/src/state/execution_graph.rs +++ b/ballista/rust/scheduler/src/state/execution_graph.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::planner::DistributedPlanner; use crate::display::DisplayableBallistaExecutionPlan; +use crate::planner::DistributedPlanner; use ballista_core::error::{BallistaError, Result}; use ballista_core::execution_plans::{ShuffleWriterExec, UnresolvedShuffleExec}; @@ -582,7 +582,8 @@ impl ExecutionGraph { if stage_complete && stage.stage_metrics.as_ref().is_some() { // The plan_metrics collected here is a snapshot clone from the plan metrics. // They are all empty now and need to combine with the stage metrics in the ExecutionStages - let mut plan_metrics = collect_plan_metrics(&stage_plan); + let mut plan_metrics = + collect_plan_metrics(stage_plan.as_ref()); let stage_metrics = stage .stage_metrics .as_ref() From 21137fb90891e7b0b99e73bcf30a701b9c5a88ac Mon Sep 17 00:00:00 2001 From: yahoNanJing <90197956+yahoNanJing@users.noreply.github.com> Date: Mon, 15 Aug 2022 14:16:07 +0800 Subject: [PATCH 10/10] Remove Keyspace::QueuedJobs (#134) * Remove Keyspace::QueuedJobs * Fix UT * Fix cargo clippy for rust 1.63 Co-authored-by: yangzhong --- ballista/rust/core/src/config.rs | 2 +- .../src/execution_plans/distributed_query.rs | 25 ++++++++++------- ballista/rust/core/src/serde/mod.rs | 2 +- .../rust/core/src/serde/scheduler/to_proto.rs | 10 +++---- ballista/rust/scheduler/src/flight_sql.rs | 11 -------- .../scheduler/src/scheduler_server/grpc.rs | 11 -------- .../rust/scheduler/src/state/backend/mod.rs | 3 +- .../rust/scheduler/src/state/task_manager.rs | 28 +++---------------- 8 files changed, 26 insertions(+), 66 deletions(-) diff --git a/ballista/rust/core/src/config.rs b/ballista/rust/core/src/config.rs index 8975e8e30b..93162bc32c 100644 --- a/ballista/rust/core/src/config.rs +++ b/ballista/rust/core/src/config.rs @@ -92,7 +92,7 @@ impl BallistaConfigBuilder { } /// Ballista configuration -#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct BallistaConfig { /// Settings stored in map for easy serde settings: HashMap, diff --git a/ballista/rust/core/src/execution_plans/distributed_query.rs b/ballista/rust/core/src/execution_plans/distributed_query.rs index 11666a8e2b..e9d8528172 100644 --- a/ballista/rust/core/src/execution_plans/distributed_query.rs +++ b/ballista/rust/core/src/execution_plans/distributed_query.rs @@ -264,32 +264,37 @@ async fn execute_query( .await .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))? .into_inner(); - let status = status.and_then(|s| s.status).ok_or_else(|| { - DataFusionError::Internal("Received empty status message".to_owned()) - })?; + let status = status.and_then(|s| s.status); let wait_future = tokio::time::sleep(Duration::from_millis(100)); - let has_status_change = prev_status.map(|x| x != status).unwrap_or(true); + let has_status_change = prev_status != status; match status { - job_status::Status::Queued(_) => { + None => { + if has_status_change { + info!("Job {} still in initialization ...", job_id); + } + wait_future.await; + prev_status = status; + } + Some(job_status::Status::Queued(_)) => { if has_status_change { info!("Job {} still queued...", job_id); } wait_future.await; - prev_status = Some(status); + prev_status = status; } - job_status::Status::Running(_) => { + Some(job_status::Status::Running(_)) => { if has_status_change { info!("Job {} is running...", job_id); } wait_future.await; - prev_status = Some(status); + prev_status = status; } - job_status::Status::Failed(err) => { + Some(job_status::Status::Failed(err)) => { let msg = format!("Job {} failed: {}", job_id, err.error); error!("{}", msg); break Err(DataFusionError::Execution(msg)); } - job_status::Status::Completed(completed) => { + Some(job_status::Status::Completed(completed)) => { let streams = completed.partition_location.into_iter().map(|p| { let f = fetch_partition(p) .map_err(|e| ArrowError::ExternalError(Box::new(e))); diff --git a/ballista/rust/core/src/serde/mod.rs b/ballista/rust/core/src/serde/mod.rs index 20979f12e9..1548af182c 100644 --- a/ballista/rust/core/src/serde/mod.rs +++ b/ballista/rust/core/src/serde/mod.rs @@ -290,7 +290,7 @@ mod tests { pub expr: ::core::option::Option, } - #[derive(Clone, PartialEq, ::prost::Message)] + #[derive(Clone, Eq, PartialEq, ::prost::Message)] pub struct TopKExecProto { #[prost(uint64, tag = "1")] pub k: u64, diff --git a/ballista/rust/core/src/serde/scheduler/to_proto.rs b/ballista/rust/core/src/serde/scheduler/to_proto.rs index 4dbf5d9678..815bc96ddf 100644 --- a/ballista/rust/core/src/serde/scheduler/to_proto.rs +++ b/ballista/rust/core/src/serde/scheduler/to_proto.rs @@ -99,12 +99,10 @@ pub fn hash_partitioning_to_proto( })) } None => Ok(None), - other => { - return Err(BallistaError::General(format!( - "scheduler::to_proto() invalid partitioning for ExecutePartition: {:?}", - other - ))) - } + other => Err(BallistaError::General(format!( + "scheduler::to_proto() invalid partitioning for ExecutePartition: {:?}", + other + ))), } } diff --git a/ballista/rust/scheduler/src/flight_sql.rs b/ballista/rust/scheduler/src/flight_sql.rs index 37704f588e..1a7275b0fc 100644 --- a/ballista/rust/scheduler/src/flight_sql.rs +++ b/ballista/rust/scheduler/src/flight_sql.rs @@ -243,17 +243,6 @@ impl FlightSqlServiceImpl { plan: &LogicalPlan, ) -> Result { let job_id = self.server.state.task_manager.generate_job_id(); - self.server - .state - .task_manager - .queue_job(&job_id) - .await - .map_err(|e| { - let msg = format!("Failed to queue job {}: {:?}", job_id, e); - error!("{}", msg); - - Status::internal(msg) - })?; let query_stage_event_sender = self .server .query_stage_event_loop diff --git a/ballista/rust/scheduler/src/scheduler_server/grpc.rs b/ballista/rust/scheduler/src/scheduler_server/grpc.rs index f80da4fab9..9ef898ac2e 100644 --- a/ballista/rust/scheduler/src/scheduler_server/grpc.rs +++ b/ballista/rust/scheduler/src/scheduler_server/grpc.rs @@ -413,17 +413,6 @@ impl SchedulerGrpc let job_id = self.state.task_manager.generate_job_id(); - self.state - .task_manager - .queue_job(&job_id) - .await - .map_err(|e| { - let msg = format!("Failed to queue job {}: {:?}", job_id, e); - error!("{}", msg); - - Status::internal(msg) - })?; - let query_stage_event_sender = self.query_stage_event_loop.get_sender().map_err(|e| { Status::internal(format!( diff --git a/ballista/rust/scheduler/src/state/backend/mod.rs b/ballista/rust/scheduler/src/state/backend/mod.rs index 4a6334abe1..b69403b2ea 100644 --- a/ballista/rust/scheduler/src/state/backend/mod.rs +++ b/ballista/rust/scheduler/src/state/backend/mod.rs @@ -54,7 +54,6 @@ pub enum Keyspace { Executors, ActiveJobs, CompletedJobs, - QueuedJobs, FailedJobs, Slots, Sessions, @@ -118,7 +117,7 @@ pub trait Watch: Stream + Send + Unpin { async fn cancel(&mut self) -> Result<()>; } -#[derive(Debug, PartialEq)] +#[derive(Debug, Eq, PartialEq)] pub enum WatchEvent { /// Contains the inserted or updated key and the new value Put(String, Vec), diff --git a/ballista/rust/scheduler/src/state/task_manager.rs b/ballista/rust/scheduler/src/state/task_manager.rs index ab9ff26e7a..cc32926373 100644 --- a/ballista/rust/scheduler/src/state/task_manager.rs +++ b/ballista/rust/scheduler/src/state/task_manager.rs @@ -28,8 +28,8 @@ use ballista_core::serde::protobuf::executor_grpc_client::ExecutorGrpcClient; use crate::state::session_manager::create_datafusion_context; use ballista_core::serde::protobuf::{ - self, job_status, task_status, FailedJob, JobStatus, PartitionId, QueuedJob, - TaskDefinition, TaskStatus, + self, job_status, task_status, FailedJob, JobStatus, PartitionId, TaskDefinition, + TaskStatus, }; use ballista_core::serde::scheduler::to_proto::hash_partitioning_to_proto; use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionLocation}; @@ -88,30 +88,12 @@ impl TaskManager ) .await?; - if let Err(e) = self.state.delete(Keyspace::QueuedJobs, job_id).await { - warn!("Failed to remove key in QueuedJobs for {}: {:?}", job_id, e); - } - Ok(()) } - /// Queue a job. When a batch job is submitted we do the physical planning asynchronously so we - /// need to add a marker so we can report on its status. - pub async fn queue_job(&self, job_id: &str) -> Result<()> { - self.state - .put(Keyspace::QueuedJobs, job_id.to_owned(), vec![0x0]) - .await - } - - /// Get the status of of a job. First look in Active/Completed jobs, and then in Queued jobs, and - /// finally in FailedJobs. + /// Get the status of of a job. First look in Active/Completed jobs, and then in Failed jobs pub async fn get_job_status(&self, job_id: &str) -> Result> { - let queue_marker = self.state.get(Keyspace::QueuedJobs, job_id).await?; - if !queue_marker.is_empty() { - Ok(Some(JobStatus { - status: Some(job_status::Status::Queued(QueuedJob {})), - })) - } else if let Ok(graph) = self.get_execution_graph(job_id).await { + if let Ok(graph) = self.get_execution_graph(job_id).await { Ok(Some(graph.status())) } else { let value = self.state.get(Keyspace::FailedJobs, job_id).await?; @@ -401,8 +383,6 @@ impl TaskManager let lock = self.state.lock(Keyspace::ActiveJobs, "").await?; with_lock(lock, self.state.delete(Keyspace::ActiveJobs, job_id)).await?; - self.state.delete(Keyspace::QueuedJobs, job_id).await?; - let status = JobStatus { status: Some(job_status::Status::Failed(FailedJob { error: error_message,