From 4ca3caa3f56e56139b2e33151ff95688232c4fe3 Mon Sep 17 00:00:00 2001 From: Davies Liu Date: Fri, 24 Oct 2014 13:22:58 -0700 Subject: [PATCH 1/4] use broadcast when task is large enough --- .../apache/spark/scheduler/DAGScheduler.scala | 19 ++++++++++++++----- .../apache/spark/scheduler/ResultTask.scala | 14 ++++++++++---- .../spark/scheduler/ShuffleMapTask.scala | 14 ++++++++++---- .../org/apache/spark/scheduler/Stage.scala | 5 +++++ .../spark/scheduler/TaskContextSuite.scala | 2 +- 5 files changed, 40 insertions(+), 14 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index f81fa6d8089fc..d3f001b297da6 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -34,7 +34,6 @@ import akka.pattern.ask import akka.util.Timeout import org.apache.spark._ -import org.apache.spark.broadcast.Broadcast import org.apache.spark.executor.TaskMetrics import org.apache.spark.partial.{ApproximateActionListener, ApproximateEvaluator, PartialResult} import org.apache.spark.rdd.RDD @@ -124,6 +123,9 @@ class DAGScheduler( /** If enabled, we may run certain actions like take() and first() locally. */ private val localExecutionEnabled = sc.getConf.getBoolean("spark.localExecution.enabled", false) + /** Broadcast the serialized tasks only when they are bigger than it */ + private val broadcastTaskMinSize = sc.getConf.getInt("spark.scheduler.broadcastTaskMinSize", 8) * 1024 + private def initializeEventProcessActor() { // blocking the thread until supervisor is started, which ensures eventProcessActor is // not null before any job is submitted @@ -817,22 +819,29 @@ class DAGScheduler( listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties)) // TODO: Maybe we can keep the taskBinary in Stage to avoid serializing it multiple times. - // Broadcasted binary for the task, used to dispatch tasks to executors. Note that we broadcast + // We serialize the task first, then broadcast it if the number of serialized bytes is + // larger than broadcastTaskMinSize, used to dispatch tasks to executors. Note that we broadcast // the serialized copy of the RDD and for each task we will deserialize it, which means each // task gets a different copy of the RDD. This provides stronger isolation between tasks that // might modify state of objects referenced in their closures. This is necessary in Hadoop // where the JobConf/Configuration object is not thread-safe. - var taskBinary: Broadcast[Array[Byte]] = null + var taskBinary: Array[Byte] = null try { // For ShuffleMapTask, serialize and broadcast (rdd, shuffleDep). // For ResultTask, serialize and broadcast (rdd, func). - val taskBinaryBytes: Array[Byte] = + taskBinary = if (stage.isShuffleMap) { closureSerializer.serialize((stage.rdd, stage.shuffleDep.get) : AnyRef).array() } else { closureSerializer.serialize((stage.rdd, stage.resultOfJob.get.func) : AnyRef).array() } - taskBinary = sc.broadcast(taskBinaryBytes) + if (taskBinary.size > broadcastTaskMinSize) { + logInfo(s"Create broadcast for taskBinary: ${taskBinary.size} > $broadcastTaskMinSize") + val broadcasted = sc.broadcast(taskBinary) + // use stage to track the life cycle of broadcast + stage.broadcastedTaskBinary = broadcasted + taskBinary = closureSerializer.serialize(broadcasted: AnyRef).array() + } } catch { // In the case of a failure during serialization, abort the stage. case e: NotSerializableException => diff --git a/core/src/main/scala/org/apache/spark/scheduler/ResultTask.scala b/core/src/main/scala/org/apache/spark/scheduler/ResultTask.scala index 4a9ff918afe25..16957ef21fed3 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ResultTask.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ResultTask.scala @@ -31,7 +31,7 @@ import org.apache.spark.rdd.RDD * See [[Task]] for more information. * * @param stageId id of the stage this task belongs to - * @param taskBinary broadcasted version of the serialized RDD and the function to apply on each + * @param taskBinary serialized RDD and the function to apply on each * partition of the given RDD. Once deserialized, the type should be * (RDD[T], (TaskContext, Iterator[T]) => U). * @param partition partition of the RDD this task is associated with @@ -41,7 +41,7 @@ import org.apache.spark.rdd.RDD */ private[spark] class ResultTask[T, U]( stageId: Int, - taskBinary: Broadcast[Array[Byte]], + taskBinary: Array[Byte], partition: Partition, @transient locs: Seq[TaskLocation], val outputId: Int) @@ -54,8 +54,14 @@ private[spark] class ResultTask[T, U]( 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) + val obj = ser.deserialize[AnyRef](ByteBuffer.wrap(taskBinary), + Thread.currentThread.getContextClassLoader) + val (rdd, func) = obj match { + case (rdd: RDD[T], func: ((TaskContext, Iterator[T]) => U)) => (rdd, func) + case b: Broadcast[Array[Byte]] => + ser.deserialize[(RDD[T], (TaskContext, Iterator[T]) => U)]( + ByteBuffer.wrap(b.value), Thread.currentThread.getContextClassLoader) + } metrics = Some(context.taskMetrics) func(context, rdd.iterator(partition, context)) diff --git a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala index 79709089c0da4..f9c04e6abd3f5 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/ShuffleMapTask.scala @@ -33,14 +33,14 @@ import org.apache.spark.shuffle.ShuffleWriter * See [[org.apache.spark.scheduler.Task]] for more information. * * @param stageId id of the stage this task belongs to - * @param taskBinary broadcast version of of the RDD and the ShuffleDependency. Once deserialized, + * @param taskBinary serialized RDD and the ShuffleDependency. Once deserialized, * the type should be (RDD[_], ShuffleDependency[_, _, _]). * @param partition partition of the RDD this task is associated with * @param locs preferred task execution locations for locality scheduling */ private[spark] class ShuffleMapTask( stageId: Int, - taskBinary: Broadcast[Array[Byte]], + taskBinary: Array[Byte], partition: Partition, @transient private var locs: Seq[TaskLocation]) extends Task[MapStatus](stageId, partition.index) with Logging { @@ -57,8 +57,14 @@ private[spark] class ShuffleMapTask( 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) + val obj = ser.deserialize[AnyRef](ByteBuffer.wrap(taskBinary), + Thread.currentThread.getContextClassLoader) + val (rdd, dep) = obj match { + case (rdd: RDD[_], dep: ShuffleDependency[_, _, _]) => (rdd, dep) + case b: Broadcast[Array[Byte]] => + ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])]( + ByteBuffer.wrap(b.value), Thread.currentThread.getContextClassLoader) + } metrics = Some(context.taskMetrics) var writer: ShuffleWriter[Any, Any] = null diff --git a/core/src/main/scala/org/apache/spark/scheduler/Stage.scala b/core/src/main/scala/org/apache/spark/scheduler/Stage.scala index 071568cdfb429..8f8d90771ecdd 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/Stage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/Stage.scala @@ -20,6 +20,7 @@ package org.apache.spark.scheduler import scala.collection.mutable.HashSet import org.apache.spark._ +import org.apache.spark.broadcast.Broadcast import org.apache.spark.rdd.RDD import org.apache.spark.storage.BlockManagerId import org.apache.spark.util.CallSite @@ -69,6 +70,10 @@ private[spark] class Stage( var resultOfJob: Option[ActiveJob] = None var pendingTasks = new HashSet[Task[_]] + /** This is used to track the life cycle of broadcast, + * then it can be release by GC once the stage is released */ + var broadcastedTaskBinary: Broadcast[Array[Byte]] = _ + private var nextAttemptId = 0 val name = callSite.shortForm diff --git a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala index 561a5e9cd90c4..8c52906bc187e 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala @@ -43,7 +43,7 @@ class TaskContextSuite extends FunSuite with BeforeAndAfter with LocalSparkConte val closureSerializer = SparkEnv.get.closureSerializer.newInstance() val func = (c: TaskContext, i: Iterator[String]) => i.next() val task = new ResultTask[String, String]( - 0, sc.broadcast(closureSerializer.serialize((rdd, func)).array), rdd.partitions(0), Seq(), 0) + 0, closureSerializer.serialize((rdd, func)).array, rdd.partitions(0), Seq(), 0) intercept[RuntimeException] { task.run(0) } From 28a9409325d8fe7299dd8e01598ef86e73f70fa9 Mon Sep 17 00:00:00 2001 From: Davies Liu Date: Fri, 24 Oct 2014 13:53:09 -0700 Subject: [PATCH 2/4] fix code style --- .../main/scala/org/apache/spark/scheduler/DAGScheduler.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index d3f001b297da6..55f291236ef32 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -124,7 +124,8 @@ class DAGScheduler( private val localExecutionEnabled = sc.getConf.getBoolean("spark.localExecution.enabled", false) /** Broadcast the serialized tasks only when they are bigger than it */ - private val broadcastTaskMinSize = sc.getConf.getInt("spark.scheduler.broadcastTaskMinSize", 8) * 1024 + private val broadcastTaskMinSize = + sc.getConf.getInt("spark.scheduler.broadcastTaskMinSize", 8) * 1024 private def initializeEventProcessActor() { // blocking the thread until supervisor is started, which ensures eventProcessActor is From 7dfe41e7449c1103840a389de87557a84f8d9d9d Mon Sep 17 00:00:00 2001 From: Davies Liu Date: Fri, 24 Oct 2014 16:57:15 -0700 Subject: [PATCH 3/4] fix style --- core/src/main/scala/org/apache/spark/scheduler/Stage.scala | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/Stage.scala b/core/src/main/scala/org/apache/spark/scheduler/Stage.scala index 8f8d90771ecdd..22a3f59f5518f 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/Stage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/Stage.scala @@ -70,8 +70,10 @@ private[spark] class Stage( var resultOfJob: Option[ActiveJob] = None var pendingTasks = new HashSet[Task[_]] - /** This is used to track the life cycle of broadcast, - * then it can be release by GC once the stage is released */ + /** + * This is used to track the life cycle of broadcast, + * then it can be release by GC once the stage is released + */ var broadcastedTaskBinary: Broadcast[Array[Byte]] = _ private var nextAttemptId = 0 From f3e20814411b36140edb726ac7e6b1dec1a8f939 Mon Sep 17 00:00:00 2001 From: Davies Liu Date: Sat, 25 Oct 2014 14:29:55 -0700 Subject: [PATCH 4/4] change config --- .../main/scala/org/apache/spark/scheduler/DAGScheduler.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala index 55f291236ef32..f22b5eaa90a5d 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala @@ -125,7 +125,7 @@ class DAGScheduler( /** Broadcast the serialized tasks only when they are bigger than it */ private val broadcastTaskMinSize = - sc.getConf.getInt("spark.scheduler.broadcastTaskMinSize", 8) * 1024 + sc.getConf.getInt("spark.scheduler.broadcastTaskMinSizeKB", 8) * 1024 private def initializeEventProcessActor() { // blocking the thread until supervisor is started, which ensures eventProcessActor is