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/15 15:32:15 UTC
(datafusion-ballista) branch main updated: feat(executor): report executor memory usage in heartbeats (#2581)
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 8db228cb6 feat(executor): report executor memory usage in heartbeats (#2581)
8db228cb6 is described below
commit 8db228cb621592d830aac861c00a5ab68d2143a9
Author: Richard Durkee <56...@users.noreply.github.com>
AuthorDate: Thu Oct 15 11:32:09 2026 -0400
feat(executor): report executor memory usage in heartbeats (#2581)
---
ballista/core/proto/ballista.proto | 9 ++
ballista/core/src/serde/generated/ballista.rs | 14 ++-
ballista/executor/README.md | 9 ++
ballista/executor/src/bin/main.rs | 6 +-
ballista/executor/src/executor.rs | 24 ++++-
ballista/executor/src/executor_process.rs | 5 +-
ballista/executor/src/executor_server.rs | 124 +++++++++++++++++++++++++-
ballista/executor/src/metrics/mod.rs | 22 +++++
ballista/scheduler/src/api/handlers.rs | 32 +++++++
9 files changed, 238 insertions(+), 7 deletions(-)
diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto
index acee01bd6..9a6fbab9a 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -748,6 +748,15 @@ message ExecutorMetric {
uint64 proc_virtual_memory = 5;
uint64 peak_physical_memory = 6;
uint64 peak_virtual_memory = 7;
+ // Bytes currently reserved across the memory pools of the executor's
+ // running tasks.
+ uint64 memory_pool_reserved = 8;
+ // The executor's total memory pool budget, in bytes. Absent when the
+ // pool is unbounded.
+ uint64 memory_pool_size = 9;
+ // Outstanding bytes allocated by the executor process, as counted by its
+ // global allocator. Includes memory the pools do not track.
+ uint64 allocated_memory = 10;
}
}
diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs
index 2e14088d7..c16086d35 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -997,7 +997,7 @@ pub struct ExecutorHeartbeat {
#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)]
pub struct ExecutorMetric {
/// TODO add more metrics
- #[prost(oneof = "executor_metric::Metric", tags = "1, 2, 3, 4, 5, 6, 7")]
+ #[prost(oneof = "executor_metric::Metric", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10")]
pub metric: ::core::option::Option<executor_metric::Metric>,
}
/// Nested message and enum types in `ExecutorMetric`.
@@ -1019,6 +1019,18 @@ pub mod executor_metric {
PeakPhysicalMemory(u64),
#[prost(uint64, tag = "7")]
PeakVirtualMemory(u64),
+ /// Bytes currently reserved across the memory pools of the executor's
+ /// running tasks.
+ #[prost(uint64, tag = "8")]
+ MemoryPoolReserved(u64),
+ /// The executor's total memory pool budget, in bytes. Absent when the
+ /// pool is unbounded.
+ #[prost(uint64, tag = "9")]
+ MemoryPoolSize(u64),
+ /// Outstanding bytes allocated by the executor process, as counted by its
+ /// global allocator. Includes memory the pools do not track.
+ #[prost(uint64, tag = "10")]
+ AllocatedMemory(u64),
}
}
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
diff --git a/ballista/executor/README.md b/ballista/executor/README.md
index 0f2dbab3e..9935babca 100644
--- a/ballista/executor/README.md
+++ b/ballista/executor/README.md
@@ -66,6 +66,15 @@ they explicitly wrap it with `ballista_executor::alloc_accounting::AccountingAll
logging is only started by the standalone binary. Embedders can opt into pool metrics through
`start_executor_process_with_memory_metrics` without installing the accounting allocator.
+Executors started this way also send these figures to the scheduler in push-mode heartbeats, as
+the `memory_pool_reserved`, `memory_pool_size` and `allocated_memory` metrics in
+`GET /api/executors` and `GET /api/executor/{executor_id}`. They are process metrics, so they're
+sent when the metrics collection policy is `proc` (the default) or `all`. `memory_pool_size` is
+left out when the pool is unbounded. Allocated bytes are only sent when a source is set with
+`ExecutorMemoryMetrics::set_allocated_bytes_source`, as the standalone binary does with
+`alloc_accounting::current_balance`. Pull-mode executors report no heartbeat metrics
+(`poll_work` carries none), so none of these appear for them.
+
## Using it as a library
`ExecutorProcessConfig` carries the same override hooks as the scheduler, so an embedder can supply
diff --git a/ballista/executor/src/bin/main.rs b/ballista/executor/src/bin/main.rs
index dcba54b2e..d20fb358a 100644
--- a/ballista/executor/src/bin/main.rs
+++ b/ballista/executor/src/bin/main.rs
@@ -92,6 +92,9 @@ async fn main() -> ballista_core::error::Result<()> {
}
let memory_metrics = ExecutorMemoryMetrics::default();
+ // This binary installs `AccountingAllocator` as its global allocator, so
+ // its balance can be reported alongside the pool usage.
+ memory_metrics.set_allocated_bytes_source(current_balance);
let logging_metrics = memory_metrics.clone();
let memory_logging = tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(10));
@@ -106,7 +109,8 @@ async fn main() -> ballista_core::error::Result<()> {
log::info!(
"Ballista executor memory usage: allocated {:.1} MiB, \
reserved {:.1} MiB, pool size {pool_size} ({} live pools)",
- current_balance() as f64 / (1024.0 * 1024.0),
+ logging_metrics.allocated_bytes().unwrap_or_default() as f64
+ / (1024.0 * 1024.0),
usage.reserved as f64 / (1024.0 * 1024.0),
usage.pool_count,
);
diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs
index b6121c2b1..eabf1b338 100644
--- a/ballista/executor/src/executor.rs
+++ b/ballista/executor/src/executor.rs
@@ -21,8 +21,8 @@ use crate::execution_engine::DefaultExecutionEngine;
use crate::execution_engine::ExecutionEngine;
use crate::execution_engine::QueryStageExecutor;
use crate::execution_loop::any_to_string;
-use crate::metrics::ExecutorMetricsCollector;
use crate::metrics::LoggingMetricsCollector;
+use crate::metrics::{ExecutorMemoryMetrics, ExecutorMetricsCollector};
use crate::runtime_cache::SessionRuntimeCache;
use ballista_core::ConfigProducer;
use ballista_core::JobId;
@@ -106,6 +106,10 @@ pub struct Executor {
/// `produce_runtime_for_session` reuses read-side state across a session's
/// tasks; when `None`, each task builds a runtime from `runtime_producer`.
session_runtime_cache: Option<Arc<dyn SessionRuntimeCache>>,
+
+ /// Executor-wide memory usage, sampled for heartbeats. Set when the
+ /// executor was started with memory metrics enabled.
+ memory_metrics: Option<ExecutorMemoryMetrics>,
}
impl Executor {
@@ -155,6 +159,7 @@ impl Executor {
tasks_drained_waker: Default::default(),
execution_engine,
session_runtime_cache: None,
+ memory_metrics: None,
}
}
/// Creates new Executor with default `ExecutionEngine`.
@@ -180,6 +185,7 @@ impl Executor {
tasks_drained_waker: Default::default(),
execution_engine: Arc::new(DefaultExecutionEngine::new()),
session_runtime_cache: None,
+ memory_metrics: None,
}
}
}
@@ -209,6 +215,22 @@ impl Executor {
self
}
+ /// Attaches (or clears) the executor-wide memory metrics reported in
+ /// heartbeats. It should be the same handle the memory pool policy
+ /// registers task pools with.
+ pub(crate) fn with_memory_metrics(
+ mut self,
+ metrics: Option<ExecutorMemoryMetrics>,
+ ) -> Self {
+ self.memory_metrics = metrics;
+ self
+ }
+
+ /// The executor-wide memory metrics, when memory metrics are enabled.
+ pub(crate) fn memory_metrics(&self) -> Option<&ExecutorMemoryMetrics> {
+ self.memory_metrics.as_ref()
+ }
+
/// Produces the runtime for a task, reusing the session's shared read-side
/// state when a session cache is attached; otherwise builds a runtime per
/// task via the `runtime_producer`. `vcores_consumed` is the number of
diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs
index ab24db881..1f72c925f 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -567,7 +567,7 @@ async fn start_executor_process_inner(
}
};
- let pool_policy = if let Some(metrics) = memory_metrics {
+ let pool_policy = if let Some(metrics) = memory_metrics.clone() {
observe_memory_pools(pool_policy, metrics)
} else {
pool_policy
@@ -632,7 +632,8 @@ async fn start_executor_process_inner(
}
}),
)
- .with_session_runtime_cache(Some(session_runtime_cache)),
+ .with_session_runtime_cache(Some(session_runtime_cache))
+ .with_memory_metrics(memory_metrics),
);
let connect_timeout = opt.scheduler_connect_timeout_seconds as u64;
diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs
index 53274c4e2..e73229221 100644
--- a/ballista/executor/src/executor_server.rs
+++ b/ballista/executor/src/executor_server.rs
@@ -67,7 +67,7 @@ use crate::cpu_bound_executor::DedicatedExecutor;
use crate::executor::Executor;
use crate::executor_process::{ExecutorProcessConfig, remove_job_data};
use crate::health::ExecutorHealth;
-use crate::metrics::ExecutorMetricCollectionPolicy;
+use crate::metrics::{ExecutorMemoryMetrics, ExecutorMetricCollectionPolicy};
use crate::shutdown::ShutdownNotifier;
use crate::{TaskCompletionExtras, TaskExecutionTimes, as_task_status};
@@ -545,6 +545,16 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
/// Getting executor's metrics
fn get_executor_metrics(&self) -> Vec<ExecutorMetric> {
+ let mut metrics = self.get_host_and_process_metrics();
+ metrics.extend(heartbeat_memory_metrics(
+ self.metric_collection_policy,
+ self.executor.memory_metrics(),
+ ));
+ metrics
+ }
+
+ /// Host and process memory, as selected by the metric collection policy.
+ fn get_host_and_process_metrics(&self) -> Vec<ExecutorMetric> {
match self.metric_collection_policy {
ExecutorMetricCollectionPolicy::SystemOnly => {
let mut executor_system = System::new_all();
@@ -1026,5 +1036,115 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorGrpc
}
}
+/// The executor's memory metrics to send in a heartbeat. These are process
+/// metrics, so they're only sent when the policy collects process metrics and
+/// the executor was started with memory metrics.
+fn heartbeat_memory_metrics(
+ policy: ExecutorMetricCollectionPolicy,
+ metrics: Option<&ExecutorMemoryMetrics>,
+) -> Vec<ExecutorMetric> {
+ match (policy, metrics) {
+ (
+ ExecutorMetricCollectionPolicy::ProcessOnly
+ | ExecutorMetricCollectionPolicy::SystemAndProcess,
+ Some(metrics),
+ ) => executor_memory_metrics(metrics),
+ _ => vec![],
+ }
+}
+
+/// Converts a snapshot of executor-wide memory usage into heartbeat metrics.
+///
+/// Reports nothing until the executor has resolved its memory pool. The pool
+/// capacity is left out when the pool is unbounded, and allocated bytes when
+/// the process's allocator doesn't count them.
+fn executor_memory_metrics(metrics: &ExecutorMemoryMetrics) -> Vec<ExecutorMetric> {
+ let Some(usage) = metrics.snapshot() else {
+ return vec![];
+ };
+ let mut values = vec![executor_metric::Metric::MemoryPoolReserved(
+ usage.reserved as u64,
+ )];
+ if let Some(capacity) = usage.pool_size {
+ values.push(executor_metric::Metric::MemoryPoolSize(capacity));
+ }
+ if let Some(allocated) = metrics.allocated_bytes() {
+ values.push(executor_metric::Metric::AllocatedMemory(allocated as u64));
+ }
+ values
+ .into_iter()
+ .map(|metric| ExecutorMetric {
+ metric: Some(metric),
+ })
+ .collect()
+}
+
#[cfg(test)]
-mod test {}
+mod test {
+ use super::*;
+ use datafusion::execution::memory_pool::{FairSpillPool, MemoryConsumer, MemoryPool};
+
+ fn values(metrics: Vec<ExecutorMetric>) -> Vec<executor_metric::Metric> {
+ metrics.into_iter().filter_map(|m| m.metric).collect()
+ }
+
+ #[test]
+ fn reports_nothing_before_the_pool_is_initialized() {
+ let metrics = ExecutorMemoryMetrics::default();
+ assert!(executor_memory_metrics(&metrics).is_empty());
+ }
+
+ #[test]
+ fn reports_reserved_pool_size_and_allocated_bytes() {
+ let metrics = ExecutorMemoryMetrics::default();
+ metrics.set_pool_size(Some(4096));
+ metrics.set_allocated_bytes_source(|| 12_345);
+ let pool: Arc<dyn MemoryPool> = Arc::new(FairSpillPool::new(4096));
+ metrics.register(&pool);
+ let reservation = MemoryConsumer::new("r").register(&pool);
+ reservation.try_grow(1000).unwrap();
+
+ assert_eq!(
+ values(executor_memory_metrics(&metrics)),
+ vec![
+ executor_metric::Metric::MemoryPoolReserved(1000),
+ executor_metric::Metric::MemoryPoolSize(4096),
+ executor_metric::Metric::AllocatedMemory(12_345),
+ ]
+ );
+ }
+
+ #[test]
+ fn sends_memory_metrics_only_when_collecting_process_metrics() {
+ let metrics = ExecutorMemoryMetrics::default();
+ metrics.set_pool_size(Some(4096));
+
+ for policy in [
+ ExecutorMetricCollectionPolicy::ProcessOnly,
+ ExecutorMetricCollectionPolicy::SystemAndProcess,
+ ] {
+ assert_eq!(heartbeat_memory_metrics(policy, Some(&metrics)).len(), 2);
+ }
+ for policy in [
+ ExecutorMetricCollectionPolicy::SystemOnly,
+ ExecutorMetricCollectionPolicy::Off,
+ ] {
+ assert!(heartbeat_memory_metrics(policy, Some(&metrics)).is_empty());
+ }
+ assert!(
+ heartbeat_memory_metrics(ExecutorMetricCollectionPolicy::ProcessOnly, None)
+ .is_empty()
+ );
+ }
+
+ #[test]
+ fn omits_pool_size_when_unbounded_and_allocated_without_a_source() {
+ let metrics = ExecutorMemoryMetrics::default();
+ metrics.set_pool_size(None);
+
+ assert_eq!(
+ values(executor_memory_metrics(&metrics)),
+ vec![executor_metric::Metric::MemoryPoolReserved(0)]
+ );
+ }
+}
diff --git a/ballista/executor/src/metrics/mod.rs b/ballista/executor/src/metrics/mod.rs
index 15de723e4..4b7209eea 100644
--- a/ballista/executor/src/metrics/mod.rs
+++ b/ballista/executor/src/metrics/mod.rs
@@ -34,11 +34,15 @@ pub struct ExecutorMemoryMetrics {
inner: Arc<Mutex<MemoryMetricsState>>,
}
+/// Reads a process's outstanding allocated bytes.
+type AllocatedBytesSource = Arc<dyn Fn() -> usize + Send + Sync>;
+
#[derive(Default)]
struct MemoryMetricsState {
initialized: bool,
pool_size: Option<u64>,
pools: Vec<Weak<dyn MemoryPool>>,
+ allocated_bytes: Option<AllocatedBytesSource>,
}
/// A point-in-time view of executor memory pool usage.
@@ -69,6 +73,24 @@ impl ExecutorMemoryMetrics {
}
}
+ /// Sets where [`Self::allocated_bytes`] reads the process's outstanding
+ /// allocated bytes from, such as
+ /// [`crate::alloc_accounting::current_balance`]. Only set it when the
+ /// process's global allocator actually does that accounting.
+ pub fn set_allocated_bytes_source(
+ &self,
+ source: impl Fn() -> usize + Send + Sync + 'static,
+ ) {
+ self.inner.lock().allocated_bytes = Some(Arc::new(source));
+ }
+
+ /// Outstanding allocated bytes in the process, or `None` when no source
+ /// was set with [`Self::set_allocated_bytes_source`].
+ pub fn allocated_bytes(&self) -> Option<usize> {
+ let source = self.inner.lock().allocated_bytes.clone();
+ source.map(|source| source())
+ }
+
/// Samples the resolved budget and current pool reservations.
///
/// Returns `None` before executor memory pool initialization. Reservation reads are
diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs
index 3e48f1d18..efdb3cff3 100644
--- a/ballista/scheduler/src/api/handlers.rs
+++ b/ballista/scheduler/src/api/handlers.rs
@@ -87,6 +87,9 @@ pub enum ExecutorMetricResponse {
ProcVirtualMemory(u64),
PeakPhysicalMemory(u64),
PeakVirtualMemory(u64),
+ MemoryPoolReserved(u64),
+ MemoryPoolSize(u64),
+ AllocatedMemory(u64),
}
impl ExecutorMetricResponse {
@@ -99,6 +102,9 @@ impl ExecutorMetricResponse {
Metric::ProcVirtualMemory(v) => Self::ProcVirtualMemory(v),
Metric::PeakPhysicalMemory(v) => Self::PeakPhysicalMemory(v),
Metric::PeakVirtualMemory(v) => Self::PeakVirtualMemory(v),
+ Metric::MemoryPoolReserved(v) => Self::MemoryPoolReserved(v),
+ Metric::MemoryPoolSize(v) => Self::MemoryPoolSize(v),
+ Metric::AllocatedMemory(v) => Self::AllocatedMemory(v),
})
}
}
@@ -687,6 +693,32 @@ pub async fn get_job_config<
mod tests {
use super::*;
+ /// The executor memory metrics' JSON names are part of the REST contract.
+ #[test]
+ fn executor_memory_metrics_serialize_with_stable_names() {
+ use ballista_core::serde::protobuf::ExecutorMetric;
+
+ let json = |metric| {
+ let response = ExecutorMetricResponse::from_proto(ExecutorMetric {
+ metric: Some(metric),
+ })
+ .unwrap();
+ serde_json::to_value(response).unwrap()
+ };
+ assert_eq!(
+ json(Metric::MemoryPoolReserved(1)),
+ serde_json::json!({"type": "memory_pool_reserved", "value": 1})
+ );
+ assert_eq!(
+ json(Metric::MemoryPoolSize(2)),
+ serde_json::json!({"type": "memory_pool_size", "value": 2})
+ );
+ assert_eq!(
+ json(Metric::AllocatedMemory(3)),
+ serde_json::json!({"type": "allocated_memory", "value": 3})
+ );
+ }
+
#[tokio::test]
// Still covers the deprecated embedded-proxy flag until it is removed.
#[allow(deprecated)]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org