Spark 源码解析 : DAGScheduler中的DAG划分与提交

Spark 源码解析 : DAGScheduler中的DAG划分与提交

一、Spark 运行架构

Spark 运行架构如下图：

各个RDD之间存在着依赖关系，这些依赖关系形成有向无环图DAG，DAGScheduler对这些依赖关系形成的DAG，进行Stage划分，划分的规则很简单，从后往前回溯，遇到窄依赖加入本stage，遇见宽依赖进行Stage切分。完成了Stage的划分,DAGScheduler基于每个Stage生成TaskSet,并将TaskSet提交给TaskScheduler。TaskScheduler 负责具体的task调度,在Worker节点上启动task。

二、源码解析：DAGScheduler中的DAG划分

    当RDD触发一个Action操作（如：colllect）后，导致SparkContext.runJob的执行。而在SparkContext的run方法中会调用DAGScheduler的run方法最终调用了DAGScheduler的submit方法：

1. def submitJob[T, U](
2. rdd: RDD[T],
3. func: (TaskContext, Iterator[T]) => U,
4. partitions: Seq[Int],
5. callSite: CallSite,
6. resultHandler: (Int, U) => Unit,
7. properties: Properties): JobWaiter[U] = {
8. // Check to make sure we are not launching a task on a partition that does not exist.
9. val maxPartitions = rdd.partitions.length
10. partitions.find(p => p >= maxPartitions || p < 0).foreach { p =>
11. throw new IllegalArgumentException(
12. "Attempting to access a non-existent partition: " + p + ". " +
13. "Total number of partitions: " + maxPartitions)
14. }
16. val jobId = nextJobId.getAndIncrement()
17. if (partitions.size == 0) {
18. // Return immediately if the job is running 0 tasks
19. return new JobWaiter[U](this, jobId, 0, resultHandler)
20. }
22. assert(partitions.size > 0)
23. val func2 = func.asInstanceOf[(TaskContext, Iterator[_]) => _]
24. val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
25. //给eventProcessLoop发送JobSubmitted消息
26. eventProcessLoop.post(JobSubmitted(
27. jobId, rdd, func2, partitions.toArray, callSite, waiter,
28. SerializationUtils.clone(properties)))
29. waiter
30. }

DAGScheduler的submit方法中，像eventProcessLoop对象发送了JobSubmitted消息。 eventProcessLoop是 DAGSchedulerEventProcessLoop 类的对象

1. private[scheduler] val eventProcessLoop = new DAGSchedulerEventProcessLoop(this)

DAGSchedulerEventProcessLoop，接收各种消息并进行处理，处理的逻辑在其doOnReceive方法中：

1. private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
2.    //Job提交
   

1. case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
2. dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)
4. case MapStageSubmitted(jobId, dependency, callSite, listener, properties) =>
5. dagScheduler.handleMapStageSubmitted(jobId, dependency, callSite, listener, properties)
7. case StageCancelled(stageId) =>
8. dagScheduler.handleStageCancellation(stageId)
10. case JobCancelled(jobId) =>
11. dagScheduler.handleJobCancellation(jobId)
13. case JobGroupCancelled(groupId) =>
14. dagScheduler.handleJobGroupCancelled(groupId)
16. case AllJobsCancelled =>
17. dagScheduler.doCancelAllJobs()
19. case ExecutorAdded(execId, host) =>
20. dagScheduler.handleExecutorAdded(execId, host)
22. case ExecutorLost(execId) =>
23. dagScheduler.handleExecutorLost(execId, fetchFailed = false)
25. case BeginEvent(task, taskInfo) =>
26. dagScheduler.handleBeginEvent(task, taskInfo)
28. case GettingResultEvent(taskInfo) =>
29. dagScheduler.handleGetTaskResult(taskInfo)
31. case completion: CompletionEvent =>
32. dagScheduler.handleTaskCompletion(completion)
34. case TaskSetFailed(taskSet, reason, exception) =>
35. dagScheduler.handleTaskSetFailed(taskSet, reason, exception)
37. case ResubmitFailedStages =>
38. dagScheduler.resubmitFailedStages()
39. }

可以把 DAGSchedulerEventProcessLoop 理解成DAGScheduler的对外的功能接口。它对外隐藏了自己内部实现的细节。无论是内部还是外部消息， DAGScheduler可以共用同一消息处理代码，逻辑清晰，处理方式统一。

接下来分析 DAGScheduler的Stage划分， handleJobSubmitted 方法首先创建ResultStage

