SparkCore — stage划分算法源码分析

DAGscheduler中的stage划分源码分析

  在之前的文章中，已经分析了stage的划分算法，这里我们到源码里面去看划分算法是怎么实现的。
  首先找到提交Job的入口（从action操作开始，找到action操作的runJob -> dagScheduler.runJob -> submitJob -> eventProcessLoop.JobSubmitted -> handleJobSubmitted）handleJobSubmitted()这个方法，下面我们具体分析源码

 private[scheduler] def handleJobSubmitted(jobId: Int,
      finalRDD: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      callSite: CallSite,
      listener: JobListener,
      properties: Properties) {
    // 使用触发Job的最后一个RDD，创建finalStage，它的类型是ResultStage
    var finalStage: ResultStage = null
    try {
      // 创建finalStage，并且将stage加入DAGScheduler内部的缓存中
      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
    }

    // 用finalStage创建一个Job，里面封装了Job的一些信息(比如partition的数量，ResultStage和ShuffleMapStage是不一样的，这里在性能调优的时候再讲)
    val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)
    // 清除RDD缓存
    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()
    // 将Job加入内存缓存中
    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被提交，并且其他stage，都放入了waitingStages里了。
    submitStage(finalStage)

    // 提交完第一个stage0后，剩余的stage，通过这个函数提交
    submitWaitingStages()
  }

  上面的代码中，首先使用触发Job的最后一个RDD，创建一个finalStage，这个stage是ResultStage（一个Job里面只有最后一个stage是ResultStage，其余的都是ShuffleMapStage），然后创建Job（里面包含了Job的一些信息），并将Job的相关信息放入缓存中，接着就创建了比较重要的submitStage(finalStage)方法，这个方法里面就包含了stage的划分和提交；而submitWaitingStages()则是提交剩余的stage，下面我们分析一下submitStage(finalStage)方法。

private def submitStage(stage: Stage) {
    // 获取JobId
    val jobId = activeJobForStage(stage)
    // job存在，开始进行任务分配
    if (jobId.isDefined) {
      logDebug("submitStage(" + stage + ")")
      // 这个stage还没有被提交过，或者正在运行，失败等
      if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
        // 使用getMissingParentStages()方法，获取这个stage的父stage
        val missing = getMissingParentStages(stage).sortBy(_.id)
        logDebug("missing: " + missing)
        // 这里会反复的递归调用，直到最初的stage，它没有父stage，此时就会提交第一个stage，
        // 也就是stage0
        if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          // 如果父stage没有父stage，那么就提交stage，这里的递归，就是stage划分算法的推动者。
          submitMissingTasks(stage, jobId.get)
        } else {
          for (parent <- missing) {
            // 递归调用submitStage()方法，去提交父Stage
            submitStage(parent)
          }
          // 并且将当前stage，放入waitingStages等待队列中
          waitingStages += stage
        }
      }
    } else {
      abortStage(stage, "No active job for stage " + stage.id, None)
    }
  }

  下面简单说是如何进行stage的划分的，看注释，首先就是根据当前这个stage找到它的父stage，假如父stage的RDD与当前stage的RDD是宽依赖的关系，那么就用这个宽依赖的RDD创建一个ShuffleMapStage，并返回；假如不存在宽依赖，那么就一直遍历下去，直到第一个RDD为止。我们看一下getMissingParentStages()方法。

 private def getMissingParentStages(stage: Stage): List[Stage] = {
	// 创建存储ShuffleMapStage的集合，以及已经遍历RDD的集合
    val missing = new HashSet[Stage]
    val visited = new HashSet[RDD[_]]
    // We are manually maintaining a stack here to prevent StackOverflowError
    // caused by recursively visiting
    // 创建待遍历RDD的栈，防止递归调用栈溢出
    val waitingForVisit = new Stack[RDD[_]]
    def visit(rdd: RDD[_]) {
      if (!visited(rdd)) {
        visited += rdd
        val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
        if (rddHasUncachedPartitions) {
          // 遍历RDD的依赖
          for (dep <- rdd.dependencies) {
            dep match {
              // 如果是宽依赖，使用宽依赖的RDD，创建一个stage.
              // 默认最后一个stage是ResultStage，之前所有的stage都是ShuffleMapStage
              case shufDep: ShuffleDependency[_, _, _] =>
                val mapStage = getShuffleMapStage(shufDep, stage.firstJobId)
                if (!mapStage.isAvailable) {
                  missing += mapStage
                }
              // 如果是窄依赖，那么将RDD放入栈中
              case narrowDep: NarrowDependency[_] =>
                waitingForVisit.push(narrowDep.rdd)
            }
          }
        }
      }
    }
    // 首先往栈中推入stage最后的一个RDD
    waitingForVisit.push(stage.rdd)
    // 进行while循环
    while (waitingForVisit.nonEmpty) {
      // 对stage的最后一个RDD，调用内部函数visit()
      // 对窄依赖会一直调用，直到出现宽依赖，那么就创建一个stage，并返回
      visit(waitingForVisit.pop())
    }
    missing.toList
  }

  可以看到这个方法也是一个递归调用，为了防止栈溢出，使用了一个stack结构waitingForVisit，我们看内部函数visit()方法，它会遍历当前RDD的依赖，假如存在shuffle依赖，那么就创建一个ShuffleMapStage，并返回，否则就一直执行下去，直到没有父RDD为止，这也就意味着整个Job只创建了一个Stage（ResultStage）。
  下面接着看submitStage的源码，通过getMissingParentStages()获取当前Stage的父Stage，假如存在，就执行到下面这块代码，这块代码就是stage划分算法的推动者：

