Spark Scheduler 模块（上）

在阅读 Spark 源代码的过程中，发现单步调试并不能很好的帮助理解程序。这样的多线程的分布式系统，更好的阅读源代码的方式是依据模块，分别理解。

 

在包 org.apache.spark 下面有很多下一级的包，如 deploy, storage, shuffle, scheduler 等。这就是一个个系统模块，本文主要介绍 scheduler 模块。

 

博客http://jerryshao.me/architecture/2013/04/21/Spark%E6%BA%90%E7%A0%81%E5%88%86%E6%9E%90%E4%B9%8B-scheduler%E6%A8%A1%E5%9D%97/ 讲述的是 Spark-0.7 版本的 Scheduler 模块，本文很大程度上参考了该博客，不过对应的代码是 Spark-1.5 版本的。

 

另外，在阅读本文前，最好先读懂 Spark 那篇经典的论文，这四个术语：RDD，窄依赖(Narrow Dependency)，宽依赖(Shuffle Dependency)，stage。

注：本文在黏贴源代码的过程中，只选择了与主干逻辑相关的部分。

 

======================== 正文开始 ==========================

 

现代的分布式计算系统，以经典的 YARN 为例，实现了资源的统一管理平台，Spark 启动时会向 YARN 申请一定的计算资源（CPU和Memory），然后自己管理这些底层的计算资源。同时，Spark 还负责上层用户程序的任务调度，即把用户程序分解成一个个小任务，运行在这些资源上面。本文主要分析 Spark 的上层任务调度。（更接近用户，容易理解）

 

跟踪 RDD.count()，可以看到真正的执行语句是 SparkContext.runJob，往下跟踪会发现最终调用的方法是 dagScheduler.runJob。再然后 DAGScheduler 把 job 变成 stage，再把 stage 转换成 task，最后调用 taskScheduler.submitTasks 把任务提交给了执行单元。这里最重要的两个类就是 DAGScheduler 和 TaskScheduler。

 

DAGScheduler

在 SparkContext 启动的时候，也创建并启动了_taskScheduler 和 _dagScheduler。 

// Create and start the scheduler
val (sched, ts) = SparkContext.createTaskScheduler(this, master)
_schedulerBackend = sched
_taskScheduler = ts
_dagScheduler = new DAGScheduler(this)

SparkContext.createTaskScheduler 方法中，依据 master 的不同，分别创建不同的 _taskScheduler，比如 master=local 对应的是类 TaskSchedulerImpl，master=yarn-cluster对应的是org.apache.spark.scheduler.cluster.YarnClusterScheduler。因为 TaskScheduler 需要运行在资源调度器（YARN, Mesos, Local）上，所以需要分别实现。而 DAGScheduler 是 Spark 自身的逻辑，只有一种实现即可。这种设计上的分离也是值得我们学习的。

 

需要注意到 DAGScheduler 中一个重要的属性:

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

这是一个使用 BlockingQueue实现的消息队列。DAGSchedulerEventProcessLoop 中的方法 onReceive 能够异步的接受来自用户的 DAGSchedulerEvent 类型的事件，分别处理。事件类型包括：

JobSubmitted
CompletionEvent
StageCancelled
JobCancelled

Job 到 Stage 的完整流程

当 SparkContext.runJob 调用了 DAGScheduler.runJob 后，传入的参数是 RDD。DAGScheduler.runJob 什么都没做，又把 RDD 传入了 DAGScheduler.submitJob：

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] = {
    val jobId = nextJobId.getAndIncrement()
    val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
    eventProcessLoop.post(JobSubmitted(
      jobId, rdd, func2, partitions.toArray, callSite, waiter,
      SerializationUtils.clone(properties)))
    waiter
  }

submitJob() 创建了一个独立的 jobId，一个传出去的 waiter，然后把发送了一个事件 JobSubmitted 给 eventProcessLoop: DAGSchedulerEventProcessLoop。上文提到这是一个异步的消息队列。于是函数调用者（DAGScheduler.runJob）就使用 waiter 等待消息即可。

当  eventProcessLoop: DAGSchedulerEventProcessLoop 接受到事件 JobSubmitted，然后调用 DAGScheduler.handleJobSubmitted

private[scheduler] def handleJobSubmitted(jobId: Int,
    finalRDD: RDD[_],
    func: (TaskContext, Iterator[_]) => _,
    partitions: Array[Int],
    callSite: CallSite,
    listener: JobListener,
    properties: Properties) {
  var finalStage: ResultStage = null
  finalStage = newResultStage(finalRDD, partitions.length, jobId, callSite)
  if (finalStage != null) {
    val job = new ActiveJob(jobId, finalStage, func, partitions, callSite, listener, properties)
    activeJobs += job
    finalStage.resultOfJob = Some(job)
    submitStage(finalStage)
  }
  submitWaitingStages()
}

首先创建了 finalStage，然后 sumbitStage 把它提交。submitWaitingStages 是提交以前失败的、现在又满足条件的作业。这里的核心是 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)) {
      val missing = getMissingParentStages(stage).sortBy(_.id)
      logDebug("missing: " + missing)
      if (missing.isEmpty) {
        logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
        submitMissingTasks(stage, jobId.get)
      } else {
        for (parent <- missing) {
          submitStage(parent)
        }
        waitingStages += stage
      }
    }
  } else {
    abortStage(stage, "No active job for stage " + stage.id, None)
  }
}

