Spark源码 -- Stage的切分

源码挖掘机

已于 2023-11-01 13:22:52 修改

阅读量211

点赞数 1

分类专栏： spark 文章标签： spark 大数据分布式

于 2023-10-30 21:53:55 首次发布

本文链接：https://blog.csdn.net/bigdatakenan/article/details/134119479

版权

spark 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

上一节我们介绍了spark作业在集群中通过sparkSubmit脚本提交job，但job提交后作业是如何运行的，这节我们详细来探讨一下作业的运行和stage的切分，下面我们通过一个wordCount代码来看看作业是如何触发的。

import org.apache.spark.{SparkConf, SparkContext}

object WordCount {
  def main(args: Array[String]): Unit = {

    val sparkConf = new SparkConf()
      .setMaster("local[*]")
      .setAppName("wordCount")

    val spark = SparkContext.getOrCreate(sparkConf)

    val sourceRDD = spark.textFile("src/main/resources/word2.txt")

    val flatMapRDD = sourceRDD.flatMap(line => line.split(" "))

    val mapRDD = flatMapRDD.map((_, 1))

    val sortRDD = mapRDD.sortByKey()

    val reduceRDD = sortRDD.reduceByKey(_ + _)
    //行动算子会触发job的执行
    reduceRDD.count()


  }
}

我们都知道spark算子分为转换算子和行动算子，spark代码只有在遇到行动算子的时候才会触发代码的执行，这里的count算子就是行动算子。我们点击count算子进入算子内部

1.SparkContext

可以看到count算子底层调用sc.runJob方法，这里的sc就是sparkContext.然后点击runJob方法，底层调用了异构方法runjob

内部调用了好几个异构方法，最终通过调用dagScheduler.runJob方法

2.DAGScheduler

1.runJob方法

此方法中通过submitJob方法提交job,并且会异步获取job的提交结果，并执行对应操作

 def runJob[T, U](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      resultHandler: (Int, U) => Unit,
      properties: Properties): Unit = {
    val start = System.nanoTime
    //通过submitJob方法提交job
    val waiter = submitJob(rdd, func, partitions, callSite, resultHandler, properties)
    //异步获取job的提交结果，根据结果执行对应操作
    ThreadUtils.awaitReady(waiter.completionFuture, Duration.Inf)
    waiter.completionFuture.value.get match {
      case scala.util.Success(_) =>
        logInfo("Job %d finished: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
      case scala.util.Failure(exception) =>
        logInfo("Job %d failed: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
        // SPARK-8644: Include user stack trace in exceptions coming from DAGScheduler.
        val callerStackTrace = Thread.currentThread().getStackTrace.tail
        exception.setStackTrace(exception.getStackTrace ++ callerStackTrace)
        throw exception
    }
  }

2.submitJob方法

该方法中主要做了两件事：

1.创建jobWaiter，用于异步获取job的提交结果

2.往eventProcessLoop队列中发送一条jobSubmitted消息

def submitJob[T, U](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      resultHandler: (Int, U) => Unit,
      properties: Properties): JobWaiter[U] = {
    // Check to make sure we are not launching a task on a partition that does not exist.
    
    //获取该rdd最大分区数
    val maxPartitions = rdd.partitions.length
    partitions.find(p => p >= maxPartitions || p < 0).foreach { p =>
      throw new IllegalArgumentException(
        "Attempting to access a non-existent partition: " + p + ". " +
          "Total number of partitions: " + maxPartitions)
    }
    
    //获取JobJd
    val jobId = nextJobId.getAndIncrement()
    if (partitions.size == 0) {
      // Return immediately if the job is running 0 tasks
      return new JobWaiter[U](this, jobId, 0, resultHandler)
    }

    assert(partitions.size > 0)
    val func2 = func.asInstanceOf[(TaskContext, Iterator[_]) => _]
    //创建jobWaiter
    val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
    //往队列中发送一条jobSubmitted消息
    eventProcessLoop.post(JobSubmitted(
      jobId, rdd, func2, partitions.toArray, callSite, waiter,
      SerializationUtils.clone(properties)))
    waiter
  }

这里的eventProcessLoop是一个基于事件的消息队列，该队列的具体类型是DAGSchedulerEventProcessLoop，该消息队列启动后会去拉取收到的消息，并调用OnReceive方法去处理

我们点进DAGSchedulerEventProcessLoop这个类中，去找一下他的OnReceive方法

可以看到OnReceive方法中调用了doOnReceive方法，这个方法中会根据不同消息类型调用dagScheduler的不同方法，jobSubmitted类型的消息会调用它的handleJobSubmitted方法

private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
    case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
      dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)

    case MapStageSubmitted(jobId, dependency, callSite, listener, properties) =>
      dagScheduler.handleMapStageSubmitted(jobId, dependency, callSite, listener, properties)

    case StageCancelled(stageId, reason) =>
      dagScheduler.handleStageCancellation(stageId, reason)

    case JobCancelled(jobId, reason) =>
      dagScheduler.handleJobCancellation(jobId, reason)

    case JobGroupCancelled(groupId) =>
      dagScheduler.handleJobGroupCancelled(groupId)

    case AllJobsCancelled =>
      dagScheduler.doCancelAllJobs()

    case ExecutorAdded(execId, host) =>
      dagScheduler.handleExecutorAdded(execId, host)

    case ExecutorLost(execId, reason) =>
      val workerLost = reason match {
        case SlaveLost(_, true) => true
        case _ => false
      }
      dagScheduler.handleExecutorLost(execId, workerLost)

    case WorkerRemoved(workerId, host, message) =>
      dagScheduler.handleWorkerRemoved(workerId, host, message)

    case BeginEvent(task, taskInfo) =>
      dagScheduler.handleBeginEvent(task, taskInfo)

    case SpeculativeTaskSubmitted(task) =>
      dagScheduler.handleSpeculativeTaskSubmitted(task)

    case GettingResultEvent(taskInfo) =>
      dagScheduler.handleGetTaskResult(taskInfo)

    case completion: CompletionEvent =>
      dagScheduler.handleTaskCompletion(completion)

    case TaskSetFailed(taskSet, reason, exception) =>
      dagScheduler.handleTaskSetFailed(taskSet, reason, exception)

    case ResubmitFailedStages =>
      dagScheduler.resubmitFailedStages()
  }

