feat: add BALLISTA_PROTOCOL_VERSION + k8s health probes

ballista/core/proto/ballista.proto

@@ -381,6 +381,11 @@ message ExecutorRegistration {
   uint32 grpc_port = 4;
   ExecutorSpecification specification = 5;
   ExecutorOperatingSystemSpecification os_info = 6;
+  // Wire-protocol version this executor was built against. The scheduler
+  // rejects heartbeats/registrations whose version does not match its own
+  // compiled-in BALLISTA_PROTOCOL_VERSION. Old executors that predate this
+  // field default to 0, which is never a valid scheduler version.
+  uint32 ballista_protocol_version = 7;
 }
 
 message ExecutorHeartbeat {

ballista/core/src/lib.rs

@@ -29,6 +29,18 @@ pub use crate::ids::JobId;
 /// The current version of Ballista, derived from the Cargo package version.
 pub const BALLISTA_VERSION: &str = env!("CARGO_PKG_VERSION");
 
+/// Wire-protocol version negotiated between scheduler and executor.
+///
+/// The scheduler compares this against `ExecutorRegistration.ballista_protocol_version`
+/// on every registration and heartbeat; mismatched executors are rejected with
+/// `Status::failed_precondition` and never receive work. Bump this on any release
+/// that changes the executor↔scheduler wire format (proto shape, task/plan
+/// encoding, launch/heartbeat semantics). Most releases will not bump it.
+///
+/// Zero is reserved as the proto-default "unset" value, produced by executors
+/// that predate this field — it never matches a real scheduler version.
+pub const BALLISTA_PROTOCOL_VERSION: u32 = 1;
+
 /// Prints the current Ballista version to stdout.
 pub fn print_version() {
     println!("Ballista version: {BALLISTA_VERSION}")

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

@@ -575,6 +575,12 @@ pub struct ExecutorRegistration {
     pub specification: ::core::option::Option<ExecutorSpecification>,
     #[prost(message, optional, tag = "6")]
     pub os_info: ::core::option::Option<ExecutorOperatingSystemSpecification>,
+    /// Wire-protocol version this executor was built against. The scheduler
+    /// rejects heartbeats/registrations whose version does not match its own
+    /// compiled-in BALLISTA_PROTOCOL_VERSION. Old executors that predate this
+    /// field default to 0, which is never a valid scheduler version.
+    #[prost(uint32, tag = "7")]
+    pub ballista_protocol_version: u32,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ExecutorHeartbeat {

ballista/executor/Cargo.toml

@@ -34,14 +34,15 @@ required-features = ["build-binary"]
 
 [features]
 arrow-ipc-optimizations = []
-build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary", "mimalloc"]
+build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary", "mimalloc", "dep:axum"]
 default = ["arrow-ipc-optimizations", "build-binary"]
 spark-compat = ["ballista-core/spark-compat"]
 
 [dependencies]
 arrow = { workspace = true }
 arrow-flight = { workspace = true }
 async-trait = { workspace = true }
+axum = { version = "0.8.9", optional = true }
 ballista-core = { path = "../core", version = "54.0.0" }
 bytesize = "2"
 clap = { workspace = true, optional = true }

ballista/executor/src/bin/main.rs

@@ -18,15 +18,18 @@
 //! Ballista Rust executor binary.
 
 use ballista_core::config::LogRotationPolicy;
+use ballista_core::error::BallistaError;
 use ballista_core::object_store::{
     runtime_env_with_s3_support, session_config_with_s3_support,
 };
 use ballista_executor::config::Config;
 use ballista_executor::executor_process::{
     ExecutorProcessConfig, start_executor_process,
 };
+use ballista_executor::health::spawn_health_server;
 use clap::Parser;
 use std::env;
+use std::net::SocketAddr;
 use std::sync::Arc;
 use tracing_subscriber::EnvFilter;
 
@@ -38,6 +41,8 @@ static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 async fn main() -> ballista_core::error::Result<()> {
     // parse command-line arguments
     let opt = Config::parse();
+    let bind_host = opt.bind_host.clone();
+    let bind_health_port = opt.bind_health_port;
 
     let mut config: ExecutorProcessConfig = opt.try_into()?;
     config.override_config_producer = Some(Arc::new(session_config_with_s3_support));
@@ -76,5 +81,25 @@ async fn main() -> ballista_core::error::Result<()> {
         tracing.init();
     }
 
-    start_executor_process(Arc::new(config)).await
+    // Kubernetes-style health probes are a concern of the standalone binary,
+    // not of the library, so wire them here on top of the config's shared
+    // ExecutorHealth handle. `/healthz` is process-liveness, `/readyz`
+    // reflects heartbeat state (see `ballista_executor::health`).
+    let health = config.health.clone();
+    let health_addr: SocketAddr = format!("{bind_host}:{bind_health_port}")
+        .parse()
+        .map_err(|e| {
+            BallistaError::General(format!("invalid --bind-health-port address: {e}"))
+        })?;
+    let (health_shutdown_tx, health_shutdown_rx) = tokio::sync::oneshot::channel();
+    let health_handle = spawn_health_server(health_addr, health, health_shutdown_rx);
+
+    let result = start_executor_process(Arc::new(config)).await;
+
+    // Ask the health server to shut down so the process can exit cleanly.
+    let _ = health_shutdown_tx.send(());
+    if let Err(e) = health_handle.await {
+        log::warn!("health server task join error: {e}");
+    }
+    result
 }

ballista/executor/src/config.rs

@@ -71,6 +71,13 @@ pub struct Config {
     /// Port for the executor's gRPC service (used for task management).
     #[arg(long, default_value_t = 50052, help = "Grpc service bind port.")]
     pub bind_grpc_port: u16,
+    /// Port for the executor's HTTP health server (/healthz and /readyz).
+    #[arg(
+        long,
+        default_value_t = 50053,
+        help = "HTTP port for the executor's Kubernetes health probes (/healthz and /readyz)."
+    )]
+    pub bind_health_port: u16,
     /// Timeout in seconds for establishing connection to scheduler (0 = fail immediately).
     #[arg(
         long,
@@ -250,6 +257,7 @@ impl TryFrom<Config> for ExecutorProcessConfig {
             override_arrow_flight_service: None,
             override_create_grpc_client_endpoint: None,
             client_ttl: opt.client_ttl,
+            health: crate::health::ExecutorHealth::new(),
         })
     }
 }

