DAGScheduler任务调度剖析

一、主要功能

• 对每个job划分stage
• 决定任务运行的首选位置（preferred locations）
• 面向stage进行任务调度
• 跟踪stage的输出
• 需要时重新提交stage

二、主要功能的展开描述

1、怎么划分stage?

先看这个数据结构，stage划分之后保存在一个HashMap中，stageId作为key，Stage对象作为value

private[scheduler] val stageIdToStage = new HashMap[Int, Stage]

spark在任务提交时才划分stage，我们可以根据任务提交调用栈链条来找到线索跟踪job提交流程我们追踪到handleJobSubmitted

finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)

我们具体看下createResultStage

  private def createResultStage(
      rdd: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      jobId: Int,
      callSite: CallSite): ResultStage = {
    ......
    //这里创建祖先stage，这说明stage是根据最终RDD，逐级往源头回溯创建stage，可以看到主要的创建过程就在这里
    val parents = getOrCreateParentStages(rdd, jobId)
    val id = nextStageId.getAndIncrement()
    val stage = new ResultStage(id, rdd, func, partitions, parents, jobId, callSite)
    stageIdToStage(id) = stage
    updateJobIdStageIdMaps(jobId, stage)
    stage
  }

接下来

// 返回 List[Stage] ，这一步操作将祖先stage全部创建完成
private def getOrCreateParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
// 这里也就容易理解是按照宽依赖为界限划分stage，在为RDD的祖先RDD生成stage时，直接获取祖先的shuffle依赖，如果没有
shuffle依赖，那整个job就一个stage
    getShuffleDependencies(rdd).map { shuffleDep =>
    // 某个宽依赖的stage可能已经存在，避免重复生成
      getOrCreateShuffleMapStage(shuffleDep, firstJobId)
    }.toList
  }
  
private def getOrCreateShuffleMapStage(
      shuffleDep: ShuffleDependency[_, _, _],
      firstJobId: Int): ShuffleMapStage = {
    shuffleIdToMapStage.get(shuffleDep.shuffleId) match {
      case Some(stage) =>
        stage

      case None =>
        // 为祖先shuffle依赖创建stage
        getMissingAncestorShuffleDependencies(shuffleDep.rdd).foreach { dep =>
          if (!shuffleIdToMapStage.contains(dep.shuffleId)) {
          // 创建ShuffleMapStage 这一stage生成宽依赖的分区 
            createShuffleMapStage(dep, firstJobId)
          }
        }
        // 为传入参数中给定的宽依赖创建stage
        createShuffleMapStage(shuffleDep, firstJobId)
    }
  }

然后我们再看下createShuffleMapStage

def createShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _], jobId: Int): ShuffleMapStage = {
    val rdd = shuffleDep.rdd
    ......
    val numTasks = rdd.partitions.length
    // 关注这里，很明显开始重复创建stage过程了：最终RDD的生成stage--->最终RDD的祖先的宽依赖的生成stage--->
    再往谱系上游寻找其祖先的宽依赖，确定其宽依赖生成的stage--->直到没有宽依赖，祖先stage划分结束
    val parents = getOrCreateParentStages(rdd, jobId)
    val id = nextStageId.getAndIncrement()
    val stage = new ShuffleMapStage(
      id, rdd, numTasks, parents, jobId, rdd.creationSite, shuffleDep, mapOutputTracker)

    stageIdToStage(id) = stage
    shuffleIdToMapStage(shuffleDep.shuffleId) = stage
    updateJobIdStageIdMaps(jobId, stage)

    //这个地方对于后边说的首选位置有帮助，先提一下，重点在于mapOutputTracker
    if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {
      // Kind of ugly: need to register RDDs with the cache and map output tracker here
      // since we can't do it in the RDD constructor because # of partitions is unknown
      logInfo("Registering RDD " + rdd.id + " (" + rdd.getCreationSite + ")")
      mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)
    }
    stage
  }

2、怎么决定任务运行的首选位置？

我们还是从任务提交的流程中来确认这一逻辑，定位到submitMissingTasks方法中