3.handleJobSubmitted方法

一个Job有多个Stage，最后一个Stage叫做：ResultStage，前面的就叫做：ShuffleMapStage
ShuffleMapStage 执行完了之后，数据可以被持久化或者shuffle给下一个stage
resultStage 执行完了之后，就按照程序的要求，把数据持久化，除非打印输出

该方法中主要做了一下几件事：

1.创建resultStage，stage分为resultStage和shuffleMapStage

2.提交stage

 /**
     * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
     *   注释： RDD DAG划分Stages：Stage的划分是从最后一个Stage开始逆推的，每遇到一个宽依赖处，就分裂成另外一个Stage
     *   依此类推直到Stage划分完毕为止。并且，只有最后一个Stage的类型是ResultStage类型。
     *   注意Dataset、DataFrame、sparkSession.sql("select ...")经过catalyst代码解析会将代码转化为RDD
     *
     *   做了两件最重要的事情：
     *   1、stage切分
     *   2、stage提交
     */
    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 {
            /**
             * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
             *   注释： Stage 切分
             *   这个 finalRDD 就是 rdd链条中的最后一个 RDD，也就是触发 sc.runJob() 方法执行的 RDD
             *   必然是针对某个 RDD 调用了一个 action 算子才触发执行的，则该 RDD 就是 finallRDD
             *
             *   stage切分的核心方法：createResultStage
             *   返回的是一个 ResultStage对象，但是这个对象中，会包含他的所有的 父stage
             *   这样做的最大好处：容错！
             *   spark DAG引擎：血脉关系  RDD之间的依赖被构建成了 dependency ,也被构建成了 stage 之间的依赖关系
             */
            // New stage creation may throw an exception if, for example, jobs are run on a
            // HadoopRDD whose underlying HDFS files have been deleted.
            finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)
            
            /**
             * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
             *   注释：
             *   注意几对概念：
             *   1、ShuffleMapStage  +  ResultStage
             *   2、ShuffleMapTask + ResultTask
             *   3、ShuffleDependency + NarrowDependency
             */
            
            /**
             * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
             *   注释：先看这四个概念：
             *   1、Application
             *   2、Job
             *   3、Stage
             *      一个Job有多个Stage，最后一个Stage叫做：ResultStage，前面的就叫做：ShuffleMapStage
             *          ShuffleMapStage 执行完了之后，数据可以被持久化或者shuffle给下一个stage
             *          ResultStage 执行完了之后，就按照程序的要求，把数据持久化，除非打印输出
             *   4、Task
             */
            
        } catch {
            case e: BarrierJobSlotsNumberCheckFailed => logWarning(
                s"The job $jobId requires to run a barrier stage that requires more slots " + "than the total number of slots in the cluster currently.")
                // If jobId doesn't exist in the map, Scala coverts its value null to 0: Int automatically.
                val numCheckFailures = barrierJobIdToNumTasksCheckFailures.compute(jobId, new BiFunction[Int, Int, Int] {
                    override def apply(key: Int, value: Int): Int = value + 1
                })
                if (numCheckFailures <= maxFailureNumTasksCheck) {
                    messageScheduler.schedule(new Runnable {
                        override def run(): Unit = eventProcessLoop
                            .post(JobSubmitted(jobId, finalRDD, func, partitions, callSite, listener, properties))
                    }, timeIntervalNumTasksCheck, TimeUnit.SECONDS)
                    return
                } else {
                    // Job failed, clear internal data.
                    barrierJobIdToNumTasksCheckFailures.remove(jobId)
                    listener.jobFailed(e)
                    return
                }
            case e: Exception => logWarning("Creating new stage failed due to exception - job: " + jobId, e)
                listener.jobFailed(e)
                return
        }
        
        // Job submitted, clear internal data.
        barrierJobIdToNumTasksCheckFailures.remove(jobId)
        
        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))
        
        /**
         * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
         *   注释： 提交 Stage, 包含了所有的父stage，这里面也是通过递归来提交执行 stage
         */
        submitStage(finalStage)
    }

