Spark源码：Task最佳位置计算方法

在之前的文章中《SparkCore — stage划分算法源码分析​​​​​​​》，创建完Stages之后，就开始提交Stages，在DAGScheduler.scala的submitStage方法中，使用submitMissingTasks，提交第一个Stage0，并且剩余的Stage加入等待队列，waitingStages，剩余的Stage使用submitWaitingStages()方法提交Stage。

1. 下面首先来分析一下submitMissingTasks()方法，

传入两个参数，Stage和jobId

                private def submitMissingTasks(stage: Stage, jobId: Int)

接着计算需要创建Task的数量，事实上它和Partition的数量是相等的，具体代码如下：

// 获取待创建Task的数量
val partitionsToCompute: Seq[Int] = {
      // 如果Stage包含Shuffle依赖
      if (stage.isShuffleMap) {
        // task的数量与RDD的partitions的数量是一样的
        (0 until stage.numPartitions).filter(id => stage.outputLocs(id) == Nil)
      } else {
        // 直接创建ResultTask
        val job = stage.resultOfJob.get
        (0 until job.numPartitions).filter(id => !job.finished(id))
      }
 }

获取Job的优先级，Job优先级是按照stage创建的先后顺序定义的，Stage0具有最高的优先级，并将当前Stage加入运行队列中。

// 获取job的优先级，说白了按照stage创建先后定义优先级，stage0具有最高优先级
val properties = if (jobIdToActiveJob.contains(jobId)) {
  jobIdToActiveJob(stage.jobId).properties
} else {
  // this stage will be assigned to "default" pool
  null
}
// 将stage加入runningStages队列
runningStages += stage

下面省略一些不是特别重要的代码。。。

// 将需要运行的Job加入监控中，然后提交stage
stage.latestInfo = StageInfo.fromStage(stage, Some(partitionsToCompute.size))
    outputCommitCoordinator.stageStart(stage.id)
listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))

// 将task序列化，然后创建其共享变量，有ShuffleMapTask与ResultTask之分
// 在worker节点的executor上再反序列化task
var taskBinary: Broadcast[Array[Byte]] = null
try {
  // For ShuffleMapTask, serialize and broadcast (rdd, shuffleDep).
  val taskBinaryBytes: Array[Byte] =
    if (stage.isShuffleMap) {
      closureSerializer.serialize((stage.rdd, stage.shuffleDep.get) : AnyRef).array()
    } else {
      closureSerializer.serialize((stage.rdd, stage.resultOfJob.get.func) : AnyRef).array()
    }
  taskBinary = sc.broadcast(taskBinaryBytes)
} catch {
    // 此处省略代码
    .....
}

上面就是将task进行序列化，创建其共享变量，并且将该Stage的RDD也一并序列化，发送给Worker节点上的executor，这样做的好处是，每个Task都有一个不同的RDD副本，具有强隔离作用，防止闭包引用被更改。

下面的代码就是为Stage创建指定数量的Task，并且针对每个Task计算它的最佳位置：

// 为stage创建指定数量的task，这里很关键的一点是，task的最佳位置计算算法
val tasks: Seq[Task[_]] = try {
  // 如果是ShuffleMap
  if (stage.isShuffleMap) {
    // 给每个partition创建一个task
    // 给每个task计算最佳位置
    partitionsToCompute.map { id =>
      // 获取Stage的RDD所在的本地节点位置信息
      val locs = getPreferredLocs(stage.rdd, id)
      val part = stage.rdd.partitions(id)
      // 对于finalStage之外的stage，它的isShuffleMap都是true
      // 所以创建ShuffleMapTask
      new ShuffleMapTask(stage.id, taskBinary, part, locs)
    }
  } else {
    // 如果不是shuffleMap，那么就是finalStage
    // finalstage它创建的是ResultTask
    val job = stage.resultOfJob.get
    partitionsToCompute.map { id =>
      // 也是计算RDD的本地位置信息
      val p: Int = job.partitions(id)
      val part = stage.rdd.partitions(p)
      val locs = getPreferredLocs(stage.rdd, p)
      new ResultTask(stage.id, taskBinary, part, locs, id)
    }
  }
} catch {
    // 此处省略代码
    .......
}

在上述代码中，为每个Stage创建指定数量的Task，其中分为ShuffleMapTask和ResultTask，他们都会计算当前Stage包含的RDD的最佳位置，最佳位置是通过getPreferredLocs()方法计算，下面分析一下这个方法。

// 该方法嵌套了另一个方法，getPreferredLocsInternal
def getPreferredLocs(rdd: RDD[_], partition: Int): Seq[TaskLocation] = {
    getPreferredLocsInternal(rdd, partition, new HashSet)
}

// 这个方法里面包含了计算Task最佳位置的方法，这是个线程安全的方法
 private def getPreferredLocsInternal(
      rdd: RDD[_],
      partition: Int,
      visited: HashSet[(RDD[_],Int)])
    : Seq[TaskLocation] =
  {
    // 先判断partition是否已经被访问，如果访问过，那么无需再找该RDD了
    if (!visited.add((rdd,partition))) {
      // Nil has already been returned for previously visited partitions.
      return Nil
    }

    // 寻找当前RDD的partition是否被缓存了，如果被缓存，则返回缓存的位置
    val cached = getCacheLocs(rdd)(partition)
    if (!cached.isEmpty) {
      return cached
    }
    
    // 寻找当前RDD的partition是否checkpoint了
    val rddPrefs = rdd.preferredLocations(rdd.partitions(partition)).toList
    if (!rddPrefs.isEmpty) {
      return rddPrefs.map(TaskLocation(_))
    }
   
    // 假如既没有缓存也没有checkpoint，
    // 就递归调用自己，看看该RDD的父RDD，看看对应的partition是否缓存或者checkpoint
    rdd.dependencies.foreach {
      case n: NarrowDependency[_] =>
        for (inPart <- n.getParents(partition)) {
          val locs = getPreferredLocsInternal(n.rdd, inPart, visited)
          if (locs != Nil) {
            return locs
          }
        }
      case _ =>
    }
    // 如果这个stage，从最后一个RDD，到最开始的RDD，partition都没有被缓存或者checkpoint
    // 那么task的最佳位置就是Nil。
    Nil
  }

Task的最佳路径，说白了就是从stage的最后一个RDD开始，去找该RDD的partition是被cache了，或者checkpoint了，那么，task的最佳位置，就是缓存cache的或者checkpoint的partition的位置，因为这样的话，task就在那个节点上运行，不需要计算之前的RDD了，否则的话RDD没有缓存或checkpoint，那么就需要从头重新计算，这样效率就会很低。

因此，从Task最佳位置计算算法来看，对于使用频率比较高的中间结果的RDD，最好被cache或者checkpoint，这样能避免从头重复计算，提高程序效率，这也是性能调优的地方。