菜鸟的Spark 源码学习之路 -3 TaskScheduler源码

上一篇文章里，我们解了SparkContext的源码：https://mp.youkuaiyun.com/postedit/81478344

我们看到，里面有很多数据结构，其中一个就是TaskScheduler，顾名思义，它就是做任务调度用的。

1. 源码概览

找到TaskScheduler源码,我们先看一下它的结构和定义

可以看到，TaskScheduler是一个trait，因此我们再看一下它的实现类：

可以看到有两个匿名实现，还有两个直接实现类。点进源码后发现，重要功能的实现都在TaskSchedulerImpl里面：

2. 源码深入

我们看一下TaskSchedulerImpl的注释怎么说。

上面讲到，这个类主要用于调度集群上的任务，这些操作都是通过SchedulerBackend来进行的，进入SchedulerBackend源码中发现，很多工作都是调用TaskSchedulerImpl的方法来完成的。这种关系类似于SchedulerBackend的一个代理类。

TaskSchedulerImpl主要处理通用逻辑，比如，决定job间的调度顺序，启动推测任务等。客户端首先需调用initialize() and start()，然后通过runTasks方法提交任务设置。

2.1 initialize()方法

我们就先从initialize()方法开始进入TaskScheduler源码学习。

//这里主要是根据schedulingMode完成schedulableBuilder的初始化工作

def initialize(backend: SchedulerBackend) {
  this.backend = backend
  schedulableBuilder = {
    schedulingMode match {
      case SchedulingMode.FIFO =>
        new FIFOSchedulableBuilder(rootPool)
      case SchedulingMode.FAIR =>
        new FairSchedulableBuilder(rootPool, conf)
      case _ =>
        throw new IllegalArgumentException(s"Unsupported $SCHEDULER_MODE_PROPERTY: " +
          s"$schedulingMode")
    }
  }
  schedulableBuilder.buildPools()
}

看一下SchedulableBuilder主要做什么工作

SchedulableBuilder 是一个构建 Schedulable 树的接口，buildPools()方法主要就是构建树节点，addTaskSetManager()方法主要是构建叶子节点。

再看一下schedulingMode，默认是FIFO

schedulableBuilder 有一个Pool类型参数，再看rootPool的初始化，

默认初始化了一个名字为空的pool，我们先看一下Pool的作用：

Pool是一个可调度的实体，是Pools或者TaskSetManagers的集合的抽象。也就是说Pool可以包含其他的Pool或者TaskSetManagers。

再回到initialize()方法，最后调用了buildPools()方法

如果是FIFOSchedulableBuilder，实际buildPools()方法什么都没有做；

如果是FairSchedulableBuilder，buildPools()做了一些较为复杂的初始化工作，这个我们后续再看。

整体来看 initialize() 方法主要就是完成SchedulableBuilder变量的初始化工作。

2.2 start() 方法

override def start() {
// 启动backend， 并根据一些判断条件确定是否启动推测执行
  backend.start()
  if (!isLocal && conf.getBoolean("spark.speculation", false)) {
    logInfo("Starting speculative execution thread")
    speculationScheduler.scheduleWithFixedDelay(new Runnable {
      override def run(): Unit = Utils.tryOrStopSparkContext(sc) {
        checkSpeculatableTasks()
      }
    }, SPECULATION_INTERVAL_MS, SPECULATION_INTERVAL_MS, TimeUnit.MILLISECONDS)
  }
}

2.3 postStartHook()

在 SparkContext初始化代码块中，先是调用了tasknScheduler的initialize(), 然后是start(), 然后是postStartHook()。我们先看看

TaskScheduler里面这个方法的解释

// Invoked after system has successfully initialized (typically in spark context).
// 在系统成功初始化（通常是在spark context的初始化）后调用
// Yarn uses this to bootstrap allocation of resources based on preferred locations,
// yarn 使用它基于启动资源分配
// wait for slave registrations, etc.
// 等待slave的注册等
def postStartHook() { }

在实现类中，它调用了一个私有方法： waitBackendReady()

override def postStartHook() {
  waitBackendReady()
}

这个方法主要是判断 backend是否已经准备好，如果没有就看context是否已经停止，如果停止则抛出异常，如果没有则继续等待其初始化。

2.4 submitTasks(taskSet: TaskSet)

接下来是submitTasks方法，TaskScheduler定义该方法如下：

// Submit a sequence of tasks to run.
def submitTasks(taskSet: TaskSet): Unit

可以看出这个是spark 提交任务的方法。实现类的源码如下：

override def submitTasks(taskSet: TaskSet) {
  val tasks = taskSet.tasks
  logInfo("Adding task set " + taskSet.id + " with " + tasks.length + " tasks")
  this.synchronized {
    // 创建一个TaskSetManager
    val manager = createTaskSetManager(taskSet, maxTaskFailures)
    val stage = taskSet.stageId
    val stageTaskSets =
      taskSetsByStageIdAndAttempt.getOrElseUpdate(stage, new HashMap[Int, TaskSetManager])
    // 为特定的stage 指定taskSetManager
    stageTaskSets(taskSet.stageAttemptId) = manager
    val conflictingTaskSet = stageTaskSets.exists { case (_, ts) =>
      ts.taskSet != taskSet && !ts.isZombie
    }
    if (conflictingTaskSet) {
      throw new IllegalStateException(s"more than one active taskSet for stage $stage:" +
        s" ${
          stageTaskSets.toSeq.map {
            _._2.taskSet.id
          }.mkString(",")
        }")
    }
    // 为schedulableBuilder添加manager，这里就是构建了schedulableBuilder的树的叶子节点
    schedulableBuilder.addTaskSetManager(manager, manager.taskSet.properties)

    if (!isLocal && !hasReceivedTask) {
      starvationTimer.scheduleAtFixedRate(new TimerTask() {
        override def run() {
          if (!hasLaunchedTask) {
            logWarning("Initial job has not accepted any resources; " +
              "check your cluster UI to ensure that workers are registered " +
              "and have sufficient resources")
          } else {
            this.cancel()
          }
        }
      }, STARVATION_TIMEOUT_MS, STARVATION_TIMEOUT_MS)
    }
    hasReceivedTask = true
  }
  backend.reviveOffers()
}

该方法有个TaskSet类型的参数，其数据结构如下：

private[spark] class TaskSet(
    val tasks: Array[Task[_]],
    val stageId: Int,
    val stageAttemptId: Int,
    val priority: Int,
    val properties: Properties) {
  val id: String = stageId + "." + stageAttemptId

  override def toString: String = "TaskSet " + id
}

其实主要是记录跟某个stage关联的诸多task和stage的相关信息。

接下来我们再看一下到底哪些地方调用了submitTasks()方法，首先就是在DAGScheduler中。