Add SparkCommonUtils.validateAttemptConfig

apache · Aug 12, 2024 · c0243d0 · c0243d0
1 parent e0affd7
commit c0243d0
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 18 deletions.
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SparkCommonUtils.java
@@ -17,11 +17,17 @@
 
 package org.apache.spark.shuffle.celeborn;
 
+import org.apache.spark.SparkConf;
 import org.apache.spark.TaskContext;
+import org.apache.spark.scheduler.DAGScheduler;
 
 public class SparkCommonUtils {
-  public static void validateMaxAttempts(int maxStageAttempts, int maxTaskAttempts)
-      throws IllegalArgumentException {
+  public static void validateAttemptConfig(SparkConf conf) throws IllegalArgumentException {
+    int maxStageAttempts =
+        conf.getInt(
+            "spark.stage.maxConsecutiveAttempts",
+            DAGScheduler.DEFAULT_MAX_CONSECUTIVE_STAGE_ATTEMPTS());
+    int maxTaskAttempts = conf.getInt("spark.task.maxFailures", 4);
     if (maxStageAttempts >= (1 << 15) || maxTaskAttempts >= (1 << 16)) {
       // The map attemptId is a non-negative number constructed from
       // both stageAttemptNumber and taskAttemptNumber.

diff --git a/...nt-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/...nt-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -27,7 +27,6 @@
 import org.apache.spark.internal.config.package$;
 import org.apache.spark.launcher.SparkLauncher;
 import org.apache.spark.rdd.DeterministicLevel;
-import org.apache.spark.scheduler.DAGScheduler;
 import org.apache.spark.shuffle.*;
 import org.apache.spark.shuffle.sort.SortShuffleManager;
 import org.apache.spark.util.Utils;
@@ -67,12 +66,7 @@ public class SparkShuffleManager implements ShuffleManager {
   private ExecutorShuffleIdTracker shuffleIdTracker = new ExecutorShuffleIdTracker();
 
   public SparkShuffleManager(SparkConf conf, boolean isDriver) {
-    int maxStageAttempts =
-        conf.getInt(
-            "spark.stage.maxConsecutiveAttempts",
-            DAGScheduler.DEFAULT_MAX_CONSECUTIVE_STAGE_ATTEMPTS());
-    int maxTaskAttempts = (Integer) conf.get(package$.MODULE$.MAX_TASK_FAILURES());
-    SparkCommonUtils.validateMaxAttempts(maxStageAttempts, maxTaskAttempts);
+    SparkCommonUtils.validateAttemptConfig(conf);
     this.conf = conf;
     this.isDriver = isDriver;
     this.celebornConf = SparkUtils.fromSparkConf(conf);

diff --git a/...nt-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java b/...nt-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java
@@ -25,7 +25,6 @@
 import org.apache.spark.internal.config.package$;
 import org.apache.spark.launcher.SparkLauncher;
 import org.apache.spark.rdd.DeterministicLevel;
-import org.apache.spark.scheduler.DAGScheduler;
 import org.apache.spark.shuffle.*;
 import org.apache.spark.shuffle.sort.SortShuffleManager;
 import org.apache.spark.sql.internal.SQLConf;
@@ -108,12 +107,7 @@ public SparkShuffleManager(SparkConf conf, boolean isDriver) {
           key,
           defaultValue);
     }
-    int maxStageAttempts =
-        conf.getInt(
-            "spark.stage.maxConsecutiveAttempts",
-            DAGScheduler.DEFAULT_MAX_CONSECUTIVE_STAGE_ATTEMPTS());
-    int maxTaskAttempts = (Integer) conf.get(package$.MODULE$.TASK_MAX_FAILURES());
-    SparkCommonUtils.validateMaxAttempts(maxStageAttempts, maxTaskAttempts);
+    SparkCommonUtils.validateAttemptConfig(conf);
     this.conf = conf;
     this.isDriver = isDriver;
     this.celebornConf = SparkUtils.fromSparkConf(conf);

diff --git a/...s/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornFetchFailureSuite.scala b/...s/spark-it/src/test/scala/org/apache/celeborn/tests/spark/CelebornFetchFailureSuite.scala
@@ -24,7 +24,6 @@ import java.util.concurrent.atomic.AtomicBoolean
 import org.apache.spark.{BarrierTaskContext, ShuffleDependency, SparkConf, SparkContextHelper, SparkException, TaskContext}
 import org.apache.spark.celeborn.ExceptionMakerHelper
 import org.apache.spark.rdd.RDD
-import org.apache.spark.{SparkConf, SparkContextHelper, TaskContext}
 import org.apache.spark.shuffle.ShuffleHandle
 import org.apache.spark.shuffle.celeborn.{CelebornShuffleHandle, ShuffleManagerHook, SparkShuffleManager, SparkUtils, TestCelebornShuffleManager}
 import org.apache.spark.sql.{Row, RowFactory, SparkSession}
@@ -36,7 +35,6 @@ import org.scalatest.funsuite.AnyFunSuite
 
 import org.apache.celeborn.client.ShuffleClient
 import org.apache.celeborn.common.CelebornConf
-import org.apache.celeborn.common.exception.CelebornIOException
 import org.apache.celeborn.common.protocol.ShuffleMode
 import org.apache.celeborn.service.deploy.worker.Worker