Spark中shuffle的觸發和排程

devos發表於2015-09-11

Spark中的shuffle是在幹嘛?

Shuffle在Spark中即是把父RDD中的KV對按照Key重新分割槽,從而得到一個新的RDD。也就是說原本同屬於父RDD同一個分割槽的資料需要進入到子RDD的不同的分割槽。

但這只是shuffle的過程,卻不是shuffle的原因。為何需要shuffle呢?

Shuffle和Stage

在分散式計算框架中,比如map-reduce,資料本地化是一個很重要的考慮,即計算需要被分發到資料所在的位置,從而減少資料的移動,提高執行效率。

Map-Reduce的輸入資料通常是HDFS中的檔案,所以資料本地化要求map任務儘量被排程到儲存了輸入檔案的節點執行。但是,有一些計算邏輯是無法簡單地獲取本地資料的,reduce的邏輯都是如此。對於reduce來說,處理函式的輸入是key相同的所有value,但是這些value所在的資料集(即map的輸出)位於不同的節點上,因此需要對map的輸出進行重新組織,使得同樣的key進入相同的reducer。 shuffle移動了大量的資料,對計算、記憶體、網路和磁碟都有巨大的消耗,因此,只有確實需要shuffle的地方才應該進行shuffle。

Stage的劃分

對於Spark來說,計算的邏輯存在於RDD的轉換邏輯中。Spark的排程器也是在依據資料本地化在排程任務,只不過此處的“本地”不僅包括磁碟檔案,也包括RDD的分割槽, Spark會使得資料儘量少地被移動,據此,DAGScheduler把一個job劃分為多個Stage,在一個Stage內部,資料是不需要移動地,資料會在本地經過一系列函式的處理,直至確實需要shuffle的地方。

 例如,在DAGScheduler的getParentStages方法中,尋找父stage時,使用瞭如下的程式碼段

        for (dep <- r.dependencies) {
          dep match {
            case shufDep: ShuffleDependency[_, _, _] =>
              parents += getShuffleMapStage(shufDep, jobId)
            case _ =>
              waitingForVisit.push(dep.rdd)
          }

即找到了ShuffleDependency才會劃分出一個最的Stage(除了沒有父RDD的RDD,比如HadoopRDD,它的dependencies為Nil)。

在上邊的程式碼中,提到了ShuffleMapStage,其實Spark的Stage只有兩個子類:ShuffleStage和 ResultStage。相應的,Task也只有兩個子類,ResultTask和ShuffleMapTask。這些類之間的聯絡,可以從DAGScheduler的submitMissingTasks方法中表現中來。下面是這個方法中的一段程式碼:

  val tasks: Seq[Task[_]] = try {
      stage match {
        case stage: ShuffleMapStage =>
          partitionsToCompute.map { id =>
            val locs = getPreferredLocs(stage.rdd, id)
            val part = stage.rdd.partitions(id)
            new ShuffleMapTask(stage.id, taskBinary, part, locs)
          }

        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 = getPreferredLocs(stage.rdd, p)
            new ResultTask(stage.id, taskBinary, part, locs, id)
          }
      }
    } catch {
      case NonFatal(e) =>
        abortStage(stage, s"Task creation failed: $e\n${e.getStackTraceString}")
        runningStages -= stage
        return
    }

這段程式碼用來生成task, 確切地說是為某個Stage生成task。從以上程式碼可以看出,為ResultStage生成的就是ResultTask, 為ShuffleMapStage生成的就是ShuffleMapTask。

ShuffleMapTask有何特殊之處呢?

對於多於一個Stage的job,肯定會存在shuffle,這也意味會有Stage的父Stage是ShuffleMapStage。ShuffleMapStage中的ShuffleMapTask的最後一個RDD的資料會被進行shuffle,這也是它與ResultTask的區別。下邊是ShuffleMapTask的runTask方法中的一段程式碼,executor會間接呼叫runTask方法

      val manager = SparkEnv.get.shuffleManager//蕕取ShuffleManager
      //獲取writer,注意會把ShuffleDependency.shuffleHander傳過去
      writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
      writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
      return writer.stop(success = true).get

writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])這一句會計算最後一個rdd的某個分割槽,然後用writer寫入這個分割槽的資料,這可以認為是shuffle中的map階段。

那麼reduce階段是如何觸發的呢?

這實際上是很自然地由Spark對RDD的計算邏輯觸發的。

Spark的運算邏輯是由對RDD的partition的計算驅動的(上一篇提到過), 即對子RDD的partition的計算會觸發對父RDD的對應partition的計算,由此觸發到第一個可以計算的RDD的分割槽。所以shuffle關係子Stage中最初始的那個RDD一定包含有和shuffle過程相關的邏輯,這種特殊的RDD有兩類,ShuffledRDD和CoGroupedRDD,(後者不一定是shuffle的結果), 也就是說reduce是由對特殊RDD的計算觸發的。下面以ShuffledRDD為例進行說明,單個RDD進行shuffle會生成這種RDD。

ShuffledRDD

ShuffledRDD的特點由三部分可以體現。首先,它包括了一些跟shuffle有關的field:

  private var serializer: Option[Serializer] = None

  private var keyOrdering: Option[Ordering[K]] = None

  private var aggregator: Option[Aggregator[K, V, C]] = None

  private var mapSideCombine: Boolean = false

其中Aggregator主要用來指明對於同一個key對應的value,如何進行aggregate,但不僅於此。這是個挺有意思的類,它的域是一系列函式。

其次,它的dependency是ShuffleDependency,因此DAGScheduler會把它當作新Stage的起點,它的父RDD被當作前一個Stage的終點。

  override def getDependencies: Seq[Dependency[_]] = {
    List(new ShuffleDependency(prev, part, serializer, keyOrdering, aggregator, mapSideCombine))
  }

