From 9a2443689c4c6ce8883c82eadf0d1f0bfc155d6e Mon Sep 17 00:00:00 2001 From: sychen Date: Fri, 24 Jul 2026 15:40:39 +0800 Subject: [PATCH] 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")