Skip to content

Commit dfde5de

Browse files
venkata91SteNicholas
authored andcommitted
[CELEBORN-2451] Make DRA safe under AUTO fallback by tying supportsReliableStorage to fallback policy
### What changes were proposed in this pull request? - Tie `supportsReliableStorage()` to the fallback policy: `true` only for `NEVER`, `false` for `AUTO`/`ALWAYS` (was keyed on deprecated `forceFallback`, so it returned `true` under `AUTO`). - Make `SparkShuffleManager` DRA warnings policy-aware; log the `registerShuffle` fallback at `error` only when shuffle tracking is also off. ### Why are the changes needed? - Fallback is per-shuffle, but `supportsReliableStorage()` is one global flag. Under `AUTO` a shuffle can fall back to local-disk `SortShuffleManager`; that output dies with a DRA-reclaimed executor, causing `FetchFailedException`. Nothing prevents the unsafe config today. - Reporting `false` under `AUTO`/`ALWAYS` lets Spark's `ExecutorAllocationManager` fail fast unless a shuffle service, shuffle tracking, or decommissioning is on, so `AUTO` is safe with `shuffleTracking.enabled=true`; `NEVER` keeps the fast path. - Trade-off: under `AUTO`, `reliableStorage=false` makes Spark recompute a lost executor's map stages even though Celeborn still holds the data; `NEVER` avoids this. Removing it needs a Spark-side change (per-shuffle reliable-storage granularity), tracked as a follow-up. ### Does this PR resolve a correctness bug? - [x] Yes ### Does this PR introduce _any_ user-facing change? - [x] Yes: with `AUTO` (default), DRA without shuffle service and without shuffle tracking now fails fast at startup. Enable tracking or set `NEVER`. ### How was this patch tested? - New `testReliableStorageFollowsFallbackPolicy`; `-Dtest=SparkShuffleManagerSuite` passes (5/5), spotless clean. Closes #3834 from venkata91/dra-auto-fallback-reliable-storage. Authored-by: Venkata Krishnan Sowrirajan <venkat.sowrirajan@gmail.com> Signed-off-by: Nicholas Jiang <programgeek@163.com>
1 parent e21b726 commit dfde5de

5 files changed

Lines changed: 148 additions & 21 deletions

File tree

‎README.md‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -320,14 +320,27 @@ spark.sql.adaptive.skewJoin.enabled true
320320
# Support Spark Dynamic Resource Allocation
321321
# Required Spark version >= 3.5.0
322322
spark.shuffle.sort.io.plugin.class org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO
323-
# Required Spark version >= 3.4.0, highly recommended to disable
324-
spark.dynamicAllocation.shuffleTracking.enabled false
323+
# Available since Spark 3.0; enabled by default since Spark 3.4.
324+
# With AUTO (default) or ALWAYS fallback, keep tracking enabled to protect local-disk shuffle data.
325+
spark.dynamicAllocation.shuffleTracking.enabled true
325326

326327
# Support ShuffleManager when defined in user jars
327328
# Required Spark version < 4.0.0 or without SPARK-45762, highly recommended to false for ShuffleManager in user-defined jar specified by --jars or spark.jars
328329
spark.executor.userClassPathFirst false
329330
```
330331

332+
With DRA enabled, `AUTO` and `ALWAYS` fallback policies require shuffle tracking unless an external
333+
shuffle service or Spark shuffle decommissioning is configured. Shuffle decommissioning requires
334+
`spark.decommission.enabled=true` and `spark.storage.decommission.shuffleBlocks.enabled=true`;
335+
also enable `spark.storage.decommission.enabled` to migrate stored blocks.
336+
To keep all shuffle data on Celeborn, set `spark.celeborn.client.spark.shuffle.fallback.policy=NEVER`.
337+
With `NEVER`, the application fails rather than falling back when Celeborn cannot serve a shuffle.
338+
Shuffle tracking can then be disabled to release idle executors promptly, but only when Spark's
339+
dynamic allocation check still passes without it: either use Spark 3.5 or later with
340+
`spark.shuffle.sort.io.plugin.class=org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO`, or
341+
enable Spark shuffle decommissioning as described above. Otherwise, for example on Spark 3.4 without
342+
an external shuffle service or shuffle decommissioning, Spark rejects dynamic allocation at startup.
343+
331344
Shuffle OSS storage related configurations:
332345
```properties
333346
# If you are using Celeborn for shuffle OSS storage, these settings will be needed.

‎client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/CelebornShuffleDataIO.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import org.slf4j.LoggerFactory;
2828

2929
import org.apache.celeborn.common.CelebornConf;
30+
import org.apache.celeborn.common.protocol.FallbackPolicy;
3031

