[AURON #2422] make RSS spill memory triggers configurable

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,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<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);
+
+    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);
+
     public static final ConfigOption<Boolean> SMJ_FALLBACK_ENABLE = new ConfigOption<>(Boolean.class)
             .withKey("auron.smjfallback.enable")
             .withCategory("Operator Supports")