// 找到stage需要进行计算的分区
val partitionsToCompute: Seq[Int] = stage.findMissingPartitions()
......
// 可以看到不管是什么类型的stage，都通过getPreferredLocs(stage.rdd, p) 来获取该stage的分区id对应的主机地址
val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {
      stage match {
      //可见分区id也作为任务id
        case s: ShuffleMapStage =>
          partitionsToCompute.map { id => (id, getPreferredLocs(stage.rdd, id))}.toMap
        case s: ResultStage =>
          partitionsToCompute.map { id =>
            val p = s.partitions(id)
            (id, getPreferredLocs(stage.rdd, p))
          }.toMap
      }

private[spark]
  def getPreferredLocs(rdd: RDD[_], partition: Int): Seq[TaskLocation] = {
    getPreferredLocsInternal(rdd, partition, new HashSet)
  }

private def getPreferredLocsInternal(
      rdd: RDD[_],
      partition: Int,
      visited: HashSet[(RDD[_], Int)]): Seq[TaskLocation] = {
    // 如果分区已经知道在哪里了，不需要再重复去找一次
    // This avoids exponential path exploration.  SPARK-695
    if (!visited.add((rdd, partition))) {
      // Nil has already been returned for previously visited partitions.
      return Nil
    }
    // 如果该分区已经缓存过，直接返回缓存位置
    // 调用次数和分区数一样，即和task数量一致
    // 函数中判断存储级别，如果存储级别为None，说明未缓存直接返回
    // 否则通过blockId从blockManagerMaster获取分区所在的主机和executorId
    val cached = getCacheLocs(rdd)(partition)
    if (cached.nonEmpty) {
      return cached
    }
    // If the RDD has some placement preferences (as is the case for input RDDs), get those
    // 根据输入源的不同，有不同的子类实现，用来获取数据的首选位置，代码先考虑checkpoint位置
    val rddPrefs = rdd.preferredLocations(rdd.partitions(partition)).toList
    if (rddPrefs.nonEmpty) {
      return rddPrefs.map(TaskLocation(_))
    }

    // 对于窄依赖，选择第一个具有首选位置的分区，将其位置当做任务的首选位置
    // 源码中也有说明，理想情况下应该根据需要传送的大小来比较来选出任务的首选位置，却并未这么做
    rdd.dependencies.foreach {
      case n: NarrowDependency[_] =>
        for (inPart <- n.getParents(partition)) {
          val locs = getPreferredLocsInternal(n.rdd, inPart, visited)
          if (locs != Nil) {
            return locs
          }
        }
      case _ =>
    }
    Nil
  }

3、面向stage调度？

面向stage调度是说DAGScheduler只在stage的层面进行调度，而不负责具体task的调度运行。

submitJob --->
JobSubmitted事件Event--->
DAGSchedulerEventProcessLoop事件处理器调用onReceive--->
doOnReceive--->dagScheduler.handleJobSubmitted--->
submitStage(finalStage)--->submitMissingTasks(stage, jobId.get)--->
若任务数大于0执行调用
taskScheduler.submitTasks(new TaskSet( tasks.toArray, stage.id, stage.latestInfo.attemptNumber, jobId, properties))--->
若任务数为0，调用markStageAsFinished(stage, None) ，然后提交子stage：submitWaitingChildStages(stage)

可以看到，在DAGScheduler中，会将stage中的任务作为任务集TaskSet提交给TaskScheduler，具体TaskScheduler的实现负责任务集中每个task的调度。

4、DAGScheduler如何跟踪stage的输出？

DAGScheduler需要追踪stage的输出，确切的讲应该是ShuffleMapStage输出，以便进行最小开销调度

private[scheduler] def markMapStageJobsAsFinished(shuffleStage: ShuffleMapStage): Unit = {
    // Mark any map-stage jobs waiting on this stage as finished
    if (shuffleStage.isAvailable && shuffleStage.mapStageJobs.nonEmpty) {
      // 前边提到过mapOutputTracker，这是一个MapOutputTrackerMaster对象，这一操作获取了给定shuffle输出的所有统计值及
      mapStatus状态信息
      // 在创建ShuffleMapStage的时候，已经将shuffleID，shuffle分区等信息注册给了mapOutputTracker
      val stats = mapOutputTracker.getStatistics(shuffleStage.shuffleDep)
      for (job <- shuffleStage.mapStageJobs) {
        markMapStageJobAsFinished(job, stats)
      }
    }
  }

4、需要时重新提交stage

以下两种情况下DAGScheduler均要重新提交stage

private[scheduler] def handleTaskCompletion(event: CompletionEvent){
......
// 因shuffle output分区lost导致的拉取失败，需要重新提交上一个stage重新计算shuffle输出
case FetchFailed(bmAddress, shuffleId, mapId, _, failureMessage) =>
......
// 因shuffle任务失败，并且stage重试尝试次数未超过最大可连续重试值
case failure: TaskFailedReason if task.isBarrier =>
......
messageScheduler.schedule(new Runnable {
   override def run(): Unit = eventProcessLoop.post(ResubmitFailedStages)
}, DAGScheduler.RESUBMIT_TIMEOUT, TimeUnit.MILLISECONDS)
......
}