3132
public class CelebornShuffleDataIO implements ShuffleDataIO {
3233

@@ -58,7 +59,10 @@ class CelebornShuffleDriverComponents extends LocalDiskShuffleDriverComponents {
5859
private final boolean supportsReliableStorage;
5960

6061
public CelebornShuffleDriverComponents(CelebornConf celebornConf) {
61-
this.supportsReliableStorage = !celebornConf.shuffleForceFallbackEnabled();
62+
// Reliable only under NEVER. AUTO/ALWAYS may fall back to local-disk shuffle, whose output
63+
// dies with the executor, so report false to stop DRA from reclaiming executors that hold it.
64+
this.supportsReliableStorage =
65+
FallbackPolicy.NEVER.equals(celebornConf.sparkShuffleFallbackPolicy());
6266
}
6367

6468
// Omitting @Override annotation to avoid compile error before Spark 3.5.0

‎client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java‎

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.concurrent.ConcurrentHashMap;
2424

2525
import org.apache.spark.*;
26+
import org.apache.spark.internal.config.ConfigEntry;
2627
import org.apache.spark.internal.config.package$;
2728
import org.apache.spark.launcher.SparkLauncher;
2829
import org.apache.spark.rdd.DeterministicLevel;
@@ -36,6 +37,7 @@
3637
import org.apache.celeborn.client.ShuffleClient;
3738
import org.apache.celeborn.client.security.CryptoHandler;
3839
import org.apache.celeborn.common.CelebornConf;
40+
import org.apache.celeborn.common.protocol.FallbackPolicy;
3941
import org.apache.celeborn.common.protocol.ShuffleMode;
4042
import org.apache.celeborn.reflect.DynMethods;
4143
import org.apache.celeborn.spark.FailedShuffleCleaner;
@@ -58,6 +60,22 @@ public class SparkShuffleManager implements ShuffleManager {
5860
private static final String SORT_SHUFFLE_MANAGER_NAME =
5961
"org.apache.spark.shuffle.sort.SortShuffleManager";
6062

63+
// Decommissioning configs do not exist in Spark 3.0. Read Spark's entries to preserve its
64+
// defaults.
65+
private static final ConfigEntry<Object> DECOMMISSION_ENABLED =
66+
DynMethods.builder("DECOMMISSION_ENABLED")
67+
.impl(package$.class)
68+
.orNoop()
69+
.build(package$.MODULE$)
70+
.invoke();
71+
72+
private static final ConfigEntry<Object> STORAGE_DECOMMISSION_SHUFFLE_BLOCKS_ENABLED =
73+
DynMethods.builder("STORAGE_DECOMMISSION_SHUFFLE_BLOCKS_ENABLED")
74+
.impl(package$.class)
75+
.orNoop()
76+
.build(package$.MODULE$)
77+
.invoke();
78+
6179
private static final boolean COLUMNAR_SHUFFLE_CLASSES_PRESENT;
6280

6381
static {
@@ -110,28 +128,39 @@ private Optional<CryptoHandler> getCryptoHandler() {
110128
return cryptoHandler;
111129
}
112130

131+
private static boolean isDecommissionShuffleBlocksEnabled(SparkConf conf) {
132+
return DECOMMISSION_ENABLED != null
133+
&& STORAGE_DECOMMISSION_SHUFFLE_BLOCKS_ENABLED != null
134+
&& (Boolean) conf.get(DECOMMISSION_ENABLED)
135+
&& (Boolean) conf.get(STORAGE_DECOMMISSION_SHUFFLE_BLOCKS_ENABLED);
136+
}
137+
138+
static boolean isUnsafeDraFallback(SparkConf conf) {
139+
return conf.getBoolean("spark.dynamicAllocation.enabled", false)
140+
&& !conf.getBoolean("spark.shuffle.service.enabled", false)
141+
&& !(Boolean) conf.get(package$.MODULE$.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED())
142+
&& !isDecommissionShuffleBlocksEnabled(conf);
143+
}
144+
113145
public SparkShuffleManager(SparkConf conf, boolean isDriver) {
114146
if (conf.getBoolean(SQLConf.LOCAL_SHUFFLE_READER_ENABLED().key(), true)) {
115147
logger.warn(
116148
"Detected {} (default is true) is enabled, it's highly recommended to disable it when "
117149
+ "use Celeborn as Remote Shuffle Service to avoid performance degradation.",
118150
SQLConf.LOCAL_SHUFFLE_READER_ENABLED().key());
119151
}
120-
if ((Boolean) conf.get(package$.MODULE$.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED())) {
121-
String key = package$.MODULE$.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED().key();
122-
Boolean defaultValue =
123-
(Boolean) package$.MODULE$.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED().defaultValue().get();
152+
this.celebornConf = SparkUtils.fromSparkConf(conf);
153+
boolean neverFallback = FallbackPolicy.NEVER.equals(celebornConf.sparkShuffleFallbackPolicy());
154+
if (!neverFallback && isUnsafeDraFallback(conf)) {
124155
logger.warn(
125-
"Detected {} (default is {}) is enabled, "
126-
+ "it's highly recommended to disable it when use Celeborn as Remote Shuffle Service "
127-
+ "to avoid performance degradation.",
128-
key,
129-
defaultValue);
156+
"DRA is enabled without the external shuffle service and fallback policy is not NEVER, "
157+
+ "but {} is disabled and shuffle decommissioning is not enabled. Enable tracking so fallback "
158+
+ "shuffle output is not lost when idle executors are reclaimed.",
159+
package$.MODULE$.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED().key());
130160
}
131161
SparkCommonUtils.validateAttemptConfig(conf);
132162
this.conf = conf;
133163
this.isDriver = isDriver;
134-
this.celebornConf = SparkUtils.fromSparkConf(conf);
135164
this.cores = executorCores(conf);
136165
this.fallbackPolicyRunner = new CelebornShuffleFallbackPolicyRunner(celebornConf);
137166
this.sendBufferPoolCheckInterval = celebornConf.clientPushSendBufferPoolExpireCheckInterval();
@@ -220,13 +249,14 @@ public <K, V, C> ShuffleHandle registerShuffle(
220249

221250
lifecycleManager.shuffleCount().increment();
222251
if (fallbackPolicyRunner.applyFallbackPolicies(dependency, lifecycleManager)) {
223-
if (conf.getBoolean("spark.dynamicAllocation.enabled", false)
224-
&& !conf.getBoolean("spark.shuffle.service.enabled", false)) {
252+
if (isUnsafeDraFallback(conf)) {
225253
logger.error(
226254
"DRA is enabled but we fallback to vanilla Spark SortShuffleManager for "
227-
+ "shuffle: {} due to fallback policy. It may cause block can not found when reducer "
228-
+ "task fetch data.",
229-
shuffleId);
255+
+ "shuffle: {} due to fallback policy, and {} is disabled while shuffle "
256+
+ "decommissioning is not enabled. Fallback shuffle blocks may be lost when idle "
257+
+ "executors are reclaimed.",
258+
shuffleId,
259+
package$.MODULE$.DYN_ALLOCATION_SHUFFLE_TRACKING_ENABLED().key());
230260
} else {
231261
logger.warn("Fallback to vanilla Spark SortShuffleManager for shuffle: {}", shuffleId);
232262
}

‎client-spark/spark-3/src/test/scala/org/apache/spark/shuffle/celeborn/CelebornShuffleManagerSuite.scala‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
package org.apache.spark.shuffle.celeborn
1919

20-
import org.apache.spark.{SparkConf, SparkContext}
20+
import org.apache.spark.{SparkConf, SparkContext, SparkVersionUtil}
2121
import org.apache.spark.sql.internal.SQLConf
2222
import org.junit
2323
import org.junit.Assert
@@ -92,6 +92,73 @@ class SparkShuffleManagerSuite extends Logging {
9292
sc.stop()
9393
}
9494

95+
@junit.Test
96+
def testReliableStorageFollowsFallbackPolicy(): Unit = {
97+
// supportsReliableStorage must be true only for NEVER; AUTO/ALWAYS may fall back to local-disk
98+
// shuffle, so reporting reliable storage would let DRA reclaim executors holding fallback data.
99+
def reliable(policy: String): Boolean = {
100+
val conf = new SparkConf(false)
101+
.set(s"spark.${CelebornConf.SPARK_SHUFFLE_FALLBACK_POLICY.key}", policy)
102+
new CelebornShuffleDataIO(conf).driver()
103+
.asInstanceOf[CelebornShuffleDriverComponents].supportsReliableStorage()
104+
}
105+
Assert.assertTrue(reliable("NEVER"))
106+
Assert.assertFalse(reliable("AUTO"))
107+
Assert.assertFalse(reliable("ALWAYS"))
108+
Assert.assertFalse(
109+
new CelebornShuffleDataIO(new SparkConf(false)).driver()
110+
.asInstanceOf[CelebornShuffleDriverComponents].supportsReliableStorage())
111+
}
112+
113+
@junit.Test
114+
def testUnsafeDraFallback(): Unit = {
115+
for {
116+
dra <- Seq(false, true)
117+
shuffleService <- Seq(false, true)
118+
tracking <- Seq(false, true)
119+
decommission <- Seq(false, true)
120+
shuffleBlocks <- Seq(false, true)
121+
} {
122+
val conf = new SparkConf(false)
123+
.set("spark.dynamicAllocation.enabled", dra.toString)
124+
.set("spark.shuffle.service.enabled", shuffleService.toString)
125+
.set("spark.dynamicAllocation.shuffleTracking.enabled", tracking.toString)
126+
.set("spark.decommission.enabled", decommission.toString)
127+
.set("spark.storage.decommission.shuffleBlocks.enabled", shuffleBlocks.toString)
128+
val decommissionSupported = SparkVersionUtil.isGreaterThan(3, 0)
129+
val expected = dra && !shuffleService && !tracking &&
130+
!(decommissionSupported && decommission && shuffleBlocks)
131+
Assert.assertEquals(
132+
conf.toDebugString,
133+
expected,
134+
SparkShuffleManager.isUnsafeDraFallback(conf))
135+
}
136+
}
137+
138+
@junit.Test
139+
def testUnsafeDraFallbackUsesSparkDefaults(): Unit = {
140+
val conf = new SparkConf(false)
141+
.set("spark.dynamicAllocation.enabled", "true")
142+
.set("spark.shuffle.service.enabled", "false")
143+
.set("spark.dynamicAllocation.shuffleTracking.enabled", "false")
144+
Assert.assertTrue(SparkShuffleManager.isUnsafeDraFallback(conf))
145+
146+
conf.set("spark.decommission.enabled", "true")
147+
// Spark 3.4 changed shuffle-block decommissioning's default from false to true.
148+
Assert.assertEquals(
149+
!SparkVersionUtil.isGreaterThan(3, 3),
150+
SparkShuffleManager.isUnsafeDraFallback(conf))
151+
152+
conf.remove("spark.decommission.enabled")
153+
conf.set("spark.storage.decommission.shuffleBlocks.enabled", "true")
154+
Assert.assertTrue(SparkShuffleManager.isUnsafeDraFallback(conf))
155+
156+
conf.remove("spark.dynamicAllocation.shuffleTracking.enabled")
157+
Assert.assertEquals(
158+
!SparkVersionUtil.isGreaterThan(3, 3),
159+
SparkShuffleManager.isUnsafeDraFallback(conf))
160+
}
161+
95162
@junit.Test
96163
def testWrongSparkConfMaxAttemptLimit(): Unit = {
97164
val conf = new SparkConf().setIfMissing("spark.master", "local")

‎docs/deploy.md‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -208,14 +208,27 @@ spark.sql.adaptive.skewJoin.enabled true
208208
# Support Spark Dynamic Resource Allocation
209209
# Required Spark version >= 3.5.0
210210
spark.shuffle.sort.io.plugin.class org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO
211-
# Required Spark version >= 3.4.0, highly recommended to disable
212-
spark.dynamicAllocation.shuffleTracking.enabled false
211+
# Available since Spark 3.0; enabled by default since Spark 3.4.
212+
# With AUTO (default) or ALWAYS fallback, keep tracking enabled to protect local-disk shuffle data.
213+
spark.dynamicAllocation.shuffleTracking.enabled true
213214

214215
# Support ShuffleManager when defined in user jars
215216
# Required Spark version < 4.0.0 or without SPARK-45762, highly recommended to false for ShuffleManager in user-defined jar specified by --jars or spark.jars
216217
spark.executor.userClassPathFirst false
217218
```
218219

220+
With DRA enabled, `AUTO` and `ALWAYS` fallback policies require shuffle tracking unless an external
221+
shuffle service or Spark shuffle decommissioning is configured. Shuffle decommissioning requires
222+
`spark.decommission.enabled=true` and `spark.storage.decommission.shuffleBlocks.enabled=true`;
223+
also enable `spark.storage.decommission.enabled` to migrate stored blocks.
224+
To keep all shuffle data on Celeborn, set `spark.celeborn.client.spark.shuffle.fallback.policy=NEVER`.
225+
With `NEVER`, the application fails rather than falling back when Celeborn cannot serve a shuffle.
226+
Shuffle tracking can then be disabled to release idle executors promptly, but only when Spark's
227+
dynamic allocation check still passes without it: either use Spark 3.5 or later with
228+
`spark.shuffle.sort.io.plugin.class=org.apache.spark.shuffle.celeborn.CelebornShuffleDataIO`, or
229+
enable Spark shuffle decommissioning as described above. Otherwise, for example on Spark 3.4 without
230+
an external shuffle service or shuffle decommissioning, Spark rejects dynamic allocation at startup.
231+
219232
## Deploy Flink client
220233

221234
**Important: Only Flink batch jobs are supported for now. Due to the Shuffle Service in Flink is cluster-granularity, if you want to use Celeborn in a session cluster, it will not be able to submit both streaming and batch job to the same cluster. We plan to get rid of this restriction for Hybrid Shuffle mode in a future release.**

0 commit comments

Comments
 (0)