Ballista Executor report plan/operators metrics to Ballista Scheduler

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 {

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<Action> for protobuf::Action {
@@ -104,3 +112,93 @@ impl TryInto<PartitionLocation> for protobuf::PartitionLocation {
         })
     }
 }
+
+impl TryInto<MetricValue> for protobuf::OperatorMetric {
+    type Error = BallistaError;
+
+    fn try_into(self) -> Result<MetricValue, Self::Error> {
+        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<MetricsSet> for protobuf::OperatorMetricsSet {
+    type Error = BallistaError;
+
+    fn try_into(self) -> Result<MetricsSet, Self::Error> {
+        let mut ms = MetricsSet::new();
+        let metrics = self
+            .metrics
+            .into_iter()
+            .map(|m| m.try_into())
+            .collect::<Result<Vec<_>, BallistaError>>()?;
+
+        for value in metrics {
+            let new_metric = Arc::new(Metric::new(value, None));
+            ms.push(new_metric)
+        }
+        Ok(ms)
+    }
+}

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<protobuf::OperatorMetric> for &MetricValue {
+    type Error = BallistaError;
+
+    fn try_into(self) -> Result<protobuf::OperatorMetric, Self::Error> {
+        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<protobuf::OperatorMetricsSet> for MetricsSet {
+    type Error = BallistaError;
+
+    fn try_into(self) -> Result<protobuf::OperatorMetricsSet, Self::Error> {
+        let metrics = self
+            .iter()
+            .map(|m| m.value().try_into())
+            .collect::<Result<Vec<_>, BallistaError>>()?;
+        Ok(protobuf::OperatorMetricsSet { metrics })
+    }
+}

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<dyn ExecutionPlan>) -> Vec<MetricsSet> {
+    let mut metrics_array = Vec::<MetricsSet>::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
+}

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<T: 'static + AsLogicalPlan, U: 'static + AsExecution
             task_id.job_id.clone(),
             task_id.stage_id as usize,
             task_id.partition_id as usize,
-            plan,
+            plan.clone(),
             task_context,
             shuffle_output_partitioning,
         ))
@@ -209,10 +211,18 @@ async fn run_received_tasks<T: 'static + AsLogicalPlan, U: 'static + AsExecution
         debug!("Statistics: {:?}", execution_result);
         available_tasks_slots.fetch_add(1, Ordering::SeqCst);
 
+        let plan_metrics = collect_plan_metrics(&plan);
+        let operator_metrics = plan_metrics
+            .into_iter()
+            .map(|m| m.try_into())
+            .collect::<Result<Vec<_>, BallistaError>>()
+            .ok();
+
         let _ = task_status_sender.send(as_task_status(
             execution_result,
             executor.metadata.id.clone(),
             task_id,
+            operator_metrics,
         ));
     });
 

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,16 +240,26 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
                 task_id.job_id.clone(),
                 task_id.stage_id as usize,
                 task_id.partition_id as usize,
-                plan,
+                plan.clone(),
                 task_context,
                 shuffle_output_partitioning,
             )
             .await;
         info!("Done with task {}", task_id_log);
         debug!("Statistics: {:?}", execution_result);
 
+        let plan_metrics = collect_plan_metrics(&plan);
+        let operator_metrics = plan_metrics
+            .into_iter()
+            .map(|m| m.try_into())
+            .collect::<Result<Vec<_>, 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();