4.createResultStage方法

该方法中主要功能：

1.获取父stage，没有shuffle则会返回一个空list

2.创建resultStage并返回

 /**
     * TODO_MA Create a ResultStage associated with the provided jobId.
     *  进行Stage切分的详细方法实现
     */
    private def createResultStage(rdd: RDD[_], func: (TaskContext, Iterator[_]) => _, partitions: Array[Int], jobId: Int,
        callSite: CallSite): ResultStage = {
        checkBarrierStageWithDynamicAllocation(rdd)
        checkBarrierStageWithNumSlots(rdd)
        checkBarrierStageWithRDDChainPattern(rdd, partitions.toSet.size)
        
        // TODO_MA 注释：获得父stage,若没有shuffle则返回空List
        // TODO_MA 注释：获取当前Stage的parent Stage，这个方法是划分Stage的核心实现
        val parents = getOrCreateParentStages(rdd, jobId)
        
        // TODO_MA 注释： finalStage=resultStage 的 stageID
        val id = nextStageId.getAndIncrement()
        
        // TODO_MA 注释：创建当前最后的ResultStage
        // TODO_MA 注释：parents 所有的 父 stage
        val stage = new ResultStage(id, rdd, func, partitions, parents, jobId, callSite)
        
        // TODO_MA 注释：将 ResultStage 与 stageId 相关联, 保存在 map 中
        // TODO_MA 注释：stageIdToStage = new HashMap[Int, Stage]
        stageIdToStage(id) = stage
        
        // TODO_MA 注释：更新该job中包含的stage
        updateJobIdStageIdMaps(jobId, stage)
        
        // TODO_MA 注释：返回ResultStage
        stage
    }

5.getOrCreateParentStages方法

主要功能：

1.获取该rdd的shuffle依赖

2.获取或者创建shuffleMapStage

private def getOrCreateParentStages(rdd: RDD[_], firstJobId: Int): List[Stage] = {
        
        /**
         * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
         *   注释：
         *   1、遍历 RDD 的依赖关系，找到第一个 ShuffleDependency (可能多个，也可能没有)然后放入 HashSet 并返回
         *   2、如果获取不到 ShuffleDependency，逻辑在此终止返回空 list
         *   3、里面会创建当前 ShuffleDependency 的所有父 ShuffleMapStage
         */
        /**
         * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
         *   注释：
         *   1、rdd 就是  finalRDD
         *   2、getShuffleDependencies 获取当前 job 的所有 shuffleDependencies
         *   3、调用 map 函数，给每一个 ShuffleDepedency 构建 ShuffleMapStage
         *   4、构建 Stage 的方法是：getOrCreateShuffleMapStage
         *   5、map 的 参数函数的 参数 shuffleDep 是一个宽依赖
         */
        getShuffleDependencies(rdd).map { shuffleDep => getOrCreateShuffleMapStage(shuffleDep, firstJobId) }.toList
    }