1. try {
2. //创建新stage可能出现异常，比如job运行依赖hdfs文文件被删除
3. finalStage = newResultStage(finalRDD, func, partitions, jobId, callSite)
4. } catch {
5. case e: Exception =>
6. logWarning("Creating new stage failed due to exception - job: " + jobId, e)
7. listener.jobFailed(e)
8. return
9. }

然后调用submitStage方法，进行stage的划分。

首先由finalRDD获取它的父RDD依赖，判断依赖类型，如果是窄依赖，则将父RDD压入栈中，如果是宽依赖，则作为父Stage。

看一下源码的具体过程:

1. private def getMissingParentStages(stage: Stage): List[Stage] = {
2. val missing = new HashSet[Stage] //存储需要返回的父Stage
3. val visited = new HashSet[RDD[_]] //存储访问过的RDD
4. //自己建立栈，以免函数的递归调用导致
5. val waitingForVisit = new Stack[RDD[_]]
6. 

7. def visit(rdd: RDD[_]) {
8. if (!visited(rdd)) {
9. visited += rdd
10. val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
11. if (rddHasUncachedPartitions) {
12. for (dep <- rdd.dependencies) {
13. dep match {
14. case shufDep: ShuffleDependency[_, _, _] =>
15. val mapStage = getShuffleMapStage(shufDep, stage.firstJobId)
16. if (!mapStage.isAvailable) {
17. missing += mapStage //遇到宽依赖，加入父stage
18. }
19. case narrowDep: NarrowDependency[_] =>
20. waitingForVisit.push(narrowDep.rdd) //窄依赖入栈,
21. }
22. }
23. }
24. }
25. }
26. 

27.    //回溯的起始RDD入栈
28. waitingForVisit.push(stage.rdd)
29. while (waitingForVisit.nonEmpty) {
30. visit(waitingForVisit.pop())
31. }
32. missing.toList
33. }

getMissingParentStages方法是由当前stage，返回他的父stage，父stage的创建由getShuffleMapStage返回，最终会调用 newOrUsedShuffleStage 方法返回ShuffleMapStage

1. private def newOrUsedShuffleStage(
2. shuffleDep: ShuffleDependency[_, _, _],
3. firstJobId: Int): ShuffleMapStage = {
4. val rdd = shuffleDep.rdd
5. val numTasks = rdd.partitions.length
6. val stage = newShuffleMapStage(rdd, numTasks, shuffleDep, firstJobId, rdd.creationSite)
7. if (mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {
8. //Stage已经被计算过，从MapOutputTracker中获取计算结果
9. val serLocs = mapOutputTracker.getSerializedMapOutputStatuses(shuffleDep.shuffleId)
10. val locs = MapOutputTracker.deserializeMapStatuses(serLocs)
11. (0 until locs.length).foreach { i =>
12. if (locs(i) ne null) {
13. // locs(i) will be null if missing
14. stage.addOutputLoc(i, locs(i))
15. }
16. }
17. } else {
18. // Kind of ugly: need to register RDDs with the cache and map output tracker here
19. // since we can't do it in the RDD constructor because # of partitions is unknown
20. logInfo("Registering RDD " + rdd.id + " (" + rdd.getCreationSite + ")")
21. mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)
22. }
23. stage
24. }

现在父Stage已经划分好，下面看看你Stage的提交逻辑

1. /** Submits stage, but first recursively submits any missing parents. */
2. private def submitStage(stage: Stage) {
3. val jobId = activeJobForStage(stage)
4. if (jobId.isDefined) {
5. logDebug("submitStage(" + stage + ")")
6. if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
7. val missing = getMissingParentStages(stage).sortBy(_.id)
8. logDebug("missing: " + missing)
9. if (missing.isEmpty) {
10. logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
11. //如果没有父stage，则提交当前stage
12. submitMissingTasks(stage, jobId.get)
13. } else {
14. for (parent <- missing) {
15. //如果有父stage，则递归提交父stage
16. submitStage(parent)
17. }
18. waitingStages += stage
19. }
20. }
21. } else {
22. abortStage(stage, "No active job for stage " + stage.id, None)
23. }
24. }

提交的过程很简单，首先当前stage获取父stage，如果父stage为空，则当前Stage为起始stage，交给submitMissingTasks处理，如果当前stage不为空，则递归调用submitStage进行提交。

到这里，DAGScheduler中的DAG划分与提交就讲完了，下次解析这些stage是如果封装成TaskSet交给TaskScheduler以及TaskSchedule的调度过程。