ballista/executor/src/execution_loop.rs

@@ -61,6 +61,7 @@ pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan,
     mut scheduler: SchedulerGrpcClient<C>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
+    health: crate::health::ExecutorHealth,
 ) -> Result<(), BallistaError>
 where
     C: tonic::client::GrpcService<tonic::body::Body>,
@@ -108,6 +109,7 @@ where
 
         match poll_work_result {
             Ok(result) => {
+                health.mark_heartbeat_ok();
                 let PollWorkResult {
                     tasks,
                     jobs_to_clean,
@@ -198,6 +200,7 @@ where
                 }
             }
             Err(error) => {
+                health.mark_heartbeat_failed();
                 warn!(
                     "Executor poll work loop failed. If this continues to happen the Scheduler might be marked as dead. Error: {error}"
                 );

ballista/executor/src/executor.rs

@@ -436,11 +436,7 @@ mod test {
 
         let executor_registration = ExecutorRegistration {
             id: "executor".to_string(),
-            port: 0,
-            grpc_port: 0,
-            specification: None,
-            host: None,
-            os_info: None,
+            ..Default::default()
         };
         let config_producer = Arc::new(default_config_producer);
         let ctx = SessionContext::new();
@@ -501,11 +497,7 @@ mod test {
     fn produce_runtime_for_session_shares_read_side_state() {
         let executor_registration = ExecutorRegistration {
             id: "executor".to_string(),
-            port: 0,
-            grpc_port: 0,
-            specification: None,
-            host: None,
-            os_info: None,
+            ..Default::default()
         };
         let config_producer = Arc::new(default_config_producer);
 
@@ -540,11 +532,7 @@ mod test {
     fn produce_runtime_for_session_falls_back_without_cache() {
         let executor_registration = ExecutorRegistration {
             id: "executor".to_string(),
-            port: 0,
-            grpc_port: 0,
-            specification: None,
-            host: None,
-            os_info: None,
+            ..Default::default()
         };
         let config_producer = Arc::new(default_config_producer);
         let base_producer: RuntimeProducer =

ballista/executor/src/executor_process.rs

@@ -61,7 +61,9 @@ use ballista_core::utils::{
     GrpcClientConfig, GrpcServerConfig, create_grpc_client_endpoint, create_grpc_server,
     default_config_producer, get_time_before,
 };
-use ballista_core::{BALLISTA_VERSION, ConfigProducer, JobId, RuntimeProducer};
+use ballista_core::{
+    BALLISTA_PROTOCOL_VERSION, BALLISTA_VERSION, ConfigProducer, JobId, RuntimeProducer,
+};
 
 use crate::client_pool::DefaultBallistaClientPool;
 use crate::execution_engine::{DefaultExecutionEngine, ExecutionEngine};
@@ -192,6 +194,11 @@ pub struct ExecutorProcessConfig {
     /// connection and a shuffle-heavy query can exhaust the host's ephemeral
     /// ports.
     pub client_ttl: u64,
+    /// Shared readiness state that the heartbeat loops flip on every RPC
+    /// outcome. Embedders leave this as `Default::default()` and never
+    /// observe it; the standalone binary passes a handle here and also spawns
+    /// an HTTP server on it (see `bin/main.rs`).
+    pub health: crate::health::ExecutorHealth,
 }
 
 impl ExecutorProcessConfig {
@@ -242,6 +249,7 @@ impl Default for ExecutorProcessConfig {
             override_arrow_flight_service: None,
             override_create_grpc_client_endpoint: None,
             client_ttl: 30,
+            health: crate::health::ExecutorHealth::new(),
         }
     }
 }
@@ -519,6 +527,12 @@ pub async fn start_executor_process(
     // Channels used to receive stop requests from Executor grpc service.
     let (stop_send, mut stop_recv) = mpsc::channel::<bool>(10);
 
+    // Shared readiness state, flipped by the heartbeat/poll_work loop. When
+    // the caller is the standalone binary, an HTTP probe server observes
+    // this same handle (see `bin/main.rs`); library embedders leave it
+    // unobserved.
+    let health = opt.health.clone();
+
     // Starting main executor process based on the TaskSchedulingPolicy
     //
     // PushStaged => starting new executor_server that waits for tasks from the schedule
@@ -534,6 +548,7 @@ pub async fn start_executor_process(
                     default_codec,
                     stop_send,
                     &shutdown_notification,
+                    health,
                 )
                 .await?,
             );
@@ -543,6 +558,7 @@ pub async fn start_executor_process(
                 scheduler.clone(),
                 executor.clone(),
                 default_codec,
+                health,
             )));
         }
     };
@@ -939,6 +955,7 @@ pub fn structure_executor_metadata(
             total_available_disk_space,
             open_files_limit,
         }),
+        ballista_protocol_version: BALLISTA_PROTOCOL_VERSION,
     }
 }