This is an automated email from the ASF dual-hosted git repository.
cxzl25 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git
The following commit(s) were added to refs/heads/master by this push:
new 377e53242 [AURON #2422] Make RSS spill memory triggers configurable
(#2423)
377e53242 is described below
commit 377e53242a0089971d3167716b9e2972c55532e2
Author: cxzl25 <[email protected]>
AuthorDate: Tue Aug 4 14:02:21 2026 +0800
[AURON #2422] Make RSS spill memory triggers configurable (#2423)
# Which issue does this PR close?
Closes #2422
# Rationale for this change
Configurable spill trigger (rss_sort_repartitioner.rs). Replace the
hardcoded `mem_used_percent() > 0.4` with two tunable thresholds; spill
fires when either is exceeded (whichever first):
- `spark.auron.rss.spill.memoryFraction` (relative, default 0.4 =
previous behavior)
- `spark.auron.rss.spill.memorySize` (absolute bytes, default 0 =
disabled; acts as a deterministic cap on the per-spill burst regardless
of executor memory)
Both are read once per executor via OnceCell (no per-batch JNI conf
reads).
Adds `MemConsumer::mem_used()` (absolute bytes accessor) to support the
absolute threshold.
# What changes are included in this PR?
Auron's rss shuffle buffers records in native memory and spills to
celeborn in bursts. At a shuffle-stage peak, many tasks spill
simultaneously; if each spill is large, the aggregate burst can overrun
a celeborn worker's inbound capacity and push the worker into a
push-pause: it stops reading, holds inbound buffers, and only recovers
after the idle-connection timeout / connection close (observed as a
multi-minute stall — tasks survive via celeborn's client-side revive but
throughput degrades).
The previous hardcoded 0.4 fraction made the per-spill burst size depend
on executor memory and not directly tunable against this.
# Are there any user-facing changes?
No
# How was this patch tested?
Production environment verification
# Was this patch authored or co-authored using generative AI tooling?
- [x] Yes
- [ ] No
---
.../apache/auron/configuration/ConfigOption.java | 37 ++++++++++++++++++++++
native-engine/auron-jni-bridge/src/conf.rs | 2 ++
native-engine/auron-memmgr/src/lib.rs | 5 +++
.../src/shuffle/rss_sort_repartitioner.rs | 15 ++++++++-
.../configuration/SparkAuronConfiguration.java | 34 ++++++++++++++++++--
5 files changed, 90 insertions(+), 3 deletions(-)
diff --git
a/auron-core/src/main/java/org/apache/auron/configuration/ConfigOption.java
b/auron-core/src/main/java/org/apache/auron/configuration/ConfigOption.java
index 04a4e0f22..e564d7abc 100644
--- a/auron-core/src/main/java/org/apache/auron/configuration/ConfigOption.java
+++ b/auron-core/src/main/java/org/apache/auron/configuration/ConfigOption.java
@@ -18,7 +18,9 @@ package org.apache.auron.configuration;
import java.util.ArrayList;
import java.util.List;
+import java.util.Optional;
import java.util.function.Function;
+import java.util.function.Predicate;
import org.apache.auron.jni.AuronAdaptor;
/**
@@ -53,6 +55,12 @@ public class ConfigOption<T> {
/** The function to compute the default value. */
private Function<AuronConfiguration, T> dynamicDefaultValueFunction;
+ /**
+ * Optional validator. Returns a non-empty error message (wrapped in
{@link Optional}) when the
+ * value is invalid, or an empty {@link Optional} when it is valid.
+ */
+ private Function<T, Optional<String>> validator;
+
/**
* Type of the value that this ConfigOption describes.
*
@@ -98,6 +106,35 @@ public class ConfigOption<T> {
return this;
}
+ /**
+ * Registers a validator that rejects values for which {@code check}
returns {@code false}. When
+ * the engine-side configuration backend reads a user-supplied value that
fails the check, it
+ * fails fast with an {@link IllegalArgumentException} carrying {@code
errorMsg}.
+ *
+ * @param check predicate that returns {@code true} for valid values.
+ * @param errorMsg human-readable message used when the check fails.
+ * @return this {@code ConfigOption} for chaining.
+ */
+ public ConfigOption<T> checkValue(Predicate<T> check, String errorMsg) {
+ this.validator = v -> check.test(v) ? Optional.empty() :
Optional.of(errorMsg);
+ return this;
+ }
+
+ /**
+ * Validates a resolved value against the registered validator, if any.
Called by the
+ * engine-side configuration backend when it reads a user-supplied value.
+ *
+ * @param value the value to validate.
+ * @return an empty {@link Optional} if valid or no validator is
registered, otherwise an
+ * {@link Optional} carrying the error message.
+ */
+ public Optional<String> validate(T value) {
+ if (validator == null) {
+ return Optional.empty();
+ }
+ return validator.apply(value);
+ }
+
/**
* Gets the configuration key.
*
diff --git a/native-engine/auron-jni-bridge/src/conf.rs
b/native-engine/auron-jni-bridge/src/conf.rs
index fa6c1d57a..9f160beb1 100644
--- a/native-engine/auron-jni-bridge/src/conf.rs
+++ b/native-engine/auron-jni-bridge/src/conf.rs
@@ -50,6 +50,8 @@ define_conf!(IntConf, TOKIO_WORKER_THREADS_PER_CPU);
define_conf!(IntConf, TASK_CPUS);
define_conf!(IntConf, SHUFFLE_COMPRESSION_TARGET_BUF_SIZE);
define_conf!(StringConf, SPILL_COMPRESSION_CODEC);
+define_conf!(DoubleConf, RSS_SPILL_MEMORY_FRACTION);
+define_conf!(LongConf, RSS_SPILL_MEMORY_SIZE);
define_conf!(BooleanConf, SMJ_FALLBACK_ENABLE);
define_conf!(IntConf, SMJ_FALLBACK_ROWS_THRESHOLD);
define_conf!(IntConf, SMJ_FALLBACK_MEM_SIZE_THRESHOLD);
diff --git a/native-engine/auron-memmgr/src/lib.rs
b/native-engine/auron-memmgr/src/lib.rs
index e9cb90311..32c89971d 100644
--- a/native-engine/auron-memmgr/src/lib.rs
+++ b/native-engine/auron-memmgr/src/lib.rs
@@ -224,6 +224,11 @@ pub trait MemConsumer: Send + Sync {
mem_used as f64 / consumer_mem_max as f64
}
+ /// Absolute bytes currently used by this consumer.
+ fn mem_used(&self) -> usize {
+ self.consumer_info().status.lock().mem_used
+ }
+
fn set_spillable(&self, spillable: bool) {
let consumer_info = self.consumer_info();
let mut consumer_status = consumer_info.status.lock();
diff --git
a/native-engine/datafusion-ext-plans/src/shuffle/rss_sort_repartitioner.rs
b/native-engine/datafusion-ext-plans/src/shuffle/rss_sort_repartitioner.rs
index b0b2c825c..61617fe57 100644
--- a/native-engine/datafusion-ext-plans/src/shuffle/rss_sort_repartitioner.rs
+++ b/native-engine/datafusion-ext-plans/src/shuffle/rss_sort_repartitioner.rs
@@ -17,11 +17,16 @@ use std::sync::Weak;
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
+use auron_jni_bridge::{
+ conf,
+ conf::{DoubleConf, LongConf},
+};
use auron_memmgr::{MemConsumer, MemConsumerInfo, MemManager};
use datafusion::{common::Result, physical_plan::metrics::Time};
use datafusion_ext_commons::arrow::array_size::BatchSize;
use futures::lock::Mutex;
use jni::objects::GlobalRef;
+use once_cell::sync::OnceCell;
use crate::shuffle::{Partitioning, ShuffleRepartitioner,
buffered_data::BufferedData};
@@ -102,7 +107,15 @@ impl ShuffleRepartitioner for RssSortShuffleRepartitioner {
// we are likely to spill more frequently because the cost of spilling
a shuffle
// repartition is lower than other consumers.
// rss shuffle spill has even lower cost than normal shuffle
- if self.mem_used_percent() > 0.4 {
+ static PCT_THRESHOLD: OnceCell<f64> = OnceCell::new();
+ static SIZE_THRESHOLD: OnceCell<i64> = OnceCell::new();
+ let pct_threshold =
+ *PCT_THRESHOLD.get_or_init(||
conf::RSS_SPILL_MEMORY_FRACTION.value().unwrap_or(0.4));
+ let size_threshold =
+ *SIZE_THRESHOLD.get_or_init(||
conf::RSS_SPILL_MEMORY_SIZE.value().unwrap_or(0));
+ if self.mem_used_percent() > pct_threshold
+ || (size_threshold > 0 && self.mem_used() > size_threshold as
usize)
+ {
self.force_spill().await?;
}
Ok(())
diff --git
a/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java
b/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java
index 366f32c13..97da32d4f 100644
---
a/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java
+++
b/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java
@@ -227,6 +227,28 @@ public class SparkAuronConfiguration extends
AuronConfiguration {
+ "Common options include lz4, snappy, and gzip.
The choice affects both spill performance and disk space usage.")
.withDefaultValue("lz4");
+ public static final ConfigOption<Double> RSS_SPILL_MEMORY_FRACTION = new
ConfigOption<>(Double.class)
+ .withKey("auron.rss.spill.memoryFraction")
+ .withCategory("Runtime Configuration")
+ .withDescription(
+ "RSS shuffle spill trigger: spill when this consumer's
memory fraction exceeds the threshold. "
+ + "A lower value makes auron spill sooner with a
smaller buffer each time, so each spill pushes "
+ + "a smaller burst to the celeborn worker and
reduces the peak in-flight bytes that can trigger "
+ + "a worker push pause. Default 0.4 preserves the
previous hardcoded behavior.")
+ .withDefaultValue(0.4)
+ .checkValue(v -> v != null && v >= 0.0 && v <= 1.0, "must be a
number in [0.0, 1.0]");
+
+ public static final ConfigOption<Long> RSS_SPILL_MEMORY_SIZE = new
ConfigOption<>(Long.class)
+ .withKey("auron.rss.spill.memorySize")
+ .withCategory("Runtime Configuration")
+ .withDescription(
+ "RSS shuffle spill trigger (absolute): force spill when
this consumer's buffered bytes exceed "
+ + "this size, regardless of the relative fraction.
This caps the per-spill burst size pushed to "
+ + "the celeborn worker deterministically (the
fraction-based threshold varies with executor memory). "
+ + "0 means disabled (only the fraction threshold
is used). Default 0 preserves previous behavior.")
+ .withDefaultValue(0L)
+ .checkValue(v -> v != null && v >= 0L, "must be >= 0");
+
public static final ConfigOption<Boolean> SMJ_FALLBACK_ENABLE = new
ConfigOption<>(Boolean.class)
.withKey("auron.smjfallback.enable")
.withCategory("Operator Supports")
@@ -544,12 +566,20 @@ public class SparkAuronConfiguration extends
AuronConfiguration {
: option instanceof SparkContextOption
? GetFromSparkType.FROM_SPARK_CONTEXT
: GetFromSparkType.FROM_SPARK_ENV;
- return Optional.ofNullable(getFromSpark(
+ T value = getFromSpark(
option.key(),
option.altKeys(),
option.getValueClass(),
() -> getOptionDefaultValue(option),
- getFromSparkType));
+ getFromSparkType);
+ if (value != null) {
+ Optional<String> error = option.validate(value);
+ if (error.isPresent()) {
+ throw new IllegalArgumentException(
+ "Invalid value for config '" + option.key() + "': " +
error.get() + " (value=" + value + ")");
+ }
+ }
+ return Optional.ofNullable(value);
}
enum GetFromSparkType {