if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          // 如果父stage没有父stage，那么就提交stage，这里的递归，就是stage划分算法的推动者。
          submitMissingTasks(stage, jobId.get)
        } else {
          for (parent <- missing) {
            // 递归调用submitStage()方法，去提交父Stage
            submitStage(parent)
          }
          // 并且将当前stage，放入waitingStages等待队列中
          waitingStages += stage
        }

  如果存在父Stage，那么遍历父Stage，并递归调用submitStage();我们以下面这幅经典的图来说明：

  如上面图所示，首先以最后一个RDDG创建finalStage（Stage3）,接着通过**getMissingParentStages()**找它的父RDD，它有两个父RDD，分别是RDDF和RDDB。先看RDDB，它与RDDG是窄依赖，接着遍历RDDB的父RDD，也即RDDA，它们之间是groupByKey操作（发生了shuffle），是宽依赖，因此创建一个stage1，这个遍历结束；接着看RDDF，它与RDDG之间的操作是join，也发生了shuffle操作，因此创建了stage2；在遍历完两个父RDD之后，就返回getMissingParentStages()函数。这时候missing列表里面包含了stage1和stage2。
  下面判断missing是否为空，由于包含了stage1和stage2，因此不为空，就接着遍历。先遍历stage1，由于stage1没有依赖，因此它的missing是空，那么这里就调用submitMissingTasks()提交stage1；接着遍历到stage2，stage2的RDDF的父RDD中没有宽依赖，因此它的missing列表也为空，提交stage2。
  接着执行到waitingStages，它将finalStage，也就stage3加入这个等待队列中。到这里整个submitStage()就运行完成，只有最后一个stage被加入了等待队列中。
  submitStage()执行完成之后，接着执行submitWaitingStages()，提交加入等待队列的stage。
  整个stage的划分以及提交就结束了。以上就是stage的划分以及提交过程，主要分析了一下流程，至于细节这块，以后再慢慢研究，更新博客。
  总结一下，stage的划分是以shuffle为界，也即宽依赖，如果RDD之间发生了shuffle，那么就会以shuffle为界创建新的stage，依次内推。而stage的提交是递归提交，最先创建的stage，会最后提交，这刚好符合RDD的处理流程的先后顺序。