[CELEBORN-2230] SparkUtils#shouldReportShuffleFetchFailure method should retrieve the number of task failures from TaskSetManager #3650
+260
−47
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -383,6 +383,11 @@ public static void cancelShuffle(int shuffleId, String reason) { | |
| .hiddenImpl(TaskSetManager.class, "taskInfos") | ||
| .defaultAlwaysNull() | ||
| .build(); | ||
| private static final DynFields.UnboundField<int[]> TASK_FAILURES = | ||
| DynFields.builder() | ||
| .hiddenImpl(TaskSetManager.class, "numFailures") | ||
| .defaultAlwaysNull() | ||
| .build(); | ||
|
|
||
| /** | ||
| * TaskSetManager - it is not designed to be used outside the spark scheduler. Please be careful. | ||
|
|
@@ -420,6 +425,39 @@ protected static Tuple2<TaskInfo, List<TaskInfo>> getTaskAttempts( | |
| } | ||
| } | ||
|
|
||
| /** | ||
| * Gets the number of task attempts that have already failed for the given task index. Note: This | ||
| * count does NOT include the current failure. To get the total failure count including the | ||
| * current attempt, you need to add 1 to the returned value. | ||
| * | ||
| * @param taskSetManager the TaskSetManager to query | ||
| * @param index the task index | ||
| * @return the number of previous failed attempts, or -1 if an error occurs | ||
| */ | ||
| @VisibleForTesting | ||
| protected static int getTaskFailureCount(TaskSetManager taskSetManager, int index) { | ||
| if (taskSetManager == null) { | ||
| LOG.error("TaskSetManager is null for task index: {}", index); | ||
| return -1; | ||
| } | ||
|
|
||
| int[] numFailures = TASK_FAILURES.bind(taskSetManager).get(); | ||
| if (numFailures == null) { | ||
| LOG.error("Failed to get numFailures array from TaskSetManager for task index: {}", index); | ||
| return -1; | ||
| } | ||
|
|
||
| if (index < 0 || index >= numFailures.length) { | ||
| LOG.error( | ||
| "Task index {} is out of bounds for numFailures array (length: {})", | ||
| index, | ||
| numFailures.length); | ||
| return -1; | ||
| } | ||
|
|
||
| return numFailures[index]; | ||
| } | ||
|
|
||
| protected static Map<String, Set<Long>> reportedStageShuffleFetchFailureTaskIds = | ||
| JavaUtils.newConcurrentHashMap(); | ||
|
|
||
|
|
@@ -480,7 +518,6 @@ public static boolean shouldReportShuffleFetchFailure(long taskId) { | |
| if (taskAttempts == null) return true; | ||
|
|
||
| TaskInfo taskInfo = taskAttempts._1(); | ||
| boolean hasRunningAttempt = false; | ||
| for (TaskInfo ti : taskAttempts._2()) { | ||
| if (ti.taskId() != taskId) { | ||
|
|
@@ -492,7 +529,6 @@ public static boolean shouldReportShuffleFetchFailure(long taskId) { | |
| taskId, | ||
| taskInfo.attemptNumber(), | ||
| ti.attemptNumber()); | ||
| } else if (ti.successful()) { | ||
| LOG.info( | ||
| "StageId={} index={} taskId={} attempt={} another attempt {} is successful.", | ||
|
|
@@ -511,36 +547,32 @@ public static boolean shouldReportShuffleFetchFailure(long taskId) { | |
| taskInfo.attemptNumber(), | ||
| ti.attemptNumber()); | ||
| hasRunningAttempt = true; | ||
| } | ||
| } | ||
| } | ||
| // The following situations should trigger a FetchFailed exception: | ||
| // 1. If total failures (previous failures + current failure) >= maxTaskFails | ||
| // 2. If no other taskAttempts are running, trigger a FetchFailed exception | ||
| // to keep the same behavior as Spark. | ||
| // Note: previousFailureCount does NOT include the current failure, | ||
| // so (previousFailureCount + 1) represents the total failure count. | ||
| int previousFailureCount = getTaskFailureCount(taskSetManager, taskInfo.index()); | ||
| // Fail-safe: if failure count cannot be determined, conservatively trigger | ||
| // FetchFailed to avoid silently swallowing the error. | ||
| if (previousFailureCount < 0) { | ||
| return true; | ||
| } | ||
|
Comment on lines
+559
to
+587
There was a problem hiding this comment. When getTaskFailureCount() fails (returns < 0), shouldReportShuffleFetchFailure() immediately returns true. This makes the pre-check aggressively report FetchFailed even if there are other running attempts and the retry limit has not been reached, which can reintroduce premature stage reruns in exactly the scenarios this change is trying to avoid (e.g., if reflective access to TaskSetManager.numFailures breaks on some Spark builds). Consider falling back to the previous attempt/status-based counting (or at least gating on !hasRunningAttempt) instead of unconditional true, and log at WARN once to avoid error spam if the field is unavailable.
Contributor
Author
There was a problem hiding this comment. I think this is unreasonable. It shouldn't check |
||
| if (previousFailureCount + 1 >= maxTaskFails || !hasRunningAttempt) { | ||
| LOG.warn( | ||
| "StageId={}, index={}, taskId={}, attemptNumber={}: Previous failure count {} " | ||
| + "(total with current: {}) reached maximum allowed failures {} " | ||
| + "or no running attempt exists.", | ||
| stageId, | ||
| taskInfo.index(), | ||
| taskId, | ||
| taskInfo.attemptNumber(), | ||
| previousFailureCount, | ||
| previousFailureCount + 1, | ||
| maxTaskFails); | ||
| return true; | ||
| } else { | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,7 +19,7 @@ package org.apache.spark.shuffle.celeborn | |
|
|
||
| import scala.collection.JavaConverters._ | ||
|
|
||
| import org.apache.spark.{SparkConf, TaskContext} | ||
| import org.apache.spark.scheduler.{TaskSchedulerImpl, TaskSetManager} | ||
| import org.apache.spark.sql.SparkSession | ||
| import org.scalatest.BeforeAndAfterEach | ||
|
|
@@ -216,6 +216,107 @@ class SparkUtilsSuite extends AnyFunSuite | |
| } | ||
| } | ||
|
|
||
| test("getTaskFailureCount") { | ||
| assert(SparkUtils.getTaskFailureCount(null, 0) == -1) | ||
|
|
||
| if (Spark3OrNewer) { | ||
| val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[2,3]") | ||
| val sparkSession = SparkSession.builder() | ||
| .config(updateSparkConf(sparkConf, ShuffleMode.HASH)) | ||
| .config("spark.sql.shuffle.partitions", 2) | ||
| .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false) | ||
| .config("spark.celeborn.client.spark.stageRerun.enabled", "true") | ||
| .config( | ||
| "spark.shuffle.manager", | ||
| "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager") | ||
| .getOrCreate() | ||
|
|
||
| try { | ||
| val sc = sparkSession.sparkContext | ||
| val jobThread = new Thread { | ||
| override def run(): Unit = { | ||
| try { | ||
| sc.parallelize(1 to 100, 2) | ||
| .repartition(1) | ||
| .mapPartitions { iter => | ||
| Thread.sleep(3000) | ||
| iter | ||
| }.collect() | ||
| } catch { | ||
| case _: InterruptedException => | ||
| } | ||
| } | ||
| } | ||
| jobThread.start() | ||
|
|
||
| val taskScheduler = sc.taskScheduler.asInstanceOf[TaskSchedulerImpl] | ||
| eventually(timeout(3.seconds), interval(100.milliseconds)) { | ||
| val taskId = 0 | ||
| val taskSetManager = SparkUtils.getTaskSetManager(taskScheduler, taskId) | ||
| assert(taskSetManager != null) | ||
| assert(SparkUtils.getTaskFailureCount(taskSetManager, 0) == 0) | ||
| assert(SparkUtils.getTaskFailureCount(taskSetManager, -1) == -1) | ||
| assert(SparkUtils.getTaskFailureCount(taskSetManager, Int.MaxValue) == -1) | ||
| } | ||
|
|
||
| sparkSession.sparkContext.cancelAllJobs() | ||
| jobThread.interrupt() | ||
| } finally { | ||
| sparkSession.stop() | ||
| } | ||
| } | ||
| } | ||
|
|
||
| test("getTaskFailureCount after real task failures") { | ||
| if (Spark3OrNewer) { | ||
| // local[1,4]: 1 core (sequential execution), max 4 task failures before stage abort | ||
| val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[1,4]") | ||
| val sparkSession = SparkSession.builder() | ||
| .config(updateSparkConf(sparkConf, ShuffleMode.HASH)) | ||
| .config("spark.sql.shuffle.partitions", 2) | ||
| .config("spark.celeborn.shuffle.forceFallback.partition.enabled", false) | ||
| .config("spark.celeborn.client.spark.stageRerun.enabled", "true") | ||
| .config( | ||
| "spark.shuffle.manager", | ||
| "org.apache.spark.shuffle.celeborn.TestCelebornShuffleManager") | ||
| .getOrCreate() | ||
|
|
||
| try { | ||
| val sc = sparkSession.sparkContext | ||
|
|
||
| val jobThread = new Thread { | ||
| override def run(): Unit = { | ||
| try { | ||
| sc.parallelize(1 to 10, 1).mapPartitions { iter => | ||
| if (TaskContext.get().attemptNumber() < 2) { | ||
| throw new RuntimeException("Simulated task failure") | ||
| } | ||
| Thread.sleep(10000) | ||
| iter | ||
| }.collect() | ||
| } catch { | ||
| case _: Exception => | ||
| } | ||
| } | ||
| } | ||
| jobThread.start() | ||
|
|
||
| val taskScheduler = sc.taskScheduler.asInstanceOf[TaskSchedulerImpl] | ||
| eventually(timeout(10.seconds), interval(100.milliseconds)) { | ||
| // taskId 0,1 failed and removed; taskId 2 is the surviving 3rd attempt | ||
| val taskSetManager = SparkUtils.getTaskSetManager(taskScheduler, 2) | ||
| assert(taskSetManager != null) | ||
| assert(SparkUtils.getTaskFailureCount(taskSetManager, 0) == 2) | ||
|
Comment on lines
+304
to
+311
There was a problem hiding this comment. This test assumes the third attempt’s taskId will be exactly 2 (after two failures). Spark task IDs are globally assigned within a SparkContext and aren’t guaranteed to align with attempt count if any other tasks/stages run (including internal ones), which can make the test brittle across Spark versions/configs. Consider deriving the taskId dynamically (e.g., capturing TaskContext.taskAttemptId() via an accumulator/Promise, or scanning taskScheduler’s taskIdToTaskSetManager for the active TaskSetManager) instead of hardcoding 2. |
||
| } | ||
|
|
||
| sparkSession.sparkContext.cancelAllJobs() | ||
| jobThread.interrupt() | ||
| } finally { | ||
| sparkSession.stop() | ||
| } | ||
| } | ||
| } | ||
|
|
||
| test("serialize/deserialize GetReducerFileGroupResponse with broadcast") { | ||
| val sparkConf = new SparkConf().setAppName("rss-demo").setMaster("local[2,3]") | ||
| val sparkSession = SparkSession.builder() | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When getTaskFailureCount() fails (returns < 0), shouldReportShuffleFetchFailure() immediately returns true. This makes the pre-check aggressively report FetchFailed even if there are other running attempts and the retry limit has not been reached, which can reintroduce premature stage reruns if reflective access to TaskSetManager.numFailures fails on some Spark builds. Consider falling back to the previous attempt/status-based counting (or at least gating on !hasRunningAttempt) rather than unconditional true, and avoid logging this as an error on every call if the field is unavailable.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this is unreasonable. It shouldn't check
!hasRunningAttempt, but should directly returntrue. Otherwise,FetchFailedwon't be triggered and app will fail.