Use task ID in output committer instead of stage / task attempt numbers.

Author: Marcelo Vanzin

core/src/main/scala/org/apache/spark/executor/CommitDeniedException.scala

@@ -26,8 +26,8 @@ private[spark] class CommitDeniedException(
     msg: String,
     jobID: Int,
     splitID: Int,
-    attemptNumber: Int)
+    taskId: Long)
   extends Exception(msg) {
 
-  def toTaskCommitDeniedReason: TaskCommitDenied = TaskCommitDenied(jobID, splitID, attemptNumber)
+  def toTaskCommitDeniedReason: TaskCommitDenied = TaskCommitDenied(jobID, splitID, taskId.toInt)
 }

core/src/main/scala/org/apache/spark/mapred/SparkHadoopMapRedUtil.scala

@@ -70,8 +70,8 @@ object SparkHadoopMapRedUtil extends Logging {
       if (shouldCoordinateWithDriver) {
         val outputCommitCoordinator = SparkEnv.get.outputCommitCoordinator
         val ctx = TaskContext.get()
-        val canCommit = outputCommitCoordinator.canCommit(ctx.stageId(), ctx.stageAttemptNumber(),
-          splitId, ctx.attemptNumber())
+        val canCommit = outputCommitCoordinator.canCommit(ctx.stageId(), splitId,
+          ctx.taskAttemptId())
 
         if (canCommit) {
           performCommit()
@@ -81,7 +81,7 @@ object SparkHadoopMapRedUtil extends Logging {
           logInfo(message)
           // We need to abort the task so that the driver can reschedule new attempts, if necessary
           committer.abortTask(mrTaskContext)
-          throw new CommitDeniedException(message, ctx.stageId(), splitId, ctx.attemptNumber())
+          throw new CommitDeniedException(message, ctx.stageId(), splitId, ctx.taskAttemptId())
         }
       } else {
         // Speculation is disabled or a user has chosen to manually bypass the commit coordination

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -1171,9 +1171,8 @@ class DAGScheduler(
 
     outputCommitCoordinator.taskCompleted(
       stageId,
-      task.stageAttemptId,
       task.partitionId,
-      event.taskInfo.attemptNumber, // this is a task attempt number
+      event.taskInfo.taskId,
       event.reason)
 
     if (!stageIdToStage.contains(task.stageId)) {

core/src/main/scala/org/apache/spark/scheduler/OutputCommitCoordinator.scala

@@ -29,9 +29,8 @@ private sealed trait OutputCommitCoordinationMessage extends Serializable
 private case object StopCoordinator extends OutputCommitCoordinationMessage
 private case class AskPermissionToCommitOutput(
     stage: Int,
-    stageAttempt: Int,
     partition: Int,
-    attemptNumber: Int)
+    taskId: Long)
 
 /**
  * Authority that decides whether tasks can commit output to HDFS. Uses a "first committer wins"
@@ -49,15 +48,11 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
   // Initialized by SparkEnv
   var coordinatorRef: Option[RpcEndpointRef] = None
 
-  // Class used to identify a committer. The task ID for a committer is implicitly defined by
-  // the partition being processed, but the coordinator needs to keep track of both the stage
-  // attempt and the task attempt, because in some situations the same task may be running
-  // concurrently in two different attempts of the same stage.
-  private case class TaskIdentifier(stageAttempt: Int, taskAttempt: Int)
+  private val NO_AUTHORIZED_COMMITTER: Long = -1L
 
   private case class StageState(numPartitions: Int) {
-    val authorizedCommitters = Array.fill[TaskIdentifier](numPartitions)(null)
-    val failures = mutable.Map[Int, mutable.Set[TaskIdentifier]]()
+    val authorizedCommitters = Array.fill[Long](numPartitions)(NO_AUTHORIZED_COMMITTER)
+    val failures = mutable.Map[Int, mutable.Set[Long]]()
   }
 
   /**
@@ -88,16 +83,14 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
    *
    * @param stage the stage number
    * @param partition the partition number
-   * @param attemptNumber how many times this task has been attempted
-   *                      (see [[TaskContext.attemptNumber()]])
+   * @param taskId the task asking for permission to commit
    * @return true if this task is authorized to commit, false otherwise
    */
   def canCommit(
       stage: Int,
-      stageAttempt: Int,
       partition: Int,
-      attemptNumber: Int): Boolean = {
-    val msg = AskPermissionToCommitOutput(stage, stageAttempt, partition, attemptNumber)
+      taskId: Long): Boolean = {
+    val msg = AskPermissionToCommitOutput(stage, partition, taskId)
     coordinatorRef match {
       case Some(endpointRef) =>
         ThreadUtils.awaitResult(endpointRef.ask[Boolean](msg),
@@ -136,9 +129,8 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
   // Called by DAGScheduler
   private[scheduler] def taskCompleted(
       stage: Int,
-      stageAttempt: Int,
       partition: Int,
-      attemptNumber: Int,
+      taskId: Long,
       reason: TaskEndReason): Unit = synchronized {
     val stageState = stageStates.getOrElse(stage, {
       logDebug(s"Ignoring task completion for completed stage")
@@ -148,16 +140,14 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
       case Success =>
       // The task output has been committed successfully
       case _: TaskCommitDenied =>
-        logInfo(s"Task was denied committing, stage: $stage.$stageAttempt, " +
-          s"partition: $partition, attempt: $attemptNumber")
+        logInfo(s"Task was denied committing, stage: $stage, partition: $partition, task: $taskId")
       case _ =>
         // Mark the attempt as failed to blacklist from future commit protocol
-        val taskId = TaskIdentifier(stageAttempt, attemptNumber)
         stageState.failures.getOrElseUpdate(partition, mutable.Set()) += taskId
         if (stageState.authorizedCommitters(partition) == taskId) {
-          logDebug(s"Authorized committer (attemptNumber=$attemptNumber, stage=$stage, " +
+          logDebug(s"Authorized committer (taskId=$taskId, stage=$stage, " +
             s"partition=$partition) failed; clearing lock")
-          stageState.authorizedCommitters(partition) = null
+          stageState.authorizedCommitters(partition) = NO_AUTHORIZED_COMMITTER
         }
     }
   }
@@ -173,40 +163,35 @@ private[spark] class OutputCommitCoordinator(conf: SparkConf, isDriver: Boolean)
   // Marked private[scheduler] instead of private so this can be mocked in tests
   private[scheduler] def handleAskPermissionToCommit(
       stage: Int,
-      stageAttempt: Int,
       partition: Int,
-      attemptNumber: Int): Boolean = synchronized {
+      taskId: Long): Boolean = synchronized {
     stageStates.get(stage) match {
-      case Some(state) if attemptFailed(state, stageAttempt, partition, attemptNumber) =>
-        logInfo(s"Commit denied for stage=$stage.$stageAttempt, partition=$partition: " +
-          s"task attempt $attemptNumber already marked as failed.")
+      case Some(state) if attemptFailed(state, partition, taskId) =>
+        logInfo(s"Commit denied for stage=$stage, partition=$partition: " +
+          s"task $taskId already marked as failed.")
         false
       case Some(state) =>
         val existing = state.authorizedCommitters(partition)
-        if (existing == null) {
-          logDebug(s"Commit allowed for stage=$stage.$stageAttempt, partition=$partition, " +
-            s"task attempt $attemptNumber")
-          state.authorizedCommitters(partition) = TaskIdentifier(stageAttempt, attemptNumber)
+        if (existing == NO_AUTHORIZED_COMMITTER) {
+          logDebug(s"Commit allowed for stage=$stage, partition=$partition, task $taskId")
+          state.authorizedCommitters(partition) = taskId
           true
         } else {
-          logDebug(s"Commit denied for stage=$stage.$stageAttempt, partition=$partition: " +
+          logDebug(s"Commit denied for stage=$stage, partition=$partition: " +
             s"already committed by $existing")
           false
         }
       case None =>
-        logDebug(s"Commit denied for stage=$stage.$stageAttempt, partition=$partition: " +
-          "stage already marked as completed.")
+        logDebug(s"Commit denied for stage=$stage, partition=$partition: stage already completed.")
         false
     }
   }
 
   private def attemptFailed(
       stageState: StageState,
-      stageAttempt: Int,
       partition: Int,
-      attempt: Int): Boolean = synchronized {
-    val failInfo = TaskIdentifier(stageAttempt, attempt)
-    stageState.failures.get(partition).exists(_.contains(failInfo))
+      taskId: Long): Boolean = synchronized {
+    stageState.failures.get(partition).exists(_.contains(taskId))
   }
 }
 
@@ -226,10 +211,9 @@ private[spark] object OutputCommitCoordinator {
     }
 
     override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
-      case AskPermissionToCommitOutput(stage, stageAttempt, partition, attemptNumber) =>
+      case AskPermissionToCommitOutput(stage, partition, taskId) =>
         context.reply(
-          outputCommitCoordinator.handleAskPermissionToCommit(stage, stageAttempt, partition,
-            attemptNumber))
+          outputCommitCoordinator.handleAskPermissionToCommit(stage, partition, taskId))
     }
   }
 }

core/src/test/scala/org/apache/spark/scheduler/OutputCommitCoordinatorSuite.scala

@@ -154,7 +154,7 @@ class OutputCommitCoordinatorSuite extends SparkFunSuite with BeforeAndAfter {
   test("Job should not complete if all commits are denied") {
     // Create a mock OutputCommitCoordinator that denies all attempts to commit
     doReturn(false).when(outputCommitCoordinator).handleAskPermissionToCommit(
-      Matchers.any(), Matchers.any(), Matchers.any(), Matchers.any())
+      Matchers.any(), Matchers.any(), Matchers.any())
     val rdd: RDD[Int] = sc.parallelize(Seq(1), 1)
     def resultHandler(x: Int, y: Unit): Unit = {}
     val futureAction: SimpleFutureAction[Unit] = sc.submitJob[Int, Unit, Unit](rdd,
@@ -170,73 +170,65 @@ class OutputCommitCoordinatorSuite extends SparkFunSuite with BeforeAndAfter {
 
   test("Only authorized committer failures can clear the authorized committer lock (SPARK-6614)") {
     val stage: Int = 1
-    val stageAttempt: Int = 1
     val partition: Int = 2
-    val authorizedCommitter: Int = 3
-    val nonAuthorizedCommitter: Int = 100
+    val authorizedCommitter: Long = 3
+    val nonAuthorizedCommitter: Long = 100
     outputCommitCoordinator.stageStart(stage, maxPartitionId = 2)
 
-    assert(outputCommitCoordinator.canCommit(stage, stageAttempt, partition, authorizedCommitter))
-    assert(!outputCommitCoordinator.canCommit(stage, stageAttempt, partition,
-      nonAuthorizedCommitter))
+    assert(outputCommitCoordinator.canCommit(stage, partition, authorizedCommitter))
+    assert(!outputCommitCoordinator.canCommit(stage, partition, nonAuthorizedCommitter))
     // The non-authorized committer fails
-    outputCommitCoordinator.taskCompleted(stage, stageAttempt, partition,
-      attemptNumber = nonAuthorizedCommitter, reason = TaskKilled("test"))
+    outputCommitCoordinator.taskCompleted(stage, partition, nonAuthorizedCommitter,
+      reason = TaskKilled("test"))
     // New tasks should still not be able to commit because the authorized committer has not failed
-    assert(!outputCommitCoordinator.canCommit(stage, stageAttempt, partition,
-      nonAuthorizedCommitter + 1))
+    assert(!outputCommitCoordinator.canCommit(stage, partition, nonAuthorizedCommitter + 1))
     // The authorized committer now fails, clearing the lock
-    outputCommitCoordinator.taskCompleted(stage, stageAttempt, partition,
-      attemptNumber = authorizedCommitter, reason = TaskKilled("test"))
+    outputCommitCoordinator.taskCompleted(stage, partition, authorizedCommitter,
+      reason = TaskKilled("test"))
     // A new task should now be allowed to become the authorized committer
-    assert(outputCommitCoordinator.canCommit(stage, stageAttempt, partition,
-      nonAuthorizedCommitter + 2))
+    assert(outputCommitCoordinator.canCommit(stage, partition, nonAuthorizedCommitter + 2))
     // There can only be one authorized committer
-    assert(!outputCommitCoordinator.canCommit(stage, stageAttempt, partition,
-      nonAuthorizedCommitter + 3))
+    assert(!outputCommitCoordinator.canCommit(stage, partition, nonAuthorizedCommitter + 3))
   }
 
   test("SPARK-19631: Do not allow failed attempts to be authorized for committing") {
     val stage: Int = 1
-    val stageAttempt: Int = 1
     val partition: Int = 1
-    val failedAttempt: Int = 0
+    val failedAttempt: Long = 0L
     outputCommitCoordinator.stageStart(stage, maxPartitionId = 1)
-    outputCommitCoordinator.taskCompleted(stage, stageAttempt, partition,
-      attemptNumber = failedAttempt,
+    outputCommitCoordinator.taskCompleted(stage, partition, failedAttempt,
       reason = ExecutorLostFailure("0", exitCausedByApp = true, None))
-    assert(!outputCommitCoordinator.canCommit(stage, stageAttempt, partition, failedAttempt))
-    assert(outputCommitCoordinator.canCommit(stage, stageAttempt, partition, failedAttempt + 1))
+    assert(!outputCommitCoordinator.canCommit(stage, partition, failedAttempt))
+    assert(outputCommitCoordinator.canCommit(stage, partition, failedAttempt + 1))
   }
 
   test("SPARK-24589: Differentiate tasks from different stage attempts") {
     var stage = 1
-    val taskAttempt = 1
     val partition = 1
 
     outputCommitCoordinator.stageStart(stage, maxPartitionId = 1)
-    assert(outputCommitCoordinator.canCommit(stage, 1, partition, taskAttempt))
-    assert(!outputCommitCoordinator.canCommit(stage, 2, partition, taskAttempt))
+    assert(outputCommitCoordinator.canCommit(stage, partition, 1L))
+    assert(!outputCommitCoordinator.canCommit(stage, partition, 2L))
 
     // Fail the task in the first attempt, the task in the second attempt should succeed.
     stage += 1
     outputCommitCoordinator.stageStart(stage, maxPartitionId = 1)
-    outputCommitCoordinator.taskCompleted(stage, 1, partition, taskAttempt,
+    outputCommitCoordinator.taskCompleted(stage, partition, 1L,
       ExecutorLostFailure("0", exitCausedByApp = true, None))
-    assert(!outputCommitCoordinator.canCommit(stage, 1, partition, taskAttempt))
-    assert(outputCommitCoordinator.canCommit(stage, 2, partition, taskAttempt))
+    assert(!outputCommitCoordinator.canCommit(stage, partition, 1L))
+    assert(outputCommitCoordinator.canCommit(stage, partition, 2L))
 
     // Commit the 1st attempt, fail the 2nd attempt, make sure 3rd attempt cannot commit,
     // then fail the 1st attempt and make sure the 4th one can commit again.
     stage += 1
     outputCommitCoordinator.stageStart(stage, maxPartitionId = 1)
-    assert(outputCommitCoordinator.canCommit(stage, 1, partition, taskAttempt))
-    outputCommitCoordinator.taskCompleted(stage, 2, partition, taskAttempt,
+    assert(outputCommitCoordinator.canCommit(stage, partition, 1L))
+    outputCommitCoordinator.taskCompleted(stage, partition, 2,
       ExecutorLostFailure("0", exitCausedByApp = true, None))
-    assert(!outputCommitCoordinator.canCommit(stage, 3, partition, taskAttempt))
-    outputCommitCoordinator.taskCompleted(stage, 1, partition, taskAttempt,
+    assert(!outputCommitCoordinator.canCommit(stage, partition, 3L))
+    outputCommitCoordinator.taskCompleted(stage, partition, 1L,
       ExecutorLostFailure("0", exitCausedByApp = true, None))
-    assert(outputCommitCoordinator.canCommit(stage, 4, partition, taskAttempt))
+    assert(outputCommitCoordinator.canCommit(stage, partition, 4L))
   }
 
   test("SPARK-24589: Make sure stage state is cleaned up") {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2.scala

@@ -109,11 +109,11 @@ object DataWritingSparkTask extends Logging {
       iter: Iterator[InternalRow],
       useCommitCoordinator: Boolean): WriterCommitMessage = {
     val stageId = context.stageId()
-    val stageAttempt = context.stageAttemptNumber()
     val partId = context.partitionId()
     val attemptId = context.attemptNumber()
     val epochId = Option(context.getLocalProperty(MicroBatchExecution.BATCH_ID_KEY)).getOrElse("0")
     val dataWriter = writeTask.createDataWriter(partId, attemptId, epochId.toLong)
+    val taskId = context.taskAttemptId()
 
     // write the data and commit this writer.
     Utils.tryWithSafeFinallyAndFailureCallbacks(block = {
@@ -123,33 +123,32 @@ object DataWritingSparkTask extends Logging {
 
       val msg = if (useCommitCoordinator) {
         val coordinator = SparkEnv.get.outputCommitCoordinator
-        val commitAuthorized = coordinator.canCommit(stageId, stageAttempt, partId, attemptId)
+        val commitAuthorized = coordinator.canCommit(stageId, partId, taskId)
         if (commitAuthorized) {
-          logInfo(s"Writer for stage $stageId / $stageAttempt, " +
-            s"task $partId.$attemptId is authorized to commit.")
+          logInfo(s"Writer for stage $stageId, part $partId, task $taskId is authorized to commit.")
           dataWriter.commit()
         } else {
-          val message = s"Stage $stageId / $stageAttempt, " +
-            s"task $partId.$attemptId: driver did not authorize commit"
+          val message = s"Stage $stageId, part $partId, task $taskId: " +
+            "driver did not authorize commit"
           logInfo(message)
           // throwing CommitDeniedException will trigger the catch block for abort
-          throw new CommitDeniedException(message, stageId, partId, attemptId)
+          throw new CommitDeniedException(message, stageId, partId, taskId)
         }
 
       } else {
         logInfo(s"Writer for partition ${context.partitionId()} is committing.")
         dataWriter.commit()
       }
 
-      logInfo(s"Writer for stage $stageId, task $partId.$attemptId committed.")
+      logInfo(s"Writer for stage $stageId, part $partId, task $taskId committed.")
 
       msg
 
     })(catchBlock = {
       // If there is an error, abort this writer
-      logError(s"Writer for stage $stageId, task $partId.$attemptId is aborting.")
+      logError(s"Writer for stage $stageId, part $partId, task $taskId is aborting.")
       dataWriter.abort()
-      logError(s"Writer for stage $stageId, task $partId.$attemptId aborted.")
+      logError(s"Writer for stage $stageId, part $partId, task $taskId aborted.")
     })
   }
 }