ballista_executor/
lib.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18#![doc = include_str!("../README.md")]
19#![warn(missing_docs)]
20
21/// Execution plan for collecting distributed query results into a single partition.
22pub mod collect;
23/// Command-line configuration for the executor binary.
24#[cfg(feature = "build-binary")]
25pub mod config;
26/// Extension point for custom query stage execution engines.
27pub mod execution_engine;
28/// Pull-based task execution loop that polls the scheduler for work.
29pub mod execution_loop;
30/// Core executor implementation for running distributed query tasks.
31pub mod executor;
32/// Executor process lifecycle management and configuration.
33pub mod executor_process;
34/// gRPC server for receiving pushed tasks from the scheduler.
35pub mod executor_server;
36/// Arrow Flight service for streaming shuffle data between executors.
37pub mod flight_service;
38/// Metrics collection for executor runtime statistics.
39pub mod metrics;
40/// Graceful shutdown coordination for executor components.
41pub mod shutdown;
42/// Signal handling for process termination.
43pub mod terminate;
44
45mod cpu_bound_executor;
46mod standalone;
47
48use ballista_core::error::BallistaError;
49use std::net::SocketAddr;
50
51pub use standalone::new_standalone_executor;
52pub use standalone::new_standalone_executor_from_builder;
53pub use standalone::new_standalone_executor_from_state;
54
55use log::info;
56
57use crate::shutdown::Shutdown;
58use ballista_core::serde::protobuf::{
59    FailedTask, OperatorMetricsSet, ShuffleWritePartition, SuccessfulTask, TaskStatus,
60    task_status,
61};
62use ballista_core::serde::scheduler::PartitionId;
63use ballista_core::utils::GrpcServerConfig;
64
65/// [ArrowFlightServerProvider] provides a function which creates a new Arrow Flight server.
66///
67/// The function should take two arguments:
68/// [SocketAddr] - the address to bind the server to
69/// [Shutdown] - a shutdown signal to gracefully shutdown the server
70/// [GrpcServerConfig] - the gRPC server configuration for timeout settings
71/// Returns a [tokio::task::JoinHandle] which will be registered as service handler
72///
73pub type ArrowFlightServerProvider = dyn Fn(
74        SocketAddr,
75        Shutdown,
76        GrpcServerConfig,
77    ) -> tokio::task::JoinHandle<Result<(), BallistaError>>
78    + Send
79    + Sync;
80
81/// Timestamps capturing the lifecycle of a task execution.
82#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
83pub struct TaskExecutionTimes {
84    /// Timestamp when the task was launched by the scheduler (milliseconds since epoch).
85    launch_time: u64,
86    /// Timestamp when task execution started on the executor (milliseconds since epoch).
87    start_exec_time: u64,
88    /// Timestamp when task execution completed (milliseconds since epoch).
89    end_exec_time: u64,
90}
91
92/// Converts a task execution result into a [`TaskStatus`] protobuf message.
93///
94/// This function wraps the outcome of task execution (success or failure)
95/// along with timing and metrics information into a status message that
96/// can be sent back to the scheduler.
97pub fn as_task_status(
98    execution_result: ballista_core::error::Result<Vec<ShuffleWritePartition>>,
99    executor_id: String,
100    task_id: usize,
101    stage_attempt_num: usize,
102    partition_id: PartitionId,
103    operator_metrics: Option<Vec<OperatorMetricsSet>>,
104    execution_times: TaskExecutionTimes,
105) -> TaskStatus {
106    let metrics = operator_metrics.unwrap_or_default();
107    match execution_result {
108        Ok(partitions) => {
109            info!(
110                "Task {:?} finished with operator_metrics array size {}",
111                task_id,
112                metrics.len()
113            );
114            TaskStatus {
115                task_id: task_id as u32,
116                job_id: partition_id.job_id,
117                stage_id: partition_id.stage_id as u32,
118                stage_attempt_num: stage_attempt_num as u32,
119                partition_id: partition_id.partition_id as u32,
120                launch_time: execution_times.launch_time,
121                start_exec_time: execution_times.start_exec_time,
122                end_exec_time: execution_times.end_exec_time,
123                metrics,
124                status: Some(task_status::Status::Successful(SuccessfulTask {
125                    executor_id,
126                    partitions,
127                })),
128            }
129        }
130        Err(e) => {
131            let error_msg = e.to_string();
132            info!("Task {task_id:?} failed: {error_msg}");
133
134            TaskStatus {
135                task_id: task_id as u32,
136                job_id: partition_id.job_id,
137                stage_id: partition_id.stage_id as u32,
138                stage_attempt_num: stage_attempt_num as u32,
139                partition_id: partition_id.partition_id as u32,
140                launch_time: execution_times.launch_time,
141                start_exec_time: execution_times.start_exec_time,
142                end_exec_time: execution_times.end_exec_time,
143                metrics,
144                status: Some(task_status::Status::Failed(FailedTask::from(e))),
145            }
146        }
147    }
148}
ballista_executor/lib.rs

ballista_executor/
lib.rs