6.getShuffleDepedendices方法

主要功能：

1.创建一个栈结构，将finalRDD压入栈中

2.取出栈中的第一个元素注册登记

3.获取该rdd的血缘关系，如果是宽依赖则放到parents集合中,如果是窄依赖则压入栈中

4.继续遍历栈中的元素，直至为空，使用深度优先遍历算法，返回parents

 * TODO_MA 注释：采用的是深度优先遍历找到 Action 算子的父依赖中的宽依赖
     * TODO_MA 注释：获得单个rdd的所有宽依赖关系(一般只有一个宽依赖；像CoGroupedRDD有多个)
     */
    private[scheduler] def getShuffleDependencies(rdd: RDD[_]): HashSet[ShuffleDependency[_, _, _]] = {
        
        // TODO_MA 注释：用来存放依赖
        val parents = new HashSet[ShuffleDependency[_, _, _]]
        
        // TODO_MA 注释：遍历过的RDD放入这个里面
        val visited = new HashSet[RDD[_]]
        
        // TODO_MA 注释：创建一个待遍历RDD的栈结构：先进后出的栈
        val waitingForVisit = new ArrayStack[RDD[_]]
        
        // TODO_MA 注释：压入 FinalRDD
        waitingForVisit.push(rdd)
        
        /**
         * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
         *   注释：循环遍历这个栈结构
         */
        while (waitingForVisit.nonEmpty) {
            // TODO_MA 注释： 取出栈顶的第一个RDD
            val toVisit = waitingForVisit.pop()
            
            // TODO_MA 注释：如果RDD没有被遍历过，则执行if内部的代码
            if (!visited(toVisit)) {
                // TODO_MA 注释：先登记
                visited += toVisit
                
                /**
                 * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
                 *   注释：
                 *   1、rdd.dependencies:通过 rdd 获得依赖关系
                 *   2、dependency.rdd:通过依赖关系获得 rdd
                 */
                // TODO_MA 注释：得到依赖，我们知道依赖中存放的有父RDD的对象
                toVisit.dependencies.foreach {
                    // TODO_MA 注释：如果这个依赖是shuffle依赖，则放入返回队列中
                    case shuffleDep: ShuffleDependency[_, _, _] => parents += shuffleDep
                    
                    // TODO_MA 注释：如果不是shuffle依赖，把其父RDD压入待访问栈中，从而进行循环
                    case dependency => waitingForVisit.push(dependency.rdd)
                }
            }
        }
        parents
    }

7.getOrCreateShuffleMapStage方法

主要功能：

1.根据shuffleDep的shuffleId获取stage,有的话就获取该stage

2.通过该宽依赖的血缘关系获取他的祖先shuffleDep

3.创建shuffleMapStage

