Comments are addressed.

Author: erenavsarogullari

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

@@ -22,7 +22,7 @@ import org.apache.spark.TaskContext
 class FakeTask(
     stageId: Int,
     partitionId: Int,
-    prefLocs: Seq[TaskLocation] = Nil) extends Task[Int](stageId, 0, partitionId) {
+    prefLocs: Seq[TaskLocation] = Nil) extends Task[Int](stageId, stageAttemptId = 0, partitionId) {
   override def runTask(context: TaskContext): Int = 0
   override def preferredLocations: Seq[TaskLocation] = prefLocs
 }
@@ -33,11 +33,11 @@ object FakeTask {
    * locations for each task (given as varargs) if this sequence is not empty.
    */
   def createTaskSet(numTasks: Int, prefLocs: Seq[TaskLocation]*): TaskSet = {
-    createTaskSet(numTasks, 0, prefLocs: _*)
+    createTaskSet(numTasks, stageAttemptId = 0, prefLocs: _*)
   }
 
   def createTaskSet(numTasks: Int, stageAttemptId: Int, prefLocs: Seq[TaskLocation]*): TaskSet = {
-    createTaskSet(numTasks, 0, stageAttemptId, prefLocs: _*)
+    createTaskSet(numTasks, stageId = 0, stageAttemptId, prefLocs: _*)
   }
 
   def createTaskSet(numTasks: Int, stageId: Int, stageAttemptId: Int, prefLocs: Seq[TaskLocation]*):
@@ -48,6 +48,6 @@ object FakeTask {
     val tasks = Array.tabulate[Task[_]](numTasks) { i =>
       new FakeTask(stageId, i, if (prefLocs.size != 0) prefLocs(i) else Nil)
     }
-    new TaskSet(tasks, stageId, stageAttemptId, 0, null)
+    new TaskSet(tasks, stageId, stageAttemptId, priority = 0, null)
   }
 }

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

@@ -974,18 +974,20 @@ class TaskSetManagerSuite extends SparkFunSuite with LocalSparkContext with Logg
     assert(manager.isZombie)
   }
 
-  test("check uniqueness of TaskSetManager name") {
+  test("SPARK-17894: Verify TaskSetManagers for different stage attempts have unique names") {
     sc = new SparkContext("local", "test")
     sched = new FakeTaskScheduler(sc, ("exec1", "host1"))
-    val taskSet = FakeTask.createTaskSet(1, 0, 0)
+    val taskSet = FakeTask.createTaskSet(numTasks = 1, stageId = 0, stageAttemptId = 0)
     val manager = new TaskSetManager(sched, taskSet, MAX_TASK_FAILURES, new ManualClock)
     assert(manager.name === "TaskSet_0.0")
 
-    val taskSet2 = FakeTask.createTaskSet(1, 0, 1)
+    // Make sure a task set with the same stage ID but different attempt ID also has a unique name
+    val taskSet2 = FakeTask.createTaskSet(numTasks = 1, stageId = 0, stageAttemptId = 1)
     val manager2 = new TaskSetManager(sched, taskSet2, MAX_TASK_FAILURES, new ManualClock)
     assert(manager2.name === "TaskSet_0.1")
 
-    val taskSet3 = FakeTask.createTaskSet(1, 1, 1)
+    // Make sure a task set with the same attempt ID but different stage ID also has a unique name
+    val taskSet3 = FakeTask.createTaskSet(numTasks = 1, stageId = 1, stageAttemptId = 1)
     val manager3 = new TaskSetManager(sched, taskSet3, MAX_TASK_FAILURES, new ManualClock)
     assert(manager3.name === "TaskSet_1.1")
   }