Spark源码解读之Stage划分和提交

上一篇讲解了Spark源码解读之Job提交，这一篇主要讲解Stage划分和提交。

调用流程：

org.apache.spark.scheduler.DAGScheduler.handleJobSubmitted

org.apache.spark.scheduler.DAGScheduler.submitStage

org.apache.spark.scheduler.DAGScheduler.submitMissingTasks

org.apache.spark.scheduler.TaskScheduler.submitTasks

一、Stage划分

Spark中会根据RDD之间的依赖关系进行Stage划分，在遇到ShuffleDependency时，会将这两个RDD划分到不同的Stage。在调用DAGScheduler的handleJobSubmitted进行Job提交后，会先进行Stage划分，源码如下：

// 参数finalRDD为触发action操作时最后一个RDD
private[scheduler] def handleJobSubmitted(jobId: Int,
finalRDD: RDD[_],
func: (TaskContext, Iterator[_]) => _,
partitions: Array[Int],
callSite: CallSite,
listener: JobListener,
properties: Properties) {
var finalStage: ResultStage = null
try {
// New stage creation may throw an exception if, for example, jobs are run on a
// HadoopRDD whose underlying HDFS files have been deleted.
// 创建finalStage
finalStage = newResultStage(finalRDD, func, partitions, jobId, callSite)
} catch {
case e: Exception =>
logWarning("Creating new stage failed due to exception - job: " + jobId, e)
listener.jobFailed(e)
return
}

val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)
clearCacheLocs()
logInfo("Got job %s (%s) with %d output partitions".format(
job.jobId, callSite.shortForm, partitions.length))
logInfo("Final stage: " + finalStage + " (" + finalStage.name + ")")
logInfo("Parents of final stage: " + finalStage.parents)
logInfo("Missing parents: " + getMissingParentStages(finalStage))

val jobSubmissionTime = clock.getTimeMillis()
jobIdToActiveJob(jobId) = job
activeJobs += job
finalStage.setActiveJob(job)
val stageIds = jobIdToStageIds(jobId).toArray
val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
listenerBus.post(
SparkListenerJobStart(job.jobId, jobSubmissionTime, stageInfos, properties))
// 提交finalStage，该方法会提交所有关联的未提交的stage
submitStage(finalStage)

submitWaitingStages()
}

可以看出，在创建finalStage时初始化了newResultStage实例。最后调用submitStage方法（详情见Stage提交部分）。newResultStage源码如下：

/**
* Create a ResultStage associated with the provided jobId.
*/
private def newResultStage(
rdd: RDD[_],
func: (TaskContext, Iterator[_]) => _,
partitions: Array[Int],
jobId: Int,
callSite: CallSite): ResultStage = {
// 获取parent Stages和id
val (parentStages: List[Stage], id: Int) = getParentStagesAndId(rdd, jobId)
// 创建stage
val stage = new ResultStage(id, rdd, func, partitions, parentStages, jobId, callSite)
// 更新stageId和Stage、JobId和stageId之间的映射关系
stageIdToStage(id) = stage
updateJobIdStageIdMaps(jobId, stage)
// 返回stage
stage
}

初始化newResultStage实例时会做两件事，一是调用getParentStagesAndId方法得到parentStages和id，二是更新stageId和Stage、JobId和stageId之间的映射关系。下面是getParentStagesAndId源码：

/**
* Helper function to eliminate some code re-use when creating new stages.
*/
private def getParentStagesAndId(rdd: RDD[_], firstJobId: Int): (List[Stage], Int) = {
// 获取parentStages
val parentStages = getParentStages(rdd, firstJobId)
// 获取一个唯一id
val id = nextStageId.getAndIncrement()
(parentStages, id)
}

getParentStagesAndId会做两件事，一是调用getParentStages得到parentStages列表，二是获取一个Stage的唯一id。getParentStages源码如下：

/**
* Get or create the list of parent stages for a given RDD.  The new Stages will be created with
* the provided firstJobId.
*/
private def getParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
val parents = new HashSet[Stage]
val visited = new HashSet[RDD[_]]
// We are manually maintaining a stack here to prevent StackOverflowError
// caused by recursively visiting
val waitingForVisit = new Stack[RDD[_]]
// 广度优先遍历方式
def visit(r: RDD[_]) {
if (!visited(r)) { // visited为空
visited += r
// Kind of ugly: need to register RDDs with the cache here since
// we can't do it in its constructor because # of partitions is unknown
for (dep <- r.dependencies) {
dep match {
// 依赖为ShuffleDependency类型时，则生成一个新的shuffle map Stage
case shufDep: ShuffleDependency[_, _, _] =>
parents += getShuffleMapStage(shufDep, firstJobId)
// 依赖为非ShuffleDependency类型时，则加入到waitingForVisit栈中
case _ =>
waitingForVisit.push(dep.rdd)
}
}
}
}
waitingForVisit.push(rdd)