最後,當ShuffledRDD的某個partition被compute時,會觸發對map輸出的fetch,以及對value的aggregate等操作,也就是reduce階段。

  override def compute(split: Partition, context: TaskContext): Iterator[(K, C)] = {
    val dep = dependencies.head.asInstanceOf[ShuffleDependency[K, V, C]]
    SparkEnv.get.shuffleManager.getReader(dep.shuffleHandle, split.index, split.index + 1, context)
      .read()
      .asInstanceOf[Iterator[(K, C)]]
  }

那麼ShuffledRDD是如何生成的呢?

當然,會引起shuffle的transformation就會生成ShuffledRDD,以reduceByKey為例。

reduceByKey實際上有很多個過載的同名方法,以最簡單的為例

  def reduceByKey(func: (V, V) => V): RDD[(K, V)] = self.withScope {
    reduceByKey(defaultPartitioner(self), func)
  }
  def reduceByKey(partitioner: Partitioner, func: (V, V) => V): RDD[(K, V)] = self.withScope {
    combineByKey[V]((v: V) => v, func, func, partitioner)
  }

reduceByKey是在某個RDD上被呼叫的,設此RDD為A,呼叫reduceByKey生成的RDD為B。那麼,以上程式碼中的partitioner是指用於生成B的Partitioner, 它指出了A中的每個kv對應該進行B的哪個分割槽。之所以需要注意這點,是因為在combineByKey中會根據這個Partitioner決定需要生成的RDD,在特定情況下reduceByKey不會導致shuffle.

下面是combineByKey中用於決定生成何種RDD的程式碼:

  if (self.partitioner == Some(partitioner)) {
      self.mapPartitions(iter => {
        val context = TaskContext.get()
        new InterruptibleIterator(context, aggregator.combineValuesByKey(iter, context))
      }, preservesPartitioning = true)
    } else {
      new ShuffledRDD[K, V, C](self, partitioner)
        .setSerializer(serializer)
        .setAggregator(aggregator)
        .setMapSideCombine(mapSideCombine)
    }

它會根據

if (self.partitioner == Some(partitioner)) 

來決定是否生成ShuffledRDD。其中self.partitioner是指A這個RDD的partitioner,它指明瞭A這個RDD中的每個key在哪個partition中。而等號右邊的partitioner,指明瞭B這個RDD的每個key在哪個partition中。當二者==時,就會用self.mapPartitions生成MapPartitionsRDD, 這和map這種transformation生成的RDD是一樣的,此時reduceByKey不會引發shuffle。

Partitioner有幾個子類,它們中的某些會override預設的equals方法(注意,Scala中的==會呼叫equals方法,這點和Java不同)。典型的,如HashPartitioner中的equals方法

  override def equals(other: Any): Boolean = other match {
    case h: HashPartitioner =>
      h.numPartitions == numPartitions
    case _ =>
      false
  }

當兩個HashPartitioner的分割槽數目一致時,就認為他們相等。但是,即是A和B有相同的Partitioner,也只決定了這兩個RDD中相同的key在同一個partition中,並不意味著A中相同的key對應的value已經被aggregate了,因此在combineByKey操作中呼叫mapPartitions方法時,指定了特殊的Iterator到Iterator的轉換方法。

 new InterruptibleIterator(context, aggregator.combineValuesByKey(iter, context))

也就是說A中partition的Iterator會被執行combineValuesByKey操作,來對value進行aggregate。對於reduceByKey,不管需不需要進行shuffle,對value進行aggregate都是要執行的。比如,在ShuffledRDD的compute方法中,會呼叫ShuffleReader的read方法。ShuffleReader當前只有一種,叫HashShuffleReader, 不管是用sort還是hash進行shuffle,reduce端都是使用的這個Reader,它會對從map端抓取資料後生成的iterator進行aggregate

    val aggregatedIter: Iterator[Product2[K, C]] = if (dep.aggregator.isDefined) {
      if (dep.mapSideCombine) {
        new InterruptibleIterator(context, dep.aggregator.get.combineCombinersByKey(iter, context))
      } else {
        new InterruptibleIterator(context, dep.aggregator.get.combineValuesByKey(iter, context))
      }
    } else {
      require(!dep.mapSideCombine, "Map-side combine without Aggregator specified!")

      // Convert the Product2s to pairs since this is what downstream RDDs currently expect
      iter.asInstanceOf[Iterator[Product2[K, C]]].map(pair => (pair._1, pair._2))
    }

在上邊combineByKey的程式碼中,可以看到它生成ShuffledRDD時,設定了aggreator,而mapSideCombine使用了預設引數,為true,所以combineCombinerByKey會被呼叫,來對已經combine好的value進行combine。

總結

通過上邊的內容,基本可以瞭解到DAGScheduler是如何處理根據shuffle劃分Stage,生成特殊的task;以及Spark執行過程中,map和reduce兩個階段是如何被觸發的。

總的是來說, RDD的轉換操作會盡量避免shuffle的出現,如果不得不shuffle,會生成特殊的RDD,它的dependencies會是ShuffleDependency。DAGScheduler在劃分Stage時,會用ShuffleDependency確定Stage的邊界,也會由此生成ShuffleMapTask來完成map端的工作。引發shuffle的transformation會生成特殊的RDD,此RDD會是shuffle中子Stage的起點,當這些RDD的compute方法被呼叫時,就會觸發reduce端操作的執行。這種特殊的RDD有兩類:

ShuffledRDD, 它只有一個父RDD,是對一個RDD進行shuffle的結果。

CoGroupedRDD, 它有多個RDD,是對多個RDD進行shuffle的結果。

 

相關文章