[AURON #2422] Make RSS spill memory triggers configurable

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<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 ConfigOption<T> withDynamicDefaultValue(Function<AuronConfiguration, T> 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<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.
      *

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);

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();

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(())

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 <T> Optional<T> getOptional(ConfigOption<T> 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<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 {