/**
     * TODO_MA 如果 shuffleIdToMapStage 中存在 shuffle，则获取 shuffle map stage。
     *  否则，如果 shuffle map stage 不存在，该方法将创建 shuffle map stage
     *  以及任何丢失的 parent shuffle map stage。
     *
     *  第一次来到这里传进来的是左到右首个shuffleDep，没有父stage，
     *  shuffleIdToMapStage也没有记录，会调 getMissingAncestorShuffleDependencies 并创建 stage
     *  之后进来可直接根据 shuffleIdToMapStage 提取到 ShuffleMapStage
     *
     * Gets a shuffle map stage if one exists in shuffleIdToMapStage.
     * Otherwise, if the shuffle map stage doesn't already exist,
     * this method will create the shuffle map stage in addition to any missing ancestor shuffle map stages.
     */
    private def getOrCreateShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _], firstJobId: Int): ShuffleMapStage = {
        
        // TODO_MA 注释： 根据每个 ShuffleDep 的 shuffleID 来获取 Stage 对象
        shuffleIdToMapStage.get(shuffleDep.shuffleId) match {
            // TODO_MA 注释：如果有，则直接返回
            case Some(stage) => stage
            
            // TODO_MA 注释：如果还有 ShuffleDependency 没有构建 ShuffleMapStage 则创建一个
                // TODO_MA 注释：如果这个stage没有构建
                // TODO_MA 注释：第一件事：先把这个stage的所有父stage都构建出来
                // TODO_MA 注释：然后再次构建当前stage
            case None =>
                
                // TODO_MA 注释：给当前 stage 的所有 父stage 创建 ShuffleMapStage
                // TODO_MA 注释：举例：C <---B <-- A
                // Create stages for all missing ancestor shuffle dependencies.
                getMissingAncestorShuffleDependencies(shuffleDep.rdd)
                    .foreach { dep => // Even though getMissingAncestorShuffleDependencies only returns shuffle dependencies
                        // that were not already in shuffleIdToMapStage, it's possible that by the time we
                        // get to a particular dependency in the foreach loop, it's been added to
                        // shuffleIdToMapStage by the stage creation process for an earlier dependency. See
                        // SPARK-13902 for more information.
                        if (!shuffleIdToMapStage.contains(dep.shuffleId)) {
                            createShuffleMapStage(dep, firstJobId)
                        }
                    }
                
                /**
                 * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
                 *   注释：为指定的 ShuffleDependency 创建一个 ShuffleMapStage
                 */
                // Finally, create a stage for the given shuffle dependency.
                createShuffleMapStage(shuffleDep, firstJobId)
        }
    }

8.createShuffleMapStage方法

主要功能：

1.获取它的parentStage

2.创建shuffleMapStage并返回

  /**
     * TODO_MA 创建一个ShuffleMapStage，它生成给定的shuffle依赖项的分区。
     *  如果先前运行的stage生成了相同的shuffle 数据，则此函数将复制先前shuffle 中仍然可用的输出位置，以避免不必要地重新生成数据。
     *
     * Creates a ShuffleMapStage that generates the given shuffle dependency's partitions. If a
     * previously run stage generated the same shuffle data, this function will copy the output
     * locations that are still available from the previous shuffle to avoid unnecessarily regenerating data.
     */
    def createShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _], jobId: Int): ShuffleMapStage = {
        val rdd = shuffleDep.rdd
        checkBarrierStageWithDynamicAllocation(rdd)
        checkBarrierStageWithNumSlots(rdd)
        checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions)
        
        val numTasks = rdd.partitions.length
        val parents = getOrCreateParentStages(rdd, jobId)
        
        // TODO_MA 注释： 用来指定一个Stage的ID
        // TODO_MA 注释： 最先构建的stage 他的key=stageID就是最小的
        val id = nextStageId.getAndIncrement()
        
        // TODO_MA 注释：生成一个 ShuffleMapStage
        val stage = new ShuffleMapStage(id, rdd, numTasks, parents, jobId, rdd.creationSite, shuffleDep, mapOutputTracker)
        
        // TODO_MA 注释： 添加
        stageIdToStage(id) = stage
        shuffleIdToMapStage(shuffleDep.shuffleId) = stage
        
        updateJobIdStageIdMaps(jobId, stage)
        
        /**
         * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
         *   注释： 把 shuffle 信息注册到 Driver上的 MapOutputTrackerMaster 的 shuffleStatuses
         */
        if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {
            // Kind of ugly: need to register RDDs with the cache and map output tracker here
            // since we can't do it in the RDD constructor because # of partitions is unknown
            logInfo(s"Registering RDD ${rdd.id} (${rdd.getCreationSite}) as input to " + s"shuffle ${shuffleDep.shuffleId}")
            
            /**
             * TODO_MA 马中华 https://blog.csdn.net/zhongqi2513
             *   注释：
             *   1、把 Shuffle 信息注册到自己 Driver 的 MapOutputTrackerMaster 上，
             *   2、生成的是 shuffleId 和 ShuffleStatus 的映射关系
             *   3、在后面提交 Job 的时候还会根据它来的验证 map stage 是否已经准备好
             */
            mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)
        }
        stage
    }

至此，spark的stage切分源码就看完啦，有兴趣的伙伴可以去debug看一下源码的具体调用过程和顺序。