remove the jobId from the map when one job failed.

dmlc · Aug 26, 2020 · f1df786 · f1df786
1 parent 8852f59
commit f1df786
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/SparkParallelismTracker.scala b/jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/SparkParallelismTracker.scala
@@ -135,6 +135,7 @@ class TaskFailedListener(killSparkContext: Boolean = true) extends SparkListener
 
             if (stageIds.contains(stageId)) {
               logger.error("Cancelling jobId:" + jobId)
+              jobIdToStageIds.remove(jobId)
               SparkContext.getOrCreate().cancelJob(jobId)
             }
           })

diff --git a/...ckages/xgboost4j-spark/src/test/scala/org/apache/spark/SparkParallelismTrackerSuite.scala b/...ckages/xgboost4j-spark/src/test/scala/org/apache/spark/SparkParallelismTrackerSuite.scala
@@ -106,7 +106,7 @@ class SparkParallelismTrackerSuite extends FunSuite with PerTest {
     }
   }
 
-  test("tracker should cancel correct job when killSparkContext=false") {
+  test("tracker should cancel the correct job when killSparkContext=false") {
     val nWorkers = 2
     val tracker = new SparkParallelismTracker(sc, 0, nWorkers, false)
     val rdd: RDD[Int] = sc.parallelize(1 to 10, nWorkers)