This is an automated email from the ASF dual-hosted git repository.
nju_yaho pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-ballista.git
The following commit(s) were added to refs/heads/master by this push:
new 12128a30 Add LaunchMultiTask rpc interface for executor (#255)
12128a30 is described below
commit 12128a30a51315ae348b6a01f3ec42d0a4b41dd5
Author: yahoNanJing <[email protected]>
AuthorDate: Thu Sep 22 09:21:16 2022 +0800
Add LaunchMultiTask rpc interface for executor (#255)
Co-authored-by: yangzhong <[email protected]>
---
ballista/rust/core/proto/ballista.proto | 30 +++++++
.../rust/core/src/serde/scheduler/from_proto.rs | 96 ++++++++++++++++++++--
ballista/rust/core/src/serde/scheduler/mod.rs | 19 ++++-
ballista/rust/core/src/serde/scheduler/to_proto.rs | 43 +++++++++-
ballista/rust/executor/src/executor_server.rs | 84 ++++++++++++-------
5 files changed, 230 insertions(+), 42 deletions(-)
diff --git a/ballista/rust/core/proto/ballista.proto
b/ballista/rust/core/proto/ballista.proto
index 7b7f748c..d998cc7b 100644
--- a/ballista/rust/core/proto/ballista.proto
+++ b/ballista/rust/core/proto/ballista.proto
@@ -534,6 +534,13 @@ message PartitionId {
uint32 partition_id = 4;
}
+// Multiple partitions in the same stage
+message PartitionIds {
+ string job_id = 1;
+ uint32 stage_id = 2;
+ repeated uint32 partition_ids = 4;
+}
+
message PartitionStats {
int64 num_rows = 1;
int64 num_batches = 2;
@@ -698,6 +705,16 @@ message TaskDefinition {
repeated KeyValuePair props = 5;
}
+// A set of tasks in the same stage
+message MultiTaskDefinition {
+ PartitionIds task_ids = 1;
+ bytes plan = 2;
+ // Output partition for shuffle writer
+ PhysicalHashRepartition output_partitioning = 3;
+ string session_id = 4;
+ repeated KeyValuePair props = 5;
+}
+
message SessionSettings {
repeated KeyValuePair configs = 1;
}
@@ -836,11 +853,22 @@ message LaunchTaskParams {
string scheduler_id = 2;
}
+message LaunchMultiTaskParams {
+ // Allow to launch a task set to an executor at once
+ repeated MultiTaskDefinition multi_tasks = 1;
+ string scheduler_id = 2;
+}
+
message LaunchTaskResult {
bool success = 1;
// TODO when part of the task set are scheduled successfully
}
+message LaunchMultiTaskResult {
+ bool success = 1;
+ // TODO when part of the task set are scheduled successfully
+}
+
message CancelTasksParams {
repeated PartitionId partition_id = 1;
}
@@ -876,6 +904,8 @@ service SchedulerGrpc {
service ExecutorGrpc {
rpc LaunchTask (LaunchTaskParams) returns (LaunchTaskResult) {}
+ rpc LaunchMultiTask (LaunchMultiTaskParams) returns (LaunchMultiTaskResult)
{}
+
rpc StopExecutor (StopExecutorParams) returns (StopExecutorResult) {}
rpc CancelTasks (CancelTasksParams) returns (CancelTasksResult) {}
diff --git a/ballista/rust/core/src/serde/scheduler/from_proto.rs
b/ballista/rust/core/src/serde/scheduler/from_proto.rs
index 728ad46b..536e2a5c 100644
--- a/ballista/rust/core/src/serde/scheduler/from_proto.rs
+++ b/ballista/rust/core/src/serde/scheduler/from_proto.rs
@@ -20,6 +20,7 @@ use datafusion::physical_plan::metrics::{
Count, Gauge, MetricValue, MetricsSet, Time, Timestamp,
};
use datafusion::physical_plan::Metric;
+use std::collections::HashMap;
use std::convert::TryInto;
use std::sync::Arc;
use std::time::Duration;
@@ -30,7 +31,7 @@ use crate::serde::protobuf::action::ActionType;
use crate::serde::protobuf::{operator_metric, NamedCount, NamedGauge,
NamedTime};
use crate::serde::scheduler::{
Action, ExecutorData, ExecutorMetadata, ExecutorSpecification, PartitionId,
- PartitionLocation, PartitionStats,
+ PartitionIds, PartitionLocation, PartitionStats, TaskDefinition,
};
impl TryInto<Action> for protobuf::Action {
@@ -51,15 +52,14 @@ impl TryInto<Action> for protobuf::Action {
}
}
-impl TryInto<PartitionId> for protobuf::PartitionId {
- type Error = BallistaError;
-
- fn try_into(self) -> Result<PartitionId, Self::Error> {
- Ok(PartitionId::new(
+#[allow(clippy::from_over_into)]
+impl Into<PartitionId> for protobuf::PartitionId {
+ fn into(self) -> PartitionId {
+ PartitionId::new(
&self.job_id,
self.stage_id as usize,
self.partition_id as usize,
- ))
+ )
}
}
@@ -94,7 +94,7 @@ impl TryInto<PartitionLocation> for
protobuf::PartitionLocation {
"partition_id in PartitionLocation is
missing.".to_owned(),
)
})?
- .try_into()?,
+ .into(),
executor_meta: self
.executor_meta
.ok_or_else(|| {
@@ -263,3 +263,83 @@ impl Into<ExecutorData> for protobuf::ExecutorData {
ret
}
}
+
+#[allow(clippy::from_over_into)]
+impl Into<PartitionIds> for protobuf::PartitionIds {
+ fn into(self) -> PartitionIds {
+ PartitionIds {
+ job_id: self.job_id.clone(),
+ stage_id: self.stage_id as usize,
+ partition_ids: self
+ .partition_ids
+ .into_iter()
+ .map(|partition_id| partition_id as usize)
+ .collect(),
+ }
+ }
+}
+
+impl TryInto<TaskDefinition> for protobuf::TaskDefinition {
+ type Error = BallistaError;
+
+ fn try_into(self) -> Result<TaskDefinition, Self::Error> {
+ let mut props = HashMap::new();
+ for kv_pair in self.props {
+ props.insert(kv_pair.key, kv_pair.value);
+ }
+
+ let task_id = self
+ .task_id
+ .ok_or_else(|| {
+ BallistaError::General("No task id in the
TaskDefinition".to_owned())
+ })?
+ .into();
+
+ Ok(TaskDefinition {
+ task_id,
+ plan: self.plan,
+ output_partitioning: self.output_partitioning,
+ session_id: self.session_id,
+ props,
+ })
+ }
+}
+
+impl TryInto<Vec<TaskDefinition>> for protobuf::MultiTaskDefinition {
+ type Error = BallistaError;
+
+ fn try_into(self) -> Result<Vec<TaskDefinition>, Self::Error> {
+ let mut props = HashMap::new();
+ for kv_pair in self.props {
+ props.insert(kv_pair.key, kv_pair.value);
+ }
+
+ let plan = self.plan;
+ let output_partitioning = self.output_partitioning;
+ let session_id = self.session_id;
+ let task_ids: PartitionIds = self
+ .task_ids
+ .ok_or_else(|| {
+ BallistaError::General(
+ "No task ids in the MultiTaskDefinition".to_owned(),
+ )
+ })?
+ .into();
+
+ Ok(task_ids
+ .partition_ids
+ .iter()
+ .map(|partition_id| TaskDefinition {
+ task_id: PartitionId::new(
+ &task_ids.job_id,
+ task_ids.stage_id,
+ *partition_id,
+ ),
+ plan: plan.clone(),
+ output_partitioning: output_partitioning.clone(),
+ session_id: session_id.clone(),
+ props: props.clone(),
+ })
+ .collect())
+ }
+}
diff --git a/ballista/rust/core/src/serde/scheduler/mod.rs
b/ballista/rust/core/src/serde/scheduler/mod.rs
index 5afdb30e..c1fa78a2 100644
--- a/ballista/rust/core/src/serde/scheduler/mod.rs
+++ b/ballista/rust/core/src/serde/scheduler/mod.rs
@@ -21,12 +21,12 @@ use datafusion::arrow::array::{
ArrayBuilder, StructArray, StructBuilder, UInt64Array, UInt64Builder,
};
use datafusion::arrow::datatypes::{DataType, Field};
-
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::Partitioning;
use serde::Serialize;
use crate::error::BallistaError;
+use crate::serde::protobuf::PhysicalHashRepartition;
pub mod from_proto;
pub mod to_proto;
@@ -268,3 +268,20 @@ impl ExecutePartitionResult {
&self.stats
}
}
+
+/// Unique identifier for the output partitions of a set of tasks.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct PartitionIds {
+ pub job_id: String,
+ pub stage_id: usize,
+ pub partition_ids: Vec<usize>,
+}
+
+#[derive(Debug, Clone)]
+pub struct TaskDefinition {
+ pub task_id: PartitionId,
+ pub plan: Vec<u8>,
+ pub output_partitioning: Option<PhysicalHashRepartition>,
+ pub session_id: String,
+ pub props: HashMap<String, String>,
+}
diff --git a/ballista/rust/core/src/serde/scheduler/to_proto.rs
b/ballista/rust/core/src/serde/scheduler/to_proto.rs
index 10c841e9..5a7ad394 100644
--- a/ballista/rust/core/src/serde/scheduler/to_proto.rs
+++ b/ballista/rust/core/src/serde/scheduler/to_proto.rs
@@ -22,10 +22,12 @@ 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::protobuf::{
+ operator_metric, KeyValuePair, NamedCount, NamedGauge, NamedTime,
+};
use crate::serde::scheduler::{
Action, ExecutorData, ExecutorMetadata, ExecutorSpecification, PartitionId,
- PartitionLocation, PartitionStats,
+ PartitionIds, PartitionLocation, PartitionStats, TaskDefinition,
};
use datafusion::physical_plan::Partitioning;
@@ -234,3 +236,40 @@ impl Into<protobuf::ExecutorData> for ExecutorData {
}
}
}
+
+#[allow(clippy::from_over_into)]
+impl Into<protobuf::PartitionIds> for PartitionIds {
+ fn into(self) -> protobuf::PartitionIds {
+ protobuf::PartitionIds {
+ job_id: self.job_id,
+ stage_id: self.stage_id as u32,
+ partition_ids: self
+ .partition_ids
+ .into_iter()
+ .map(|partition_id| partition_id as u32)
+ .collect(),
+ }
+ }
+}
+
+#[allow(clippy::from_over_into)]
+impl Into<protobuf::TaskDefinition> for TaskDefinition {
+ fn into(self) -> protobuf::TaskDefinition {
+ let props = self
+ .props
+ .iter()
+ .map(|(k, v)| KeyValuePair {
+ key: k.to_owned(),
+ value: v.to_owned(),
+ })
+ .collect::<Vec<_>>();
+
+ protobuf::TaskDefinition {
+ task_id: Some(self.task_id.into()),
+ plan: self.plan,
+ output_partitioning: self.output_partitioning,
+ session_id: self.session_id,
+ props,
+ }
+ }
+}
diff --git a/ballista/rust/executor/src/executor_server.rs
b/ballista/rust/executor/src/executor_server.rs
index 7bde7f7f..4544d8c2 100644
--- a/ballista/rust/executor/src/executor_server.rs
+++ b/ballista/rust/executor/src/executor_server.rs
@@ -35,10 +35,11 @@ use ballista_core::serde::protobuf::executor_grpc_server::{
use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
use ballista_core::serde::protobuf::{
executor_metric, executor_status, CancelTasksParams, CancelTasksResult,
- ExecutorMetric, ExecutorStatus, HeartBeatParams, LaunchTaskParams,
LaunchTaskResult,
- RegisterExecutorParams, StopExecutorParams, StopExecutorResult,
TaskDefinition,
- TaskStatus, UpdateTaskStatusParams,
+ ExecutorMetric, ExecutorStatus, HeartBeatParams, LaunchMultiTaskParams,
+ LaunchMultiTaskResult, LaunchTaskParams, LaunchTaskResult,
RegisterExecutorParams,
+ StopExecutorParams, StopExecutorResult, TaskStatus, UpdateTaskStatusParams,
};
+use ballista_core::serde::scheduler::TaskDefinition;
use ballista_core::serde::{AsExecutionPlan, BallistaCodec};
use ballista_core::utils::{
collect_plan_metrics, create_grpc_client_connection, create_grpc_server,
@@ -286,7 +287,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorServer<T,
) -> Result<(), BallistaError> {
let scheduler_id = curator_task.scheduler_id;
let task = curator_task.task;
- let task_id = task.task_id.unwrap();
+ let task_id = task.task_id;
let task_id_log = format!(
"{}/{}/{}",
task_id.job_id, task_id.stage_id, task_id.partition_id
@@ -295,10 +296,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorServer<T,
let runtime = self.executor.runtime.clone();
let session_id = task.session_id;
- let mut task_props = HashMap::new();
- for kv_pair in task.props {
- task_props.insert(kv_pair.key, kv_pair.value);
- }
+ let task_props = task.props;
let mut task_scalar_functions = HashMap::new();
let mut task_aggregate_functions = HashMap::new();
@@ -364,7 +362,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorServer<T,
let task_status = as_task_status(
execution_result,
executor_id.clone(),
- task_id,
+ task_id.into(),
Some(operator_metrics),
);
@@ -551,28 +549,22 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> TaskRunnerPool<T,
}
};
if let Some(curator_task) = maybe_task {
- if let Some(task_id) = &curator_task.task.task_id {
- let task_id_log = format!(
- "{}/{}/{}",
- task_id.job_id, task_id.stage_id,
task_id.partition_id
- );
- info!("Received task {:?}", &task_id_log);
-
- let server = executor_server.clone();
- dedicated_executor.spawn(async move {
-
server.run_task(curator_task).await.unwrap_or_else(|e| {
- error!(
- "Fail to run the task {:?} due to {:?}",
- task_id_log, e
- );
- });
+ let task_id = &curator_task.task.task_id;
+ let task_id_log = format!(
+ "{}/{}/{}",
+ task_id.job_id, task_id.stage_id, task_id.partition_id
+ );
+ info!("Received task {:?}", &task_id_log);
+
+ let server = executor_server.clone();
+ dedicated_executor.spawn(async move {
+ server.run_task(curator_task).await.unwrap_or_else(|e|
{
+ error!(
+ "Fail to run the task {:?} due to {:?}",
+ task_id_log, e
+ );
});
- } else {
- error!(
- "There's no task id in the task definition {:?}",
- curator_task
- );
- }
+ });
} else {
info!("Channel is closed and will exit the task receive
loop");
drop(task_runner_complete);
@@ -600,7 +592,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorGrpc
task_sender
.send(CuratorTaskDefinition {
scheduler_id: scheduler_id.clone(),
- task,
+ task: task
+ .try_into()
+ .map_err(|e| Status::invalid_argument(format!("{}",
e)))?,
})
.await
.unwrap();
@@ -608,6 +602,34 @@ impl<T: 'static + AsLogicalPlan, U: 'static +
AsExecutionPlan> ExecutorGrpc
Ok(Response::new(LaunchTaskResult { success: true }))
}
+ /// by this interface, it can reduce the deserialization cost for multiple
tasks
+ /// belong to the same job stage running on the same one executor
+ async fn launch_multi_task(
+ &self,
+ request: Request<LaunchMultiTaskParams>,
+ ) -> Result<Response<LaunchMultiTaskResult>, Status> {
+ let LaunchMultiTaskParams {
+ multi_tasks,
+ scheduler_id,
+ } = request.into_inner();
+ let task_sender = self.executor_env.tx_task.clone();
+ for multi_task in multi_tasks {
+ let multi_task: Vec<TaskDefinition> = multi_task
+ .try_into()
+ .map_err(|e| Status::invalid_argument(format!("{}", e)))?;
+ for task in multi_task {
+ task_sender
+ .send(CuratorTaskDefinition {
+ scheduler_id: scheduler_id.clone(),
+ task,
+ })
+ .await
+ .unwrap();
+ }
+ }
+ Ok(Response::new(LaunchMultiTaskResult { success: true }))
+ }
+
async fn stop_executor(
&self,
request: Request<StopExecutorParams>,
