(spark) branch branch-4.3 updated: [SPARK-58913][CORE] Add `TaskSchedulerImpl.hasPipelinedTaskSets`

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon-hyun pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.3 by this push:
     new e87de7971388 [SPARK-58913][CORE] Add `TaskSchedulerImpl.hasPipelinedTaskSets`
e87de7971388 is described below

commit e87de79713887b801ed675221a50d3366adc0655
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Thu Aug 20 11:29:37 2026 -0700

    [SPARK-58913][CORE] Add `TaskSchedulerImpl.hasPipelinedTaskSets`
    
    This PR adds `TaskSchedulerImpl.hasPipelinedTaskSets`, a query for whether any live task set belongs to a pipelined group:
    
    ```scala
    private[spark] def hasPipelinedTaskSets: Boolean = synchronized {
      taskSetsByStageIdAndAttempt.values.exists(_.values.exists { tsm =>
        !tsm.isZombie && tsm.taskSet.isPipelined
      })
    }
    ```
    
    It sits next to `outstandingTasksForOtherWorkInProfile`, the other pipelined-group helper over the same map, and follows its conventions: taken under the existing lock, and skipping zombie (superseded) attempts.
    
    Pipeline Dependency is added at Apache Spark 4.3.0.
    - https://github.com/apache/spark/pull/57286
    - https://github.com/apache/spark/pull/57361
    
    A pipelined group's shuffle data is transient and lives only on the executors running the group, so a caller about to disturb those executors -- graceful decommission in #58054 -- needs to know that such a group is running. No existing API answers that: `TaskSet.isPipelined` is per task set, and `taskSetsByStageIdAndAttempt` is private to `TaskSchedulerImpl`.
    
    The check is best-effort by design. It reports what the *task* scheduler currently holds, which is narrower than "a pipelined job is active": it is false before the group's first task set is submitted, and false again once the last member's `TaskSetManager` has gone zombie but the `DAGScheduler` has not yet processed the final completion event.
    
    No. This is a new `private[spark]` API.
    
    Pass the CIs.
    
    Generated-by: Claude Opus 5
    
    Closes #58171 from dongjoon-hyun/SPARK-58913.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
    (cherry picked from commit 020da8cc090e40f5700481d4b82cff955c57fd63)
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../apache/spark/scheduler/TaskSchedulerImpl.scala | 23 +++++++++
 .../spark/scheduler/TaskSchedulerImplSuite.scala   | 55 ++++++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index 887545b5da1b..4600993d2664 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -193,6 +193,29 @@ private[spark] class TaskSchedulerImpl(
     }.sum
   }
 
+  /**
+   * Whether any live task set belongs to a pipelined group (see `TaskSet.isPipelined`). Such a
+   * group's shuffle data is transient and lives only on its executors, so a caller about to
+   * disturb the executors needs to know that a group is running.
+   *
+   * Zombie (superseded) attempts are skipped: their tasks are no longer scheduled, so a stale
+   * pipelined attempt left in the map must not make the scheduler look like it still has a
+   * pipelined group in flight. Matches the !isZombie filtering used elsewhere on this map.
+   *
+   * Best-effort: this reports what the TASK scheduler currently holds, which is narrower than "a
+   * pipelined job is active". It is false before the group's first task set is submitted, and
+   * false again once the last member's TaskSetManager has gone zombie but the DAGScheduler has
+   * not yet processed the final completion event.
+   *
+   * Synchronized like every other access to this map, so it is safe to call from any thread --
+   * including off the DAGScheduler event loop (e.g. SparkContext, on a user thread).
+   */
+  private[spark] def hasPipelinedTaskSets: Boolean = synchronized {
+    taskSetsByStageIdAndAttempt.values.exists(_.values.exists { tsm =>
+      !tsm.isZombie && tsm.taskSet.isPipelined
+    })
+  }
+
   // The set of executors we have on each host; this is used to compute hostsAlive, which
   // in turn is used to decide when we can attain data locality on a given host
   protected val hostToExecutors = new HashMap[String, HashSet[String]]
diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
index 81782df283ee..03d78b4d70b7 100644
--- a/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
+++ b/core/src/test/scala/org/apache/spark/scheduler/TaskSchedulerImplSuite.scala
@@ -2866,4 +2866,59 @@ class TaskSchedulerImplSuite extends SparkFunSuite with LocalSparkContext
     assert(taskScheduler.outstandingTasksForOtherWorkInProfile(defaultRp, Set(0)) === 0)
   }
 
+  /**
+   * A TaskSet marked as a pipelined-group member. FakeTask.createTaskSet always leaves
+   * `isPipelined` at its default (false), so build the TaskSet directly here.
+   */
+  private def pipelinedTaskSet(numTasks: Int, stageId: Int, stageAttemptId: Int = 0): TaskSet = {
+    val tasks = Array.tabulate[Task[_]](numTasks)(i => new FakeTask(stageId, i, Nil))
+    new TaskSet(tasks, stageId, stageAttemptId, priority = 0, null,
+      ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, None, isPipelined = true)
+  }
+
+  test("SPARK-58913: hasPipelinedTaskSets is true only while a pipelined-group member is live") {
+    val taskScheduler = setupScheduler()
+    assert(!taskScheduler.hasPipelinedTaskSets, "no task set has been submitted yet")
+
+    // A regular task set leaves isPipelined at false and must not register as a group member.
+    taskScheduler.submitTasks(FakeTask.createTaskSet(2, stageId = 0, stageAttemptId = 0))
+    assert(!taskScheduler.hasPipelinedTaskSets,
+      "a regular task set is not a pipelined-group member")
+
+    // One pipelined member anywhere in the map is enough, even next to regular task sets.
+    taskScheduler.submitTasks(pipelinedTaskSet(2, stageId = 1))
+    assert(taskScheduler.hasPipelinedTaskSets,
+      "a live pipelined task set must be reported even alongside regular ones")
+  }
+
+  test("SPARK-58913: hasPipelinedTaskSets ignores zombie attempts") {
+    // A zombie (superseded by a retry/kill) attempt no longer schedules tasks, so a stale pipelined
+    // attempt left in the map must not make the scheduler look like a group is still in flight.
+    val taskScheduler = setupScheduler()
+    taskScheduler.submitTasks(pipelinedTaskSet(2, stageId = 0, stageAttemptId = 0))
+    assert(taskScheduler.hasPipelinedTaskSets)
+
+    taskScheduler.taskSetManagerForAttempt(0, 0).get.isZombie = true
+    assert(!taskScheduler.hasPipelinedTaskSets,
+      "a zombie pipelined attempt must not count as a live group member")
+
+    // The live retry of the same stage is a group member again.
+    taskScheduler.submitTasks(pipelinedTaskSet(2, stageId = 0, stageAttemptId = 1))
+    assert(taskScheduler.hasPipelinedTaskSets,
+      "the live retry attempt must be reported even though the zombie attempt is skipped")
+  }
+
+  test("SPARK-58913: hasPipelinedTaskSets is false once every member's task set has finished") {
+    // The task scheduler drops a task set once it finishes, so the flag goes false at that point
+    // even though the job that owns the group is not done yet (the DAGScheduler has still to
+    // process the final completion). This is the documented narrowing of the task-scheduler view.
+    val taskScheduler = setupScheduler()
+    taskScheduler.submitTasks(pipelinedTaskSet(1, stageId = 0, stageAttemptId = 0))
+    assert(taskScheduler.hasPipelinedTaskSets)
+
+    val tsm = taskScheduler.taskSetManagerForAttempt(0, 0).get
+    taskScheduler.taskSetFinished(tsm)
+    assert(!taskScheduler.hasPipelinedTaskSets,
+      "a finished task set is no longer held by the task scheduler")
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org