From 9a2443689c4c6ce8883c82eadf0d1f0bfc155d6e Mon Sep 17 00:00:00 2001 From: sychen Date: Fri, 24 Jul 2026 15:40:39 +0800 Subject: [PATCH 1/2] feat: make RSS spill memory triggers configurable --- native-engine/auron-jni-bridge/src/conf.rs | 2 ++ native-engine/auron-memmgr/src/lib.rs | 5 +++++ .../src/shuffle/rss_sort_repartitioner.rs | 15 +++++++++++++- .../SparkAuronConfiguration.java | 20 +++++++++++++++++++ 4 files changed, 41 insertions(+), 1 deletion(-) 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 = OnceCell::new(); + static SIZE_THRESHOLD: OnceCell = 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..eb2b38e2c 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,26 @@ 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 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); + + public static final ConfigOption 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); + public static final ConfigOption SMJ_FALLBACK_ENABLE = new ConfigOption<>(Boolean.class) .withKey("auron.smjfallback.enable") .withCategory("Operator Supports") From 6312fef8d7832de90d6feae60e357b57c3c3029b Mon Sep 17 00:00:00 2001 From: sychen Date: Tue, 28 Jul 2026 20:43:03 +0800 Subject: [PATCH 2/2] feat(config): add value validation to ConfigOption --- .../auron/configuration/ConfigOption.java | 37 +++++++++++++++++++ .../SparkAuronConfiguration.java | 18 +++++++-- 2 files changed, 51 insertions(+), 4 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 @@ 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 { /** The function to compute the default value. */ private Function 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> validator; + /** * Type of the value that this ConfigOption describes. * @@ -98,6 +106,35 @@ public ConfigOption withDynamicDefaultValue(Function d 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 checkValue(Predicate 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 validate(T value) { + if (validator == null) { + return Optional.empty(); + } + return validator.apply(value); + } + /** * Gets the configuration key. * 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 eb2b38e2c..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 @@ -235,7 +235,8 @@ public class SparkAuronConfiguration extends AuronConfiguration { + "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); + .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 RSS_SPILL_MEMORY_SIZE = new ConfigOption<>(Long.class) .withKey("auron.rss.spill.memorySize") @@ -245,7 +246,8 @@ public class SparkAuronConfiguration extends AuronConfiguration { + "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); + .withDefaultValue(0L) + .checkValue(v -> v != null && v >= 0L, "must be >= 0"); public static final ConfigOption SMJ_FALLBACK_ENABLE = new ConfigOption<>(Boolean.class) .withKey("auron.smjfallback.enable") @@ -564,12 +566,20 @@ public Optional getOptional(ConfigOption option) { : 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 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 {