private def getMissingParentStages(stage: Stage): List[Stage] = {
  val missing = new HashSet[Stage]
  val visited = new HashSet[RDD[_]]
  val waitingForVisit = new Stack[RDD[_]]
  def visit(rdd: RDD[_]) {
    if (!visited(rdd)) {
      visited += rdd
      val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
      if (rddHasUncachedPartitions) {
        for (dep <- rdd.dependencies) {
          dep match {
            case shufDep: ShuffleDependency[_, _, _] =>
              val mapStage = getShuffleMapStage(shufDep, stage.firstJobId)
              if (!mapStage.isAvailable) {
                missing += mapStage
              }
            case narrowDep: NarrowDependency[_] =>
              waitingForVisit.push(narrowDep.rdd)
          }
        }
      }
    }
  }
  waitingForVisit.push(stage.rdd)
  while (waitingForVisit.nonEmpty) {
    visit(waitingForVisit.pop())
  }
  missing.toList
}

这里的 parent stage 是通过 rdd 的依赖关系递归遍历获得的。对于宽依赖（Shuffle Dependency），Spark 会产生新的 mapStage作为 finalStage 的一个 missingParent，对于窄依赖（Narrow Dependency），Spark 不会产生新的 stage。这里对于 stage 的划分就是实现了论文中的方式。这种划分方法，好处在于快速恢复。如果一个 rdd 的 partition 丢失，该 partition 如果只有窄依赖，则其 parent rdd 也只需要计算相应的一个 partition 就能实现数据恢复。

正是这种 stage 的划分策略，才有了所谓的 DAG 图。当 stage 的 DAG 产生以后，会按树状结构，拓扑有序的执行一个个 stage，即当父 stage 都完成计算后，才可以进行子 stage 的计算：

/** Called when stage's parents are available and we can now do its task. */
private def submitMissingTasks(stage: Stage, jobId: Int) {
  var taskBinary: Broadcast[Array[Byte]] = null
  try {
    // 根据不同的 stage，生成相应的 task 的二进制内容，广播出去
    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.resultOfJob.get.func): AnyRef).array()
    }
    taskBinary = sc.broadcast(taskBinaryBytes)
  } catch {
    // ..
  }

  // 一个 stage 产生的 task 数目等于 partition 数目
  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.resultOfJob.get
        partitionsToCompute.map { id =>
          val p: Int = job.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 {
    //..
  }

  // 把这个 stage 对应的 tasks 提交到 taskSchduler 上
  if (tasks.size > 0) {
    taskScheduler.submitTasks(new TaskSet(
      tasks.toArray, stage.id, stage.latestInfo.attemptId, stage.firstJobId, properties))
    stage.latestInfo.submissionTime = Some(clock.getTimeMillis())
  }
}

至此，job 在 DAGScheduler 里完成了 stage DAG 图的构建，stage 到 tasks 组的转换，最后提交到 taskScheduler 上。当 task 被执行完毕，DAGScheduler.handleTaskCompletion 会被调用:

private[scheduler] def handleTaskCompletion(event: CompletionEvent) {
  val task = event.task
  val stageId = task.stageId
  event.reason match {
    case Success =>
      task match {
          case rt: ResultTask[_, _] =>
          ...
          case smt: ShuffleMapTask =>    
          ...
      }

    case Resubmitted =>
      ...
    case TaskResultLost =>
      ...
    case other =>
      ...
  }
}

遍历各种情况，分别处理。

RDD 的计算

RDD 的计算是在 task 中完成的，根据窄依赖和宽依赖，又分为 ResultTask 和 ShuffleMapTask，分别看一下具体执行过程：

// ResultTask
override def runTask(context: TaskContext): U = {
  // Deserialize the RDD and the func using the broadcast variables.
  val ser = SparkEnv.get.closureSerializer.newInstance()
  val (rdd, func) = ser.deserialize[(RDD[T], (TaskContext, Iterator[T]) => U)](
      ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
  func(context, rdd.iterator(partition, context))
}

override def runTask(context: TaskContext): MapStatus = {
  // Deserialize the RDD using the broadcast variable.
  val ser = SparkEnv.get.closureSerializer.newInstance()
  val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
      ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
  var writer: ShuffleWriter[Any, Any] = null
  try {
    val manager = SparkEnv.get.shuffleManager
    writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
    writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
    writer.stop(success = true).get
  } catch {
      ...
    }
  }

ResultTask 和 ShuffleMapTask 都会调用 SparkEnv.get.closureSerializer 对 taskBinary 进行反序列化操作，也都调用了 RDD.iterator 来计算和转换 RDD。不同之处在于， ResultTask 之后调用 func() 计算结果，而 ShuffleMapTask 把结果存入 blockManager 中用来 Shuffle。

至此，大致分析了 DAGScheduler 的执行过程。下一篇，我们再讲 TaskScheduler。