Spark Task执行流程源码分析系列之二: 任务构建&调度&资源分配

上一节我们介绍了Task各个环节用到的主要数据结构，本节我们来看看Spark中一个Task是如何构建起来的，又是如何获取到资源，然后提交给集群相应的资源进行启动的。

任务构建&提交

Spark job内部是通过DAG来维护血缘关系的，通过shuffle算子进行stage的划分，上游stage计算完成后，下游stage才能进行，在一个stage中有多个任务需要执行，划分完stage后就会对同一个stage的任务集合进行提交，然后分配资源执行任务，我们先来看下任务提交入口，步骤如下:

1. 首先清空需要计算的stage待处理分区的索引的集合，找出当前stage还没有计算的分区<一个分区是一个Task>；
2. 将当前stage加入到runningStages集合中，并启动对当前stage输出提交到HDFS的协调机制；
3. 计算每个需要计算分区对应任务的偏好分区位置，以方便调度时候找到最合适的位置信息；
4. 对任务进行序列化并广播，ShuffleMapTask会对Stage的rdd和ShuffleDependency进行序列化，ResultTask则是对Stage的rdd和对RDD的分区进行计算的函数func进行序列化；
5. 构建Task集合TaskSet，根据stage的类型创建ShuffleMapTask或者ResultTask集合；
6. 如果集合长度大于0，说明当前stage还有没有未执行的任务，交由TaskScheduler进行调度执行；如果集合长度为0，表明这个stage已经完成了，可以触发下游stage进行执行尝试(由于下一个stage可能依赖多个上游stage，所以也不一定会直接执行)。

// org.apache.spark.scheduler.DAGScheduler 
private def submitMissingTasks(stage: Stage, jobId: Int) {
     
  // 清空当前Stage的pendingPartitions，便于记录需要计算的分区任务。
  stage.pendingPartitions.clear()
  // 找出当前Stage的所有分区中还没有完成计算的分区的索引
  val partitionsToCompute: Seq[Int] = stage.findMissingPartitions()
  // 获取ActiveJob的properties。properties包含了当前Job的调度、group、描述等属性信息。
  val properties = jobIdToActiveJob(jobId).properties

  // 将stage添加到runningStages集合中，表示其正在运行
  runningStages += stage
  // 启动对当前Stage的输出提交到HDFS的协调机制
  stage match {
   
    case s: ShuffleMapStage =>
    outputCommitCoordinator.stageStart(stage = s.id, maxPartitionId = s.numPartitions - 1)
    case s: ResultStage =>
    outputCommitCoordinator.stageStart(stage = s.id, maxPartitionId = s.rdd.partitions.length - 1)
  }
  // 获取还没有完成计算的每一个分区的偏好位置
  val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {
   
    stage match {
   
      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
    }
  } catch {
   
    // 如果发生任何异常，则调用Stage的makeNewStageAttempt()方法开始一次新的Stage执行尝试
    case NonFatal(e) =>
    ...
    return
  }

  // 开始Stage的执行尝试,对这次stage进行分装分配attemptId
  stage.makeNewStageAttempt(partitionsToCompute.size, taskIdToLocations.values.toSeq)
  // 向事件总线投递SparkListenerStageSubmitted事件
  listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))
  
  // 对任务进行序列化并广播
  var taskBinary: Broadcast[Array[Byte]] = null
  try {
   
    val taskBinaryBytes: Array[Byte] = stage match {
   
      // 对Stage的rdd和ShuffleDependency进行序列化
      case stage: ShuffleMapStage =>
      JavaUtils.bufferToArray(closureSerializer.serialize((stage.rdd, stage.shuffleDep): AnyRef))
      // 对Stage的rdd和对RDD的分区进行计算的函数func进行序列化
      case stage: ResultStage =>
      JavaUtils.bufferToArray(closureSerializer.serialize((stage.rdd, stage.func): AnyRef))
    }

    // 广播任务的序列化对象
    taskBinary = sc.broadcast(taskBinaryBytes)
  } catch {
   
    case e: NotSerializableException =>
    ...
    return
    case NonFatal(e) =>
    ...
    return
  }

  // 创建Task序列
  val tasks: Seq[Task[_]] = try {
   
    stage match {
   
      case stage: ShuffleMapStage => // 为ShuffleMapStage的每一个分区创建一个ShuffleMapTask
      partitionsToCompute.map {
    id 
        val locs = taskIdToLocations(id)  // 对应分区的偏好位置序列
        val part = stage.rdd.partitions(id)   // RDD的分区
        // 创建ShuffleMapTask
        new ShuffleMapTask(stage.id, stage.latestInfo.attemptId,
                           taskBinary, part, locs, stage.latestInfo.taskMetrics, properties, Option(jobId),
                           Option(sc.applicationId), sc.applicationAttemptId)
      }

      case stage: ResultStage => // 为ResultStage的每一个分区创建一个ResultTask
      partitionsToCompute.map {
    id =>
        val p: Int = stage.partitions(id)
        val part = stage.rdd.partitions(p)  // RDD的分区
        val locs = taskIdToLocations(id)  // 分区偏好位置序列
        // 创建ResultTask
        new ResultTask(stage.id, stage.latestInfo.attemptId,
                       taskBinary, part, locs, id, properties, stage.latestInfo.taskMetrics,
                       Option(jobId), Option(sc.applicationId), sc.applicationAttemptId)
      }
    }
  } catch {
   
    case NonFatal(e) =>
    ...
    return
  }

  if (tasks.size > 0) {
    // Task数量大于0
    // 将提交的分区添加到pendingPartitions集合中，表示它们正在等待处理
    stage.pendingPartitions ++= tasks.map(_.partitionId)
    // 为这批Task创建TaskSet，调用TaskScheduler的submitTasks方法提交此批Task
    taskScheduler.submitTasks(new TaskSet(
      tasks.toArray, stage.id, stage.latestInfo.attemptId, jobId, properties))
    // 记录最后一次提交时间
    stage.latestInfo.submissionTime = Some(clock.getTimeMillis())
  } else {
    // Task数量为0，没有创建任何Task
    // 将当前Stage标记为完成
    markStageAsFinished(stage, None)
    // 提交当前Stage的子Stage
    submitWaitingChildStages(stage)
  }
}

DAGScheduler向TaskScheduler提交了TaskSet之后，TaskSchedulerImpl会为每个TaskSet创建一个TaskSetManager对象，该对象包含TaskSet所有 tasks，并管理这些tasks的调度，执行以及失败重试等，TaskSetManager新建后，会加入到调度池中，进行调度执行，最后会通过scheduleBackend进行资源的申请来运行这些job。

// org.apache.spark.scheduler.TaskSchedulerImpl
override def submitTasks(taskSet: TaskSet) {
   
  val tasks = taskSet.tasks  // 获取TaskSet中的所有Task
  this.synchronized {
   
    val manager = createTaskSetManager(taskSet, maxTaskFailures)  // 创建TaskSetManager
    val stage = taskSet.stageId  // TaskSet的Stage
    // 更新taskSetsByStageIdAndAttempt中记录的推测执行信息
    val stageTaskSets = taskSetsByStageIdAndAttempt.getOrElseUpdate(stage, new HashMap[Int, TaskSetManager])
    stageTaskSets(taskSet.stageAttemptId) = manager
    // 判断是否有冲突的TaskSet，taskSetsByStageIdAndAttempt中不应该存在同属于当前Stage，但是TaskSet却不相同的情况
    val conflictingTaskSet = stageTaskSets.exists {
    case (_, ts) =>
      ts.taskSet != taskSet && !ts.isZombie
    }
    if (conflictingTaskSet) {
   
      throw new IllegalStateException(s"more than one active taskSet for stage $stage:" +
                                      s" ${stageTaskSets.toSeq.map{_._2.taskSet.id}.mkString(",