You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/10/03 13:30:11 UTC
(datafusion-ballista) branch main updated: feat: error when the client and scheduler have different major versions (#2514)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 9ee325c9f feat: error when the client and scheduler have different major versions (#2514)
9ee325c9f is described below
commit 9ee325c9fea6a4345be2f403f2e8023e99b91eb0
Author: Andy Grove <ag...@apache.org>
AuthorDate: Sat Oct 3 07:30:05 2026 -0600
feat: error when the client and scheduler have different major versions (#2514)
* feat: error when the client and scheduler have different major versions
A client and a scheduler from different major versions can run a query
and return wrong results, because DataFusion's plan encoding can change
between majors in ways that still decode. Until now a mismatch only
logged a warning.
The client now sends its version in a `ballista-version` gRPC header
with `ExecuteQuery` and `ExecuteQueryPush`. The scheduler rejects a
different major version with `failed_precondition` before it decodes the
plan, and sends its own version back. A side that sends no header is
treated as older than 55.0.0, so 54.x clients are rejected once the
scheduler is on 55.0.0.
The client also checks the scheduler's reply. A scheduler older than
55.0.0 doesn't check versions and has already queued the job, so the
client cancels it before returning the error.
Closes #2513.
* docs: state that clients and schedulers stay compatible within a major version
The version check treats the major version as the unit of compatibility, so a change that would break an older client of the same major version has to wait for the next major release. Say so next to the check, next to `BALLISTA_PROTOCOL_VERSION` (which can change in any release), and in the upgrade guide.
* fix: remove redundant rustdoc link targets
---
.../core/src/execution_plans/distributed_query.rs | 378 +++++++++++++++++++--
ballista/core/src/lib.rs | 7 +
ballista/core/src/version.rs | 161 +++++++++
ballista/scheduler/src/scheduler_server/grpc.rs | 188 +++++++++-
docs/source/upgrading/55.0.0.md | 35 ++
examples/examples/standalone-substrait.rs | 10 +-
6 files changed, 743 insertions(+), 36 deletions(-)
diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs
index adc573f9c..b8adec0af 100644
--- a/ballista/core/src/execution_plans/distributed_query.rs
+++ b/ballista/core/src/execution_plans/distributed_query.rs
@@ -19,15 +19,18 @@ use crate::JobId;
use crate::client::BallistaClient;
use crate::config::BallistaConfig;
use crate::error::BallistaError;
-use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt};
+use crate::extension::{
+ BallistaConfigGrpcEndpoint, BallistaGrpcMetadataInterceptor, SessionConfigExt,
+};
use crate::serde::protobuf::get_job_status_result::FlightProxy;
use crate::serde::protobuf::{
- ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult, KeyValuePair,
- PartitionLocation, execute_query_params::Query, execute_query_result, job_status,
- scheduler_grpc_client::SchedulerGrpcClient,
+ CancelJobParams, ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult,
+ KeyValuePair, PartitionLocation, execute_query_params::Query, execute_query_result,
+ job_status, scheduler_grpc_client::SchedulerGrpcClient,
};
use crate::serde::protobuf::{ExecutorMetadata, SuccessfulJob};
use crate::utils::{GrpcClientConfig, create_grpc_client_endpoint};
+use crate::version;
use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::tree_node::TreeNodeRecursion;
@@ -55,6 +58,9 @@ use std::fmt::Debug;
use std::marker::PhantomData;
use std::sync::Arc;
use std::time::Duration;
+use tonic::metadata::MetadataMap;
+use tonic::service::interceptor::InterceptedService;
+use tonic::transport::Channel;
use url::Url;
/// This operator sends a logical plan to a Ballista scheduler for execution and
@@ -451,20 +457,47 @@ pub async fn execute_physical_plan<U: 'static + AsExecutionPlan>(
Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
}
-/// Logs a warning if the scheduler's advertised version (from the `server`
-/// response header, e.g. `ballista/1.2.3`) doesn't match this client's.
-fn check_scheduler_version<T>(response: &tonic::Response<T>) {
- let expected = format!("ballista/{}", crate::BALLISTA_VERSION);
- if let Some(server) = response
- .metadata()
- .get("server")
- .and_then(|v| v.to_str().ok())
- && server != expected
+type SchedulerClient =
+ SchedulerGrpcClient<InterceptedService<Channel, BallistaGrpcMetadataInterceptor>>;
+
+/// Wraps a job submission with this client's version, which schedulers
+/// require from 55.0.0 on. See [`crate::version`].
+fn job_submission_request(
+ query: ExecuteQueryParams,
+) -> tonic::Request<ExecuteQueryParams> {
+ let mut request = tonic::Request::new(query);
+ version::insert_version_header(request.metadata_mut());
+ request
+}
+
+/// Checks the version the scheduler reported in its response to a job
+/// submission. A different major version is an error, and a different minor
+/// or patch version only logs a warning.
+fn check_scheduler_version(client_version: &str, metadata: &MetadataMap) -> Result<()> {
+ let scheduler_version = version::version_from_metadata(metadata);
+ version::check_compatibility(Some(client_version), scheduler_version)
+ .map_err(DataFusionError::Execution)?;
+ if let Some(scheduler_version) = scheduler_version
+ && scheduler_version != client_version
{
warn!(
- "Scheduler version mismatch: scheduler reports '{server}', client is '{expected}'"
+ "Scheduler version {scheduler_version} differs from client version {client_version}"
);
}
+ Ok(())
+}
+
+/// Cancels a job whose results the client won't read. A failure is only
+/// logged, since the caller is already returning an error.
+async fn cancel_job(scheduler: &mut SchedulerClient, job_id: &str) {
+ if let Err(e) = scheduler
+ .cancel_job(CancelJobParams {
+ job_id: job_id.to_owned(),
+ })
+ .await
+ {
+ warn!("Failed to cancel job {job_id}: {e}");
+ }
}
/// Client will periodically invoke scheduler to check
@@ -517,21 +550,31 @@ async fn execute_query_pull(
.max_decoding_message_size(max_message_size);
let query_response = scheduler
- .execute_query(query)
+ .execute_query(job_submission_request(query))
.await
.map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
- check_scheduler_version(&query_response);
+ let version_check =
+ check_scheduler_version(crate::BALLISTA_VERSION, query_response.metadata());
let query_result = query_response.into_inner();
let query_result = match query_result.result.unwrap() {
execute_query_result::Result::Success(success_result) => success_result,
execute_query_result::Result::Failure(failure_result) => {
+ // A version mismatch is the likelier cause, so report it first.
+ version_check?;
return Err(DataFusionError::Execution(format!(
"Fail to execute query due to {failure_result:?}"
)));
}
};
+ if let Err(e) = version_check {
+ // A scheduler older than 55.0.0 doesn't check versions, so it has
+ // already queued the job.
+ cancel_job(&mut scheduler, &query_result.job_id).await;
+ return Err(e);
+ }
+
assert_eq!(
session_id, query_result.session_id,
"Session id inconsistent between Client and Server side in DistributedQueryExec."
@@ -690,12 +733,30 @@ async fn execute_query_push(
.max_decoding_message_size(max_message_size);
let query_push_response = scheduler
- .execute_query_push(query)
+ .execute_query_push(job_submission_request(query))
.await
.map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
- check_scheduler_version(&query_push_response);
+ let version_check =
+ check_scheduler_version(crate::BALLISTA_VERSION, query_push_response.metadata());
let mut query_status_stream = query_push_response.into_inner();
+ if let Err(e) = version_check {
+ // A scheduler older than 55.0.0 doesn't check versions, so it has
+ // already queued the job. The job id only arrives with the first
+ // status update, which can take until planning finishes, so cancel
+ // the job in the background rather than hold up the error.
+ tokio::spawn(async move {
+ if let Some(Ok(GetJobStatusResult {
+ status: Some(status),
+ ..
+ })) = query_status_stream.next().await
+ {
+ cancel_job(&mut scheduler, &status.job_id).await;
+ }
+ });
+ return Err(e);
+ }
+
let mut prev_status: Option<job_status::Status> = None;
loop {
@@ -907,20 +968,295 @@ async fn fetch_partition(
#[cfg(test)]
mod test {
- use crate::JobId;
use crate::config::BallistaConfig;
use crate::execution_plans::distributed_query::{
- DistributedQueryExec, get_client_host_port,
+ DistributedQueryExec, check_scheduler_version, execute_query_pull,
+ execute_query_push, get_client_host_port, job_submission_request,
};
- use crate::serde::protobuf::ExecutorMetadata;
+ use crate::extension::SessionConfigExt;
use crate::serde::protobuf::get_job_status_result::FlightProxy;
+ use crate::serde::protobuf::scheduler_grpc_server::{
+ SchedulerGrpc, SchedulerGrpcServer,
+ };
+ use crate::serde::protobuf::{
+ CancelJobParams, CancelJobResult, CleanJobDataParams, CleanJobDataResult,
+ CreateUpdateSessionParams, CreateUpdateSessionResult, ExecuteQueryParams,
+ ExecuteQueryResult, ExecuteQuerySuccessResult, ExecutorMetadata,
+ ExecutorStoppedParams, ExecutorStoppedResult, GetJobMetricsParams,
+ GetJobMetricsResult, GetJobStatusParams, GetJobStatusResult, HeartBeatParams,
+ HeartBeatResult, JobStatus, PollWorkParams, PollWorkResult,
+ RegisterExecutorParams, RegisterExecutorResult, RemoveSessionParams,
+ RemoveSessionResult, UpdateTaskStatusParams, UpdateTaskStatusResult,
+ execute_query_result,
+ };
+ use crate::utils::GrpcClientConfig;
+ use crate::{BALLISTA_VERSION, JobId};
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::displayable;
+ use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
use datafusion::prelude::SessionConfig;
use datafusion_proto::protobuf::LogicalPlanNode;
+ use futures::Stream;
+ use parking_lot::Mutex;
+ use std::pin::Pin;
use std::sync::Arc;
+ use std::time::{Duration, Instant};
+ use tokio_stream::wrappers::TcpListenerStream;
+ use tonic::metadata::MetadataMap;
+ use tonic::{Request, Response, Status};
+
+ #[test]
+ fn job_submission_request_carries_client_version() {
+ let request = job_submission_request(ExecuteQueryParams::default());
+
+ assert_eq!(
+ request
+ .metadata()
+ .get("ballista-version")
+ .and_then(|v| v.to_str().ok()),
+ Some(BALLISTA_VERSION)
+ );
+ }
+
+ #[test]
+ fn scheduler_version_is_read_from_response_header() {
+ let mut metadata = MetadataMap::new();
+ metadata.insert("ballista-version", "55.3.0".parse().unwrap());
+
+ assert!(check_scheduler_version("55.0.0", &metadata).is_ok());
+ assert!(check_scheduler_version("56.0.0", &metadata).is_err());
+ }
+
+ #[test]
+ fn scheduler_without_version_header_is_incompatible() {
+ // Schedulers older than 55.0.0 don't report their version.
+ assert!(check_scheduler_version("55.0.0", &MetadataMap::new()).is_err());
+ }
+
+ /// A scheduler that accepts every job without checking the client's
+ /// version, reports version 1.0.0, and records the jobs it is asked to
+ /// cancel.
+ #[derive(Clone, Default)]
+ struct IncompatibleScheduler {
+ cancelled: Arc<Mutex<Vec<String>>>,
+ }
+
+ impl IncompatibleScheduler {
+ /// Serves the scheduler on a local port and returns its URL.
+ async fn start(&self) -> String {
+ let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let url = format!("http://{}", listener.local_addr().unwrap());
+ tokio::spawn(
+ tonic::transport::Server::builder()
+ .add_service(SchedulerGrpcServer::new(self.clone()))
+ .serve_with_incoming(TcpListenerStream::new(listener)),
+ );
+ url
+ }
+
+ fn cancelled(&self) -> Vec<String> {
+ self.cancelled.lock().clone()
+ }
+ }
+
+ fn reporting_version_1<T>(message: T) -> Response<T> {
+ let mut response = Response::new(message);
+ response
+ .metadata_mut()
+ .insert("ballista-version", "1.0.0".parse().unwrap());
+ response
+ }
+
+ type GrpcResult<T> = std::result::Result<Response<T>, Status>;
+
+ fn not_called<T>() -> GrpcResult<T> {
+ Err(Status::unimplemented("not called before the version check"))
+ }
+
+ #[tonic::async_trait]
+ impl SchedulerGrpc for IncompatibleScheduler {
+ type ExecuteQueryPushStream = Pin<
+ Box<
+ dyn Stream<Item = std::result::Result<GetJobStatusResult, Status>> + Send,
+ >,
+ >;
+
+ async fn execute_query(
+ &self,
+ request: Request<ExecuteQueryParams>,
+ ) -> GrpcResult<ExecuteQueryResult> {
+ let session_id = request.into_inner().session_id;
+ Ok(reporting_version_1(ExecuteQueryResult {
+ operation_id: String::new(),
+ result: Some(execute_query_result::Result::Success(
+ ExecuteQuerySuccessResult {
+ job_id: "pull-job".to_owned(),
+ session_id,
+ },
+ )),
+ }))
+ }
+
+ async fn execute_query_push(
+ &self,
+ _request: Request<ExecuteQueryParams>,
+ ) -> GrpcResult<Self::ExecuteQueryPushStream> {
+ let first_status = GetJobStatusResult {
+ status: Some(JobStatus {
+ job_id: "push-job".to_owned(),
+ ..Default::default()
+ }),
+ flight_proxy: None,
+ };
+ Ok(reporting_version_1(Box::pin(futures::stream::iter([Ok(
+ first_status,
+ )]))))
+ }
+
+ async fn cancel_job(
+ &self,
+ request: Request<CancelJobParams>,
+ ) -> GrpcResult<CancelJobResult> {
+ self.cancelled.lock().push(request.into_inner().job_id);
+ Ok(Response::new(CancelJobResult { cancelled: true }))
+ }
+
+ async fn poll_work(
+ &self,
+ _: Request<PollWorkParams>,
+ ) -> GrpcResult<PollWorkResult> {
+ not_called()
+ }
+
+ async fn register_executor(
+ &self,
+ _: Request<RegisterExecutorParams>,
+ ) -> GrpcResult<RegisterExecutorResult> {
+ not_called()
+ }
+
+ async fn heart_beat_from_executor(
+ &self,
+ _: Request<HeartBeatParams>,
+ ) -> GrpcResult<HeartBeatResult> {
+ not_called()
+ }
+
+ async fn update_task_status(
+ &self,
+ _: Request<UpdateTaskStatusParams>,
+ ) -> GrpcResult<UpdateTaskStatusResult> {
+ not_called()
+ }
+
+ async fn create_update_session(
+ &self,
+ _: Request<CreateUpdateSessionParams>,
+ ) -> GrpcResult<CreateUpdateSessionResult> {
+ not_called()
+ }
+
+ async fn remove_session(
+ &self,
+ _: Request<RemoveSessionParams>,
+ ) -> GrpcResult<RemoveSessionResult> {
+ not_called()
+ }
+
+ async fn get_job_status(
+ &self,
+ _: Request<GetJobStatusParams>,
+ ) -> GrpcResult<GetJobStatusResult> {
+ not_called()
+ }
+
+ async fn get_job_metrics(
+ &self,
+ _: Request<GetJobMetricsParams>,
+ ) -> GrpcResult<GetJobMetricsResult> {
+ not_called()
+ }
+
+ async fn executor_stopped(
+ &self,
+ _: Request<ExecutorStoppedParams>,
+ ) -> GrpcResult<ExecutorStoppedResult> {
+ not_called()
+ }
+
+ async fn clean_job_data(
+ &self,
+ _: Request<CleanJobDataParams>,
+ ) -> GrpcResult<CleanJobDataResult> {
+ not_called()
+ }
+ }
+
+ fn job_submission() -> ExecuteQueryParams {
+ ExecuteQueryParams {
+ session_id: "session".to_owned(),
+ ..Default::default()
+ }
+ }
+
+ #[tokio::test]
+ async fn pull_cancels_job_accepted_by_incompatible_scheduler() {
+ let scheduler = IncompatibleScheduler::default();
+ let url = scheduler.start().await;
+ let config = BallistaConfig::default();
+
+ let result = execute_query_pull(
+ url,
+ "session".to_owned(),
+ job_submission(),
+ config.grpc_client_max_message_size(),
+ GrpcClientConfig::from(&config),
+ Arc::new(ExecutionPlanMetricsSet::new()),
+ Arc::new(Mutex::new(None)),
+ 0,
+ SessionConfig::new_with_ballista(),
+ )
+ .await;
+
+ let Err(err) = result else {
+ panic!("an incompatible scheduler should be an error");
+ };
+ assert!(err.to_string().contains("scheduler 1.0.0"), "{err}");
+ assert_eq!(scheduler.cancelled(), ["pull-job"]);
+ }
+
+ #[tokio::test]
+ async fn push_cancels_job_accepted_by_incompatible_scheduler() {
+ let scheduler = IncompatibleScheduler::default();
+ let url = scheduler.start().await;
+ let config = BallistaConfig::default();
+
+ let result = execute_query_push(
+ url,
+ job_submission(),
+ config.grpc_client_max_message_size(),
+ GrpcClientConfig::from(&config),
+ Arc::new(ExecutionPlanMetricsSet::new()),
+ Arc::new(Mutex::new(None)),
+ 0,
+ SessionConfig::new_with_ballista(),
+ )
+ .await;
+
+ let Err(err) = result else {
+ panic!("an incompatible scheduler should be an error");
+ };
+ assert!(err.to_string().contains("scheduler 1.0.0"), "{err}");
+
+ // The job is cancelled in the background once its id arrives.
+ let deadline = Instant::now() + Duration::from_secs(5);
+ while scheduler.cancelled().is_empty() && Instant::now() < deadline {
+ tokio::time::sleep(Duration::from_millis(10)).await;
+ }
+ assert_eq!(scheduler.cancelled(), ["push-job"]);
+ }
#[test]
fn test_client_host_port() {
diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs
index ea52654cf..6c7bc0ff1 100644
--- a/ballista/core/src/lib.rs
+++ b/ballista/core/src/lib.rs
@@ -39,6 +39,11 @@ pub const BALLISTA_VERSION: &str = env!("CARGO_PKG_VERSION");
///
/// Zero is reserved as the proto-default "unset" value, produced by executors
/// that predate this field — it never matches a real scheduler version.
+///
+/// Clients are not covered by this check. They are checked by major version
+/// instead (see [`crate::version`]), so a change that also reaches clients,
+/// such as the shuffle fetch they use to read results, must keep working for
+/// older clients of the same major version.
pub const BALLISTA_PROTOCOL_VERSION: u32 = 3;
/// Prints the current Ballista version to stdout.
@@ -97,6 +102,8 @@ pub mod serde;
pub mod sort_key;
/// General utility functions for Ballista operations.
pub mod utils;
+/// Version checks between Ballista clients and schedulers.
+pub mod version;
///
/// [RuntimeProducer] is a factory which creates runtime [RuntimeEnv]
diff --git a/ballista/core/src/version.rs b/ballista/core/src/version.rs
new file mode 100644
index 000000000..d9ef3ac51
--- /dev/null
+++ b/ballista/core/src/version.rs
@@ -0,0 +1,161 @@
+// 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.
+
+//! Version checks between Ballista clients and schedulers.
+//!
+//! Each Ballista major version moves to a new DataFusion major version, and
+//! DataFusion's plan encoding can change between majors in ways that still
+//! decode but mean something different. A client and a scheduler with
+//! different major versions could run a query and return wrong results, so
+//! they refuse to work together.
+//!
+//! The client sends its [`BALLISTA_VERSION`] in the
+//! [`BALLISTA_VERSION_HEADER`](crate::version::BALLISTA_VERSION_HEADER) gRPC
+//! header when it submits a job. The scheduler rejects the submission if the
+//! major versions differ, and otherwise sends its own version back in the
+//! response, so the client can also detect a scheduler that is too old to
+//! check.
+//!
+//! The major version is the unit of compatibility. Any client works with any
+//! scheduler of the same major version, whatever their minor and patch
+//! versions, so a change that would break an older client of the same major
+//! version has to wait for the next major release. That covers everything a
+//! client relies on: job submission, job status, and the shuffle fetch it uses
+//! to read results from executors.
+//! [`BALLISTA_PROTOCOL_VERSION`] is
+//! different. Schedulers and executors are upgraded together, so it can change
+//! in any release.
+
+use tonic::metadata::{MetadataMap, MetadataValue};
+
+use crate::BALLISTA_VERSION;
+
+/// gRPC metadata key that carries a Ballista version between a client and a
+/// scheduler.
+pub const BALLISTA_VERSION_HEADER: &str = "ballista-version";
+
+/// The first major version whose clients and schedulers send
+/// [`BALLISTA_VERSION_HEADER`]. A peer that doesn't send it is older.
+pub const FIRST_MAJOR_WITH_VERSION_HEADER: u64 = 55;
+
+/// Checks that a client and a scheduler can work together.
+///
+/// Returns an error message naming both versions when their major versions
+/// differ. `None` means that side didn't send [`BALLISTA_VERSION_HEADER`], so
+/// it is older than [`FIRST_MAJOR_WITH_VERSION_HEADER`]. That is a mismatch
+/// once the other side is on that major version or later.
+pub fn check_compatibility(
+ client: Option<&str>,
+ scheduler: Option<&str>,
+) -> Result<(), String> {
+ let compatible = match (client, scheduler) {
+ (Some(client), Some(scheduler)) => {
+ let client_major = major_version(client);
+ client_major.is_some() && client_major == major_version(scheduler)
+ }
+ (Some(version), None) | (None, Some(version)) => major_version(version)
+ .is_some_and(|major| major < FIRST_MAJOR_WITH_VERSION_HEADER),
+ (None, None) => true,
+ };
+ if compatible {
+ Ok(())
+ } else {
+ Err(format!(
+ "Ballista {} cannot be used with {}. The client and the scheduler must \
+ have the same major version.",
+ describe("client", client),
+ describe("scheduler", scheduler),
+ ))
+ }
+}
+
+/// Adds this build's [`BALLISTA_VERSION`] to `metadata` as
+/// [`BALLISTA_VERSION_HEADER`].
+pub fn insert_version_header(metadata: &mut MetadataMap) {
+ metadata.insert(
+ BALLISTA_VERSION_HEADER,
+ MetadataValue::from_static(BALLISTA_VERSION),
+ );
+}
+
+/// Returns the version a peer sent in [`BALLISTA_VERSION_HEADER`], if any.
+pub fn version_from_metadata(metadata: &MetadataMap) -> Option<&str> {
+ metadata
+ .get(BALLISTA_VERSION_HEADER)
+ .and_then(|value| value.to_str().ok())
+}
+
+fn major_version(version: &str) -> Option<u64> {
+ version.split('.').next()?.parse().ok()
+}
+
+fn describe(role: &str, version: Option<&str>) -> String {
+ match version {
+ Some(version) => format!("{role} {version}"),
+ None => format!(
+ "{role} older than {FIRST_MAJOR_WITH_VERSION_HEADER}.0.0 \
+ (no `{BALLISTA_VERSION_HEADER}` header)"
+ ),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::check_compatibility;
+
+ #[test]
+ fn compares_major_versions() {
+ // (client, scheduler, compatible)
+ let cases = [
+ (Some("55.0.0"), Some("55.0.0"), true),
+ (Some("55.0.0"), Some("55.2.1"), true),
+ (Some("54.1.0"), Some("55.0.0"), false),
+ (Some("56.0.0"), Some("55.3.0"), false),
+ // A side that sends no version is older than 55.0.0.
+ (None, Some("55.0.0"), false),
+ (None, Some("56.1.0"), false),
+ (Some("55.0.0"), None, false),
+ // Before 55.0.0 a missing version can't be told apart from a
+ // matching one.
+ (None, Some("54.0.0"), true),
+ (Some("54.0.0"), None, true),
+ (Some("banana"), Some("55.0.0"), false),
+ (Some("55.0.0"), Some(""), false),
+ (Some("dev"), Some("dev"), false),
+ ];
+ for (client, scheduler, compatible) in cases {
+ assert_eq!(
+ check_compatibility(client, scheduler).is_ok(),
+ compatible,
+ "client={client:?} scheduler={scheduler:?}"
+ );
+ }
+ }
+
+ #[test]
+ fn mismatch_names_both_versions() {
+ let err = check_compatibility(Some("54.1.0"), Some("55.0.0")).unwrap_err();
+ assert!(err.contains("client 54.1.0"), "{err}");
+ assert!(err.contains("scheduler 55.0.0"), "{err}");
+ }
+
+ #[test]
+ fn missing_version_names_the_header() {
+ let err = check_compatibility(None, Some("55.0.0")).unwrap_err();
+ assert!(err.contains("ballista-version"), "{err}");
+ }
+}
diff --git a/ballista/scheduler/src/scheduler_server/grpc.rs b/ballista/scheduler/src/scheduler_server/grpc.rs
index c1be5c075..9bd357adc 100644
--- a/ballista/scheduler/src/scheduler_server/grpc.rs
+++ b/ballista/scheduler/src/scheduler_server/grpc.rs
@@ -16,7 +16,6 @@
// under the License.
use axum::extract::ConnectInfo;
-use ballista_core::BALLISTA_PROTOCOL_VERSION;
use ballista_core::config::BALLISTA_JOB_NAME;
use ballista_core::error::{BallistaError, Result as BResult};
use ballista_core::extension::SessionConfigHelperExt;
@@ -35,6 +34,8 @@ use ballista_core::serde::protobuf::{
execute_query_failure_result, execute_query_result, executor_metric::Metric,
};
use ballista_core::serde::scheduler::ExecutorMetadata;
+use ballista_core::version;
+use ballista_core::{BALLISTA_PROTOCOL_VERSION, BALLISTA_VERSION};
use datafusion_proto::logical_plan::AsLogicalPlan;
use datafusion_proto::physical_plan::AsExecutionPlan;
use futures::{Stream, StreamExt};
@@ -59,6 +60,7 @@ use datafusion::prelude::SessionContext;
use std::ops::Deref;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
+use tonic::metadata::MetadataMap;
use tonic::{Request, Response, Status};
use crate::metrics::record_protocol_mismatch;
@@ -82,6 +84,24 @@ fn check_protocol_version(metadata: &ExecutorRegistration) -> Result<(), Status>
)))
}
+/// Rejects a job submission from a client whose major version differs from
+/// the scheduler's. See [`ballista_core::version`].
+fn check_client_version(metadata: &MetadataMap) -> Result<(), Status> {
+ let client_version = version::version_from_metadata(metadata);
+ version::check_compatibility(client_version, Some(BALLISTA_VERSION)).map_err(|msg| {
+ info!("Rejecting job submission: {msg}");
+ Status::failed_precondition(msg)
+ })
+}
+
+/// Wraps a response to a job submission, adding the scheduler's version so
+/// the client can check it.
+fn versioned_response<R>(message: R) -> Response<R> {
+ let mut response = Response::new(message);
+ version::insert_version_header(response.metadata_mut());
+ response
+}
+
fn executor_registration_error(e: BallistaError) -> Status {
let msg = format!("Fail to do executor registration due to: {e}");
error!("{msg}");
@@ -427,6 +447,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
request: tonic::Request<ExecuteQueryParams>,
) -> std::result::Result<tonic::Response<Self::ExecuteQueryPushStream>, tonic::Status>
{
+ check_client_version(request.metadata())?;
let query_params = request.into_inner();
if let ExecuteQueryParams {
query: Some(query),
@@ -491,7 +512,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
"execution query (PUSH) job submitted - session_id: {session_id}, operation_id: {operation_id}, job_name: {job_name}, job_id: {job_id}"
);
- Ok(Response::new(Box::pin(stream)))
+ Ok(versioned_response(Box::pin(stream)))
} else {
Err(Status::internal(
"Error processing request, invalid message",
@@ -503,6 +524,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
&self,
request: Request<ExecuteQueryParams>,
) -> Result<Response<ExecuteQueryResult>, Status> {
+ check_client_version(request.metadata())?;
let query_params = request.into_inner();
if let ExecuteQueryParams {
query: Some(query),
@@ -533,7 +555,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
Err(e) => {
let msg = format!("Could not parse plan: {e}");
error!("{msg}");
- return Ok(Response::new(ExecuteQueryResult {
+ return Ok(versioned_response(ExecuteQueryResult {
operation_id,
result: Some(execute_query_result::Result::Failure(
ExecuteQueryFailureResult {
@@ -565,7 +587,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
"execution query, job submitted - session_id: {session_id}, operation_id: {operation_id}, job_name: {job_name}"
);
- Ok(Response::new(ExecuteQueryResult {
+ Ok(versioned_response(ExecuteQueryResult {
operation_id,
result: Some(execute_query_result::Result::Success(
ExecuteQuerySuccessResult {
@@ -906,30 +928,33 @@ mod test {
use std::sync::Arc;
use std::time::Duration;
+ use datafusion::prelude::SessionContext;
+ use datafusion_proto::logical_plan::{AsLogicalPlan, DefaultLogicalExtensionCodec};
use datafusion_proto::protobuf::LogicalPlanNode;
use datafusion_proto::protobuf::PhysicalPlanNode;
- use tonic::{Code, Request};
+ use tonic::{Code, Request, Response};
#[cfg(feature = "substrait")]
use {
- ballista_core::serde::protobuf::ExecuteQueryParams,
- ballista_core::serde::protobuf::execute_query_params::Query,
- datafusion::prelude::{SessionConfig, SessionContext},
+ ballista_core::version::insert_version_header,
+ datafusion::prelude::SessionConfig,
datafusion_substrait::serializer::serialize_bytes,
};
use crate::config::SchedulerConfig;
use crate::metrics::default_metrics_collector;
- use ballista_core::BALLISTA_PROTOCOL_VERSION;
use ballista_core::error::BallistaError;
use ballista_core::serde::BallistaCodec;
+ use ballista_core::serde::protobuf::execute_query_params::Query;
use ballista_core::serde::protobuf::{
- ExecutorRegistration, ExecutorStatus, ExecutorStoppedParams, HeartBeatParams,
- PollWorkParams, RegisterExecutorParams, executor_status,
+ ExecuteQueryParams, ExecutorRegistration, ExecutorStatus, ExecutorStoppedParams,
+ HeartBeatParams, PollWorkParams, RegisterExecutorParams, execute_query_result,
+ executor_status,
};
use ballista_core::serde::scheduler::{
ExecutorOperatingSystemSpecification, ExecutorSpecification,
};
+ use ballista_core::{BALLISTA_PROTOCOL_VERSION, BALLISTA_VERSION};
use crate::state::SchedulerState;
use crate::test_utils::await_condition;
@@ -1464,6 +1489,144 @@ mod test {
);
}
+ /// Bytes that don't decode as a plan. A submission carrying them fails in
+ /// plan parsing, unless the version check rejects it first.
+ const UNPARSEABLE_PLAN: [u8; 4] = [0xFF; 4];
+
+ async fn test_scheduler()
+ -> Result<SchedulerServer<LogicalPlanNode, PhysicalPlanNode>, BallistaError> {
+ let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+ SchedulerServer::new(
+ "localhost:50050".to_owned(),
+ test_cluster_context(),
+ BallistaCodec::default(),
+ Arc::new(SchedulerConfig::default()),
+ default_metrics_collector().unwrap(),
+ );
+ scheduler.init().await?;
+ Ok(scheduler)
+ }
+
+ /// `SELECT 1`, encoded the way a client submits a logical plan.
+ async fn select_one() -> Vec<u8> {
+ let plan = SessionContext::new()
+ .sql("SELECT 1")
+ .await
+ .unwrap()
+ .into_unoptimized_plan();
+ let mut buf = vec![];
+ LogicalPlanNode::try_from_logical_plan(&plan, &DefaultLogicalExtensionCodec {})
+ .unwrap()
+ .try_encode(&mut buf)
+ .unwrap();
+ buf
+ }
+
+ /// A job submission for `plan` from a client reporting `client_version`.
+ fn job_submission(
+ plan: Vec<u8>,
+ client_version: &str,
+ ) -> Request<ExecuteQueryParams> {
+ let mut request = Request::new(ExecuteQueryParams {
+ query: Some(Query::LogicalPlan(plan)),
+ settings: vec![],
+ session_id: uuid::Uuid::new_v4().to_string(),
+ operation_id: uuid::Uuid::now_v7().to_string(),
+ });
+ request
+ .metadata_mut()
+ .insert("ballista-version", client_version.parse().unwrap());
+ request
+ }
+
+ fn reported_version<T>(response: &Response<T>) -> Option<&str> {
+ response
+ .metadata()
+ .get("ballista-version")
+ .and_then(|v| v.to_str().ok())
+ }
+
+ #[tokio::test]
+ async fn execute_query_rejects_client_with_different_major_version()
+ -> Result<(), BallistaError> {
+ let scheduler = test_scheduler().await?;
+
+ let request = job_submission(UNPARSEABLE_PLAN.to_vec(), "1.0.0");
+ let status = scheduler.execute_query(request).await.unwrap_err();
+
+ assert_eq!(status.code(), Code::FailedPrecondition);
+ assert!(
+ status.message().contains("client 1.0.0"),
+ "{}",
+ status.message()
+ );
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn execute_query_push_rejects_client_with_different_major_version()
+ -> Result<(), BallistaError> {
+ let scheduler = test_scheduler().await?;
+
+ let request = job_submission(UNPARSEABLE_PLAN.to_vec(), "1.0.0");
+ let Err(status) = scheduler.execute_query_push(request).await else {
+ panic!("a client with a different major version should be rejected");
+ };
+
+ assert_eq!(status.code(), Code::FailedPrecondition);
+ assert!(
+ status.message().contains("client 1.0.0"),
+ "{}",
+ status.message()
+ );
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn execute_query_reports_scheduler_version() -> Result<(), BallistaError> {
+ let scheduler = test_scheduler().await?;
+
+ let request = job_submission(select_one().await, BALLISTA_VERSION);
+ let response = scheduler.execute_query(request).await?;
+
+ assert_eq!(reported_version(&response), Some(BALLISTA_VERSION));
+ assert!(matches!(
+ response.into_inner().result,
+ Some(execute_query_result::Result::Success(_))
+ ));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn execute_query_reports_scheduler_version_when_plan_fails_to_parse()
+ -> Result<(), BallistaError> {
+ // The client checks the version before it looks at the result.
+ let scheduler = test_scheduler().await?;
+
+ let request = job_submission(UNPARSEABLE_PLAN.to_vec(), BALLISTA_VERSION);
+ let response = scheduler.execute_query(request).await?;
+
+ assert_eq!(reported_version(&response), Some(BALLISTA_VERSION));
+ assert!(matches!(
+ response.into_inner().result,
+ Some(execute_query_result::Result::Failure(_))
+ ));
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn execute_query_push_reports_scheduler_version() -> Result<(), BallistaError> {
+ let scheduler = test_scheduler().await?;
+
+ let request = job_submission(select_one().await, BALLISTA_VERSION);
+ let Ok(response) = scheduler.execute_query_push(request).await else {
+ panic!("a client with the same version should be accepted");
+ };
+
+ assert_eq!(reported_version(&response), Some(BALLISTA_VERSION));
+ Ok(())
+ }
+
#[tokio::test]
#[cfg(feature = "substrait")]
async fn test_substrait_compatibility() -> Result<(), BallistaError> {
@@ -1528,12 +1691,13 @@ mod test {
)
.await?;
- let execute_query_request = Request::new(ExecuteQueryParams {
+ let mut execute_query_request = Request::new(ExecuteQueryParams {
session_id: uuid::Uuid::new_v4().to_string(),
settings: vec![],
operation_id: uuid::Uuid::now_v7().to_string(),
query: Some(Query::SubstraitPlan(serialized_substrait_plan)),
});
+ insert_version_header(execute_query_request.metadata_mut());
let response = scheduler.execute_query(execute_query_request).await?;
response
.into_inner()
diff --git a/docs/source/upgrading/55.0.0.md b/docs/source/upgrading/55.0.0.md
index 774c334c7..15a7254fe 100644
--- a/docs/source/upgrading/55.0.0.md
+++ b/docs/source/upgrading/55.0.0.md
@@ -155,6 +155,41 @@ rejected heartbeat as a heartbeat failure and shuts itself down after five
consecutive ones, so a stale executor is restarted by its orchestrator until its
image is bumped to match rather than sitting idle in a ready state.
+#### Clients and schedulers must have the same major version
+
+Each Ballista major version moves to a new DataFusion major version, and
+DataFusion's plan encoding can change between majors in ways that still decode
+but mean something different. A client and a scheduler from different major
+versions could run a query and return wrong results, so from `55.0.0` they
+refuse to work together. This applies to the Python client too: the `ballista`
+package on PyPI must have the same major version as the cluster.
+
+A `55.0.0` client sends its version in a `ballista-version` gRPC header when it
+submits a job, and the scheduler rejects the job with
+`Status::failed_precondition` if the major versions differ. A `54.x` client
+sends no version, so a `55.0.0` scheduler rejects every job it submits:
+
+```text
+Ballista client older than 55.0.0 (no `ballista-version` header) cannot be used
+with scheduler 55.0.0. The client and the scheduler must have the same major
+version.
+```
+
+The scheduler sends its own version back in the same header. A `55.0.0` client
+that gets no version back is talking to a `54.x` scheduler, which has already
+queued the job without checking. The client cancels that job and returns an
+error. A different minor or patch version is not an error, and the client only
+logs a warning.
+
+Ballista keeps clients and schedulers of the same major version compatible. Any
+`55.x` client works with any `55.x` cluster, and a change that would break an
+older client waits for the next major release.
+
+If you wrote your own client against the scheduler's gRPC API, it must send the
+header on every `ExecuteQuery` and `ExecuteQueryPush` request, set to the
+Ballista version it was built against. For a `tonic` request,
+`ballista_core::version::insert_version_header` does this.
+
### API changes
#### `SchedulerConfig` now carries a scheduler identifier
diff --git a/examples/examples/standalone-substrait.rs b/examples/examples/standalone-substrait.rs
index da35b7cf1..cf0674333 100644
--- a/examples/examples/standalone-substrait.rs
+++ b/examples/examples/standalone-substrait.rs
@@ -26,6 +26,7 @@ use ballista_core::serde::protobuf::{
};
use ballista_core::serde::scheduler::ShuffleLayout;
use ballista_core::utils::{GrpcClientConfig, create_grpc_client_connection};
+use ballista_core::version::insert_version_header;
use ballista_examples::test_util;
use datafusion::arrow::array::RecordBatch;
use datafusion::arrow::error::ArrowError;
@@ -228,15 +229,18 @@ impl SubstraitSchedulerClient {
.max_encoding_message_size(self.max_message_size)
.max_decoding_message_size(self.max_message_size);
- let execute_query_params = ExecuteQueryParams {
+ let mut request = tonic::Request::new(ExecuteQueryParams {
session_id: self.session_id.clone(),
settings: vec![],
operation_id: uuid::Uuid::now_v7().to_string(),
query: Some(SubstraitPlan(plan)),
- };
+ });
+ // The scheduler rejects job submissions from clients that don't send
+ // their Ballista version.
+ insert_version_header(request.metadata_mut());
let response = scheduler
- .execute_query(execute_query_params)
+ .execute_query(request)
.await
.map_err(|e| {
DataFusionError::Execution(format!("Failed to execute query: {e:?}"))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org