while (waitingForVisit.nonEmpty) {
// 调用visit方法
visit(waitingForVisit.pop())
}
parents.toList
}

getParentStages会先创建一个类型为RDD的栈waitingForVisit，然后遍历waitingForVisit，如果该RDD的依赖为ShuffleDependency类型，则调用getShuffleMapStage方法得到一个shuffle map stage，否则将该RDD的父RDD加入到waitingForVisit中。getShuffleMapStage源码如下：

/**
* Get or create a shuffle map stage for the given shuffle dependency's map side.
*/
private def getShuffleMapStage(
shuffleDep: ShuffleDependency[_, _, _],
firstJobId: Int): ShuffleMapStage = {
shuffleToMapStage.get(shuffleDep.shuffleId) match {
case Some(stage) => stage
case None =>
// We are going to register ancestor shuffle dependencies
getAncestorShuffleDependencies(shuffleDep.rdd).foreach { dep =>
shuffleToMapStage(dep.shuffleId) = newOrUsedShuffleStage(dep, firstJobId)
}
// Then register current shuffleDep
val stage = newOrUsedShuffleStage(shuffleDep, firstJobId)
shuffleToMapStage(shuffleDep.shuffleId) = stage
stage
}
}

getShuffleMapStage会获取或者创建一个shuffle map stage。

二、Stage提交

Stage划分完成后，会进行Stage提交，Stage提交首先会调用submitStage方法，源码如下：

/** Submits stage, but first recursively submits any missing parents. */
private def submitStage(stage: Stage) {
val jobId = activeJobForStage(stage)
if (jobId.isDefined) {
logDebug("submitStage(" + stage + ")")
if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
// 获取未提交的父Stage
val missing = getMissingParentStages(stage).sortBy(_.id)
logDebug("missing: " + missing)
if (missing.isEmpty) { // 所有的父Stage都已提交
logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
submitMissingTasks(stage, jobId.get) // 提交该Stage
} else {// 父Stage会提交
for (parent <- missing) {
submitStage(parent) /// 提交父Stage
}
waitingStages += stage
}
}
} else {
abortStage(stage, "No active job for stage " + stage.id, None)
}
}

submitStage会检测该Stage的父Stage是否提交，如果有父Stage未提交，则会递归调用submitStage；如果父Stage都已提交，则会调用submitMissingTasks方法提交该Stage。submitMissingTasks源码如下：

/** Called when stage's parents are available and we can now do its task. */
private def submitMissingTasks(stage: Stage, jobId: Int) {
logDebug("submitMissingTasks(" + stage + ")")
// Get our pending tasks and remember them in our pendingTasks entry
stage.pendingPartitions.clear()

// First figure out the indexes of partition ids to compute.
// 得到需要计算的partitions
val partitionsToCompute: Seq[Int] = stage.findMissingPartitions()

// Create internal accumulators if the stage has no accumulators initialized.
// Reset internal accumulators only if this stage is not partially submitted
// Otherwise, we may override existing accumulator values from some tasks
if (stage.internalAccumulators.isEmpty || stage.numPartitions == partitionsToCompute.size) {
stage.resetInternalAccumulators()
}

// Use the scheduling pool, job group, description, etc. from an ActiveJob associated
// with this Stage
val properties = jobIdToActiveJob(jobId).properties

runningStages += stage
// SparkListenerStageSubmitted should be posted before testing whether tasks are
// serializable. If tasks are not serializable, a SparkListenerStageCompleted event
// will be posted, which should always come after a corresponding SparkListenerStageSubmitted
// event.

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)
}

// 创建一个Map：taskIdToLocations，存储的是id->Seq[TaskLocation]的映射关系，这里的id表示task所包含的RDD的partition id，TaskLocation表示任务位置
// 实现时，对stage中需要计算的RDD的分区调用PreferredLocations来获取优先位置信息，映射成id->Seq[TaskLocation]的关系
val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {
stage match {
case s: ShuffleMapStage =>
partitionsToCompute.map { id => (id, getPreferredLocs(stage.rdd, id))}.toMap
case s: ResultStage =>
val job = s.activeJob.get
partitionsToCompute.map { id =>
val p = s.partitions(id)
(id, getPreferredLocs(stage.rdd, p))
}.toMap
}
} catch {
case NonFatal(e) =>
stage.makeNewStageAttempt(partitionsToCompute.size)
listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))
abortStage(stage, s"Task creation failed: $e\n${e.getStackTraceString}", Some(e))
runningStages -= stage
return
}

stage.makeNewStageAttempt(partitionsToCompute.size, taskIdToLocations.values.toSeq)
listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))

// TODO: Maybe we can keep the taskBinary in Stage to avoid serializing it multiple times.
// Broadcasted binary for the task, used to dispatch tasks to executors. Note that we broadcast
// the serialized copy of the RDD and for each task we will deserialize it, which means each
// task gets a different copy of the RDD. This provides stronger isolation between tasks that
// might modify state of objects referenced in their closures. This is necessary in Hadoop
// where the JobConf/Configuration object is not thread-safe.
var taskBinary: Broadcast[Array[Byte]] = null
try {
// For ShuffleMapTask, serialize and broadcast (rdd, shuffleDep).
// 对于ShuffleMapTask，序列化并广播，广播的是rdd和shuffleDep
// For ResultTask, serialize and broadcast (rdd, func).
// 对于ResultTask，序列化并广播，广播的是rdd和func
val taskBinaryBytes: Array[Byte] = stage match {
case stage: ShuffleMapStage =>
closureSerializer.serialize((stage.rdd, stage.shuffleDep): AnyRef).array()
case stage: ResultStage =>
closureSerializer.serialize((stage.rdd, stage.func): AnyRef).array()
}

taskBinary = sc.broadcast(taskBinaryBytes)
} catch {
// In the case of a failure during serialization, abort the stage.
case e: NotSerializableException =>
abortStage(stage, "Task not serializable: " + e.toString, Some(e))
runningStages -= stage

// Abort execution
return
case NonFatal(e) =>
abortStage(stage, s"Task serialization failed: $e\n${e.getStackTraceString}", Some(e))
runningStages -= stage
return
}

// 针对stage的每个分区构造task，形成tasks:ShuffleMapStage生成ShuffleMapTasks，ResultStage生成ResultTasks
val tasks: Seq[Task[_]] = try {
stage match {
case stage: ShuffleMapStage =>
partitionsToCompute.map { id =>
val locs = taskIdToLocations(id)
val part = stage.rdd.partitions(id)
new ShuffleMapTask(stage.id, stage.latestInfo.attemptId,
taskBinary, part, locs, stage.internalAccumulators)
}

case stage: ResultStage =>
val job = stage.activeJob.get
partitionsToCompute.map { id =>
val p: Int = stage.partitions(id)
val part = stage.rdd.partitions(p)
val locs = taskIdToLocations(id)
new ResultTask(stage.id, stage.latestInfo.attemptId,
taskBinary, part, locs, id, stage.internalAccumulators)
}
}
} catch {
case NonFatal(e) =>
abortStage(stage, s"Task creation failed: $e\n${e.getStackTraceString}", Some(e))
runningStages -= stage
return
}

// 如果存在tasks，则利用taskScheduler.submitTasks()提交task，否则标记stage已完成
if (tasks.size > 0) {
logInfo("Submitting " + tasks.size + " missing tasks from " + stage + " (" + stage.rdd + ")")
stage.pendingPartitions ++= tasks.map(_.partitionId)
logDebug("New pending partitions: " + stage.pendingPartitions)
// 调用taskScheduler.submitTasks提交task
taskScheduler.submitTasks(new TaskSet(
tasks.toArray, stage.id, stage.latestInfo.attemptId, jobId, properties))
// 记录提交时间
stage.latestInfo.submissionTime = Some(clock.getTimeMillis())
} else {
// Because we posted SparkListenerStageSubmitted earlier, we should mark
// the stage as completed here in case there are no tasks to run
markStageAsFinished(stage, None)

val debugString = stage match {
case stage: ShuffleMapStage =>
s"Stage ${stage} is actually done; " +
s"(available: ${stage.isAvailable}," +
s"available outputs: ${stage.numAvailableOutputs}," +
s"partitions: ${stage.numPartitions})"
case stage : ResultStage =>
s"Stage ${stage} is actually done; (partitions: ${stage.numPartitions})"
}
logDebug(debugString)
}
}

submitMissingTasks会做以下几个事：

1. 清空stage的pendingPartitions

2. 得到需要计算的partition id索引，放入partitionsToCompute

3. 将stage加入到runningStages中 

4. 启动一个stage

5. 得到task中执行的位置，即计算stage的每个RDD的partition的优先位置，存入taskIdToLocations

6. 对stage进行序列化并广播

7. （重要）针对stage的每个RDD的partition构造task，存入tasks

8. 存在tasks，则调用taskScheduler.submitTasks()提交task，否则标记stage已完成。