feat: add BALLISTA_PROTOCOL_VERSION + k8s health probes - #2088
Conversation
Introduces a compile-time `BALLISTA_PROTOCOL_VERSION` that gates executor↔scheduler communication and adds `/healthz` and `/readyz` endpoints so live-upgrade rollouts have something for Kubernetes to observe. Scheduler rejects executors whose protocol version does not match with `FailedPrecondition`, an `info!` log, and a `ballista_scheduler_rejected_by_protocol_version_total` counter. Missing registration metadata on heartbeats is treated as a mismatch. Executor flips /readyz off on RPC failure and self-terminates after 5 consecutive heartbeat failures, so a mismatched pod crash-loops until its image is bumped rather than silently sitting idle. Health probes are always mounted regardless of the `rest-api` feature: scheduler on the existing axum router, executor on a new `--bind-health-port` (default 50053). Livez is deliberately dumb — scheduler flapping must not restart executors in a cascade. Fixes apache#2071. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
the
documentation
|
@phillipleblanc FYI |
phillipleblanc
left a comment
phillipleblanc
left a comment
There was a problem hiding this comment.
I think the protocol version makes sense. I have one comment about the health service specifically below, but in general I think this points to a tension in what Ballista is and how we expect consumers to use it.
We use Ballista as a library that we embed into our existing process (i.e. we call scheduler_process::create_scheduler() and serve this on our existing gRPC port rather than let Ballista create it).
I think it would be useful to have a split between the Ballista standalone processes (i.e. ballista-scheduler, ballista-executor) vs the libraries that those binaries (and other consumers) can use to build their own distributed query engines. So for the health probes specifically, I think it would make sense to have it live in the binary/process startup, not as part of the library component. That way its still usable out of the box and is a good reference implementation for anyone that is integrating with ballista, but doesn't force them to take this specific implementation.
| /// Port for the executor's gRPC service. | ||
| pub grpc_port: u16, | ||
| /// Port for the executor's HTTP health server (`/healthz` and `/readyz`). | ||
| pub health_port: u16, |
There was a problem hiding this comment.
In our project we embed Ballista in our process that already defines health/readiness probes that are specific to our implementation. This would add a second health surface that we can't turn off without removing this code in our fork.
Perhaps similar to the rest-api feature flag, we could add a health-probes flag that controls this? (or bundle this with the rest-api feature flag?) Alternatively a separate option to control whether ballista starts the health service at all that we could disable.
There was a problem hiding this comment.
@phillipleblanc , I wasn't aware you (or anyone) was using it that way. Given your use case (should we add this to a persona @andygrove ?) I think you are absolutely correct about architecture, and I was negligent in the design. PTAL at the current state. I have:
- removed axum from the executor library
- exposed a boolean is_healthy() method in both scheduler and executor
- moved axum & the health check server into the executor binary
- not changed the scheduler (this is the one I'm not sure about)
My thought is that if schedulers are already starting a REST API, it is natural to put the health check endpoint there so as to not require another port binding. However, I can see how this might be a separate concern from the REST API.
Happy to entertain either option...
There was a problem hiding this comment.
@avantgardnerio no problem! Originally Ballista was started as a standalone binary process, and it was relatively recently that @milenkovicm started to focus it to be more of a library (which is one of the reasons we decided to go with Ballista rather than something like datafusion-distributed which is also a library, but has different design goals)
I agree that putting the health APIs on the scheduler rest API make sense. Everything else looks good now, thanks for reworking this!
The library was spinning up its own axum HTTP server on `--bind-health-port` from inside `start_executor_process`, adding a hard axum dep and a second health surface that library embedders had no way to disable short of forking the code. Health probes are a deployment concern, so push them to where the deployment lives: `ballista/executor/src/bin/main.rs`. The library keeps `ExecutorHealth` — a cheap `Arc<AtomicBool>` that the heartbeat loops flip on every RPC outcome — and exposes it on `ExecutorProcessConfig` so callers can observe it however they like. The axum server and its handlers are gated behind `#[cfg(feature = "build-binary")]`, and the `axum` dep moves from unconditional to `optional = true` under `build-binary`. With `--no-default-features`, `ballista-executor` no longer has a direct edge to axum in `cargo tree`. Verified end-to-end against the local native cluster: scheduler `/readyz` reports "ready: 2 executors registered" and both executors' `/healthz` and `/readyz` return 200, so the handle is threaded through bin -> config -> heartbeat loop -> probe correctly. Addresses apache#2088 (comment) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- examples/mtls-cluster.rs: add ballista_protocol_version to ExecutorRegistration and pass ExecutorHealth to poll_loop so the example compiles under `--all-targets` clippy. - docs/.../kubernetes.md: prettier reformatting — _emphasis_ instead of *emphasis*, and pad the rolling-upgrade table columns to match the longest row. - python/Cargo.lock: `cd python && cargo update` to resync the lockfile with the workspace manifests, so `cargo metadata --locked` inside `python/` succeeds again. The TPC-DS SF1 CI failure was infrastructure (`No space left on device` on the runner), not a code issue — re-running the job should clear it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Phillip's fork wraps the Ballista scheduler in a larger process whose health check is a superset of ours — his handler already reports app-level state and would like to AND the scheduler's readiness in as one factor. Add a public `SchedulerServer::is_ready() -> bool` that returns true when at least `min_ready_executors` executors have live heartbeats. The built-in `/readyz` handler now delegates to it (single source of truth), and embedders can call the same method from whatever health surface they expose. This is orthogonal to the `/healthz` and `/readyz` routes mounted on the axum server — those stay unconditional, so operators who don't embed still get the k8s probes for free. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirror what SchedulerServer::is_ready() offers on the executor side. The heartbeat loops already flip the internal atomic on every RPC outcome, and the built-in `/readyz` axum handler reads it via this method under `build-binary`. Making the method `pub` and dropping the cfg gate lets an embedder observe the same signal directly and fold it into their own health surface, without linking axum. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
phillipleblanc
left a comment
phillipleblanc
left a comment
There was a problem hiding this comment.
Thanks! This looks good to me now.
|
https://datafusion.apache.org/ballista/contributors-guide/user-personas.html Persona 3 should cover @phillipleblanc case if im not mistaken |
3 participants
Summary
Closes #2071.
My notes:
With that out of the way, onto the Claude ramblings:
Adds a compile-time
BALLISTA_PROTOCOL_VERSION: u32 = 1inballista_corethat gates executor↔scheduler traffic, plusKubernetes-style
/healthzand/readyzprobes on both binaries solive upgrades have something to observe.
Protocol version
uint32 ballista_protocol_version = 7onExecutorRegistration.poll_work.Status::failed_precondition, aninfo!log, and incrementsballista_scheduler_rejected_by_protocol_version_total(prometheusfeature).
/readyzflips off on RPC failure, and after 5 consecutiveheartbeat failures the executor self-terminates so a mismatched pod
crash-loops until its image is bumped rather than sitting idle.
Health probes
/healthzis always 200./readyzreturns 200 iff at least--min-ready-executorsexecutors are registered (default 1; set to 0for dev/single-node).
--bind-health-port(default50053),since Arrow Flight isn't HTTP.
/healthzis always 200 — deliberatelyindependent of scheduler connectivity to prevent restart cascades.
/readyzreflects the last heartbeat.Rolling upgrade semantics (documented in
kubernetes.md)Bump both Deployments' image tags together in the same release; each
rolls independently under its own
maxSurge, and the scheduler's/readyzgate keepsServicetraffic on the old scheduler until thenew one has matching-version executors registered.
Bump policy: bump
BALLISTA_PROTOCOL_VERSIONonly on releases thatchange the executor↔scheduler wire format. Most releases won't.
Test plan
cargo test -p ballista-scheduler --lib(236 pass, incl. 8 new)cargo test -p ballista-executor --lib(45 pass)cargo clippy --all-targets -p ballista-core -p ballista-executor -p ballista-scheduler(clean)cargo fmt --all/healthzand/readyzendpoints return 200; scheduler/readyzshows
ready: 2 executors registered.New unit tests cover:
poll_workrejected on mismatched protocol version/healthzalways OK/readyzfails with no executors/readyzOK whenmin_ready_executors = 0/readyzOK after an active heartbeat is recordedNotes for reviewers
executor" and is rejected once we start at
BALLISTA_PROTOCOL_VERSION = 1. This is intentional per the issue's"no dual-implementation window" design.
exists — so during a rolling upgrade, old-scheduler ↔ new-executor
keeps working. Only new-scheduler ↔ old-executor is rejected.
rest-apifeature or--disable-rest-api, since they'reoperational, not part of the REST API surface.
🤖 Generated with Claude Code