From 4ac1896c6cc082c6dd1c427cb89d36a3cc233eba Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Mon, 6 Jul 2015 17:44:16 -0700
Subject: [PATCH 01/27] Refactor checkpoint interface for modularity

This commit makes two classes abstract: `RDDCheckpointData` and
`CheckpointRDD`. It implements the existing fault-tolerant
checkpointing by subclassing these abstract classes. The goal
of this commit is to retain as much functionality as possible.
Much of the code is just moved from one file to another.

The following commits will add an implementation for unreliable
checkpointing.
---
 .../org/apache/spark/ContextCleaner.scala     |  11 +-
 .../scala/org/apache/spark/SparkContext.scala |   2 +-
 .../org/apache/spark/rdd/CheckpointRDD.scala  | 141 +--------------
 .../main/scala/org/apache/spark/rdd/RDD.scala |  37 +++-
 .../apache/spark/rdd/RDDCheckpointData.scala  | 105 ++++-------
 .../spark/rdd/ReliableCheckpointRDD.scala     | 171 ++++++++++++++++++
 .../spark/rdd/ReliableRDDCheckpointData.scala | 109 +++++++++++
 .../apache/spark/ContextCleanerSuite.scala    |  12 +-
 8 files changed, 363 insertions(+), 225 deletions(-)
 create mode 100644 core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
 create mode 100644 core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala

diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
index 37198d887b07b..825081eef2443 100644
--- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala
+++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
@@ -22,7 +22,7 @@ import java.lang.ref.{ReferenceQueue, WeakReference}
 import scala.collection.mutable.{ArrayBuffer, SynchronizedBuffer}
 
 import org.apache.spark.broadcast.Broadcast
-import org.apache.spark.rdd.{RDDCheckpointData, RDD}
+import org.apache.spark.rdd.{ReliableRDDCheckpointData, RDD}
 import org.apache.spark.util.Utils
 
 /**
@@ -231,11 +231,16 @@ private[spark] class ContextCleaner(sc: SparkContext) extends Logging {
     }
   }
 
-  /** Perform checkpoint cleanup. */
+  /**
+   * Perform checkpoint cleanup.
+   *
+   * Note that this only cleans up reliable checkpoint files.
+   * Unsafe checkpoint files are just cached RDD blocks and are cleaned up separately.
+   */
   def doCleanCheckpoint(rddId: Int): Unit = {
     try {
       logDebug("Cleaning rdd checkpoint data " + rddId)
-      RDDCheckpointData.clearRDDCheckpointData(sc, rddId)
+      ReliableRDDCheckpointData.clearRDDCheckpointData(sc, rddId)
       listeners.foreach(_.checkpointCleaned(rddId))
       logInfo("Cleaned rdd checkpoint data " + rddId)
     }
diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index d2547eeff2b4e..cec8e4a1351aa 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1195,7 +1195,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
   }
 
   protected[spark] def checkpointFile[T: ClassTag](path: String): RDD[T] = withScope {
-    new CheckpointRDD[T](this, path)
+    new ReliableCheckpointRDD[T](this, path)
   }
 
   /** Build the union of a list of RDDs. */
diff --git a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
index e17bd47905d7a..6681bdd14940e 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
@@ -17,156 +17,19 @@
 
 package org.apache.spark.rdd
 
-import java.io.IOException
-
 import scala.reflect.ClassTag
 
-import org.apache.hadoop.fs.Path
-
-import org.apache.spark._
-import org.apache.spark.broadcast.Broadcast
-import org.apache.spark.deploy.SparkHadoopUtil
-import org.apache.spark.util.{SerializableConfiguration, Utils}
+import org.apache.spark.{Partition, SparkContext}
 
 private[spark] class CheckpointRDDPartition(val index: Int) extends Partition
 
 /**
  * This RDD represents a RDD checkpoint file (similar to HadoopRDD).
  */
-private[spark]
-class CheckpointRDD[T: ClassTag](sc: SparkContext, val checkpointPath: String)
+private[spark] abstract class CheckpointRDD[T: ClassTag](sc: SparkContext)
   extends RDD[T](sc, Nil) {
 
-  private val broadcastedConf = sc.broadcast(new SerializableConfiguration(sc.hadoopConfiguration))
-
-  @transient private val fs = new Path(checkpointPath).getFileSystem(sc.hadoopConfiguration)
-
-  override def getCheckpointFile: Option[String] = Some(checkpointPath)
-
-  override def getPartitions: Array[Partition] = {
-    val cpath = new Path(checkpointPath)
-    val numPartitions =
-    // listStatus can throw exception if path does not exist.
-    if (fs.exists(cpath)) {
-      val dirContents = fs.listStatus(cpath).map(_.getPath)
-      val partitionFiles = dirContents.filter(_.getName.startsWith("part-")).map(_.toString).sorted
-      val numPart = partitionFiles.length
-      if (numPart > 0 && (! partitionFiles(0).endsWith(CheckpointRDD.splitIdToFile(0)) ||
-          ! partitionFiles(numPart-1).endsWith(CheckpointRDD.splitIdToFile(numPart-1)))) {
-        throw new SparkException("Invalid checkpoint directory: " + checkpointPath)
-      }
-      numPart
-    } else 0
-
-    Array.tabulate(numPartitions)(i => new CheckpointRDDPartition(i))
-  }
-
-  override def getPreferredLocations(split: Partition): Seq[String] = {
-    val status = fs.getFileStatus(new Path(checkpointPath,
-      CheckpointRDD.splitIdToFile(split.index)))
-    val locations = fs.getFileBlockLocations(status, 0, status.getLen)
-    locations.headOption.toList.flatMap(_.getHosts).filter(_ != "localhost")
-  }
-
-  override def compute(split: Partition, context: TaskContext): Iterator[T] = {
-    val file = new Path(checkpointPath, CheckpointRDD.splitIdToFile(split.index))
-    CheckpointRDD.readFromFile(file, broadcastedConf, context)
-  }
-
   // CheckpointRDD should not be checkpointed again
   override def checkpoint(): Unit = { }
   override def doCheckpoint(): Unit = { }
 }
-
-private[spark] object CheckpointRDD extends Logging {
-  def splitIdToFile(splitId: Int): String = {
-    "part-%05d".format(splitId)
-  }
-
-  def writeToFile[T: ClassTag](
-      path: String,
-      broadcastedConf: Broadcast[SerializableConfiguration],
-      blockSize: Int = -1
-    )(ctx: TaskContext, iterator: Iterator[T]) {
-    val env = SparkEnv.get
-    val outputDir = new Path(path)
-    val fs = outputDir.getFileSystem(broadcastedConf.value.value)
-
-    val finalOutputName = splitIdToFile(ctx.partitionId)
-    val finalOutputPath = new Path(outputDir, finalOutputName)
-    val tempOutputPath =
-      new Path(outputDir, "." + finalOutputName + "-attempt-" + ctx.attemptNumber)
-
-    if (fs.exists(tempOutputPath)) {
-      throw new IOException("Checkpoint failed: temporary path " +
-        tempOutputPath + " already exists")
-    }
-    val bufferSize = env.conf.getInt("spark.buffer.size", 65536)
-
-    val fileOutputStream = if (blockSize < 0) {
-      fs.create(tempOutputPath, false, bufferSize)
-    } else {
-      // This is mainly for testing purpose
-      fs.create(tempOutputPath, false, bufferSize, fs.getDefaultReplication, blockSize)
-    }
-    val serializer = env.serializer.newInstance()
-    val serializeStream = serializer.serializeStream(fileOutputStream)
-    Utils.tryWithSafeFinally {
-      serializeStream.writeAll(iterator)
-    } {
-      serializeStream.close()
-    }
-
-    if (!fs.rename(tempOutputPath, finalOutputPath)) {
-      if (!fs.exists(finalOutputPath)) {
-        logInfo("Deleting tempOutputPath " + tempOutputPath)
-        fs.delete(tempOutputPath, false)
-        throw new IOException("Checkpoint failed: failed to save output of task: "
-          + ctx.attemptNumber + " and final output path does not exist")
-      } else {
-        // Some other copy of this task must've finished before us and renamed it
-        logInfo("Final output path " + finalOutputPath + " already exists; not overwriting it")
-        fs.delete(tempOutputPath, false)
-      }
-    }
-  }
-
-  def readFromFile[T](
-      path: Path,
-      broadcastedConf: Broadcast[SerializableConfiguration],
-      context: TaskContext
-    ): Iterator[T] = {
-    val env = SparkEnv.get
-    val fs = path.getFileSystem(broadcastedConf.value.value)
-    val bufferSize = env.conf.getInt("spark.buffer.size", 65536)
-    val fileInputStream = fs.open(path, bufferSize)
-    val serializer = env.serializer.newInstance()
-    val deserializeStream = serializer.deserializeStream(fileInputStream)
-
-    // Register an on-task-completion callback to close the input stream.
-    context.addTaskCompletionListener(context => deserializeStream.close())
-
-    deserializeStream.asIterator.asInstanceOf[Iterator[T]]
-  }
-
-  // Test whether CheckpointRDD generate expected number of partitions despite
-  // each split file having multiple blocks. This needs to be run on a
-  // cluster (mesos or standalone) using HDFS.
-  def main(args: Array[String]) {
-    import org.apache.spark._
-
-    val Array(cluster, hdfsPath) = args
-    val env = SparkEnv.get
-    val sc = new SparkContext(cluster, "CheckpointRDD Test")
-    val rdd = sc.makeRDD(1 to 10, 10).flatMap(x => 1 to 10000)
-    val path = new Path(hdfsPath, "temp")
-    val conf = SparkHadoopUtil.get.newConfiguration(new SparkConf())
-    val fs = path.getFileSystem(conf)
-    val broadcastedConf = sc.broadcast(new SerializableConfiguration(conf))
-    sc.runJob(rdd, CheckpointRDD.writeToFile[Int](path.toString, broadcastedConf, 1024) _)
-    val cpRDD = new CheckpointRDD[Int](sc, path.toString)
-    assert(cpRDD.partitions.length == rdd.partitions.length, "Number of partitions is not the same")
-    assert(cpRDD.collect.toList == rdd.collect.toList, "Data of partitions not the same")
-    fs.delete(path, true)
-  }
-}
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 9f7ebae3e9af3..ed3da0a38c0f3 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1446,7 +1446,7 @@ abstract class RDD[T: ClassTag](
 
   /**
    * Mark this RDD for checkpointing. It will be saved to a file inside the checkpoint
-   * directory set with SparkContext.setCheckpointDir() and all references to its parent
+   * directory set with `SparkContext#setCheckpointDir` and all references to its parent
    * RDDs will be removed. This function must be called before any job has been
    * executed on this RDD. It is strongly recommended that this RDD is persisted in
    * memory, otherwise saving it on a file will require recomputation.
@@ -1459,20 +1459,45 @@ abstract class RDD[T: ClassTag](
       // children RDD partitions point to the correct parent partitions. In the future
       // we should revisit this consideration.
       RDDCheckpointData.synchronized {
-        checkpointData = Some(new RDDCheckpointData(this))
+        checkpointData = Some(new ReliableRDDCheckpointData(this))
       }
     }
   }
 
   /**
-   * Return whether this RDD has been checkpointed or not
+   * Mark this RDD for unsafe checkpointing using a local file system.
+   *
+   * This method is for users who wish to truncate the lineage of the RDD while skipping the
+   * expensive step of replicating the materialized data in a distributed file system. This is
+   * useful for RDDs with long lineages that need to be truncated periodically (e.g. GraphX).
+   *
+   * It is unsafe because we sacrifice fault-tolerance for performance. In particular, we
+   * persist the materialized values only to an ephemeral local storage. The effect is that
+   * if an executor fails during the computation, the checkpointed data will no longer be
+   * accessible.
+   *
+   * The actual persisting occurs after the first job involving this RDD has completed.
+   * The checkpoint directory set through `SparkContext#setCheckpointDir` is not read.
+   */
+  def localCheckpoint(): Unit = {
+
+  }
+
+  /**
+   * Return whether this RDD has been checkpointed or not, either reliably or locally.
    */
   def isCheckpointed: Boolean = checkpointData.exists(_.isCheckpointed)
 
   /**
-   * Gets the name of the file to which this RDD was checkpointed
+   * Gets the name of the directory to which this RDD was checkpointed.
+   * This is not defined if the RDD is checkpointed locally.
    */
-  def getCheckpointFile: Option[String] = checkpointData.flatMap(_.getCheckpointFile)
+  def getCheckpointFile: Option[String] = {
+    checkpointData match {
+      case Some(reliable: ReliableRDDCheckpointData[T]) => reliable.getCheckpointDir
+      case _ => None
+    }
+  }
 
   // =======================================================================
   // Other internal methods and fields
@@ -1543,7 +1568,7 @@ abstract class RDD[T: ClassTag](
       if (!doCheckpointCalled) {
         doCheckpointCalled = true
         if (checkpointData.isDefined) {
-          checkpointData.get.doCheckpoint()
+          checkpointData.get.checkpoint()
         } else {
           dependencies.foreach(_.rdd.doCheckpoint())
         }
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
index 4f954363bed8e..5df420cce1c65 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
@@ -19,10 +19,7 @@ package org.apache.spark.rdd
 
 import scala.reflect.ClassTag
 
-import org.apache.hadoop.fs.Path
-
-import org.apache.spark._
-import org.apache.spark.util.SerializableConfiguration
+import org.apache.spark.Partition
 
 /**
  * Enumeration to manage state transitions of an RDD through checkpointing
@@ -39,39 +36,30 @@ private[spark] object CheckpointState extends Enumeration {
  * as well as, manages the post-checkpoint state by providing the updated partitions,
  * iterator and preferred locations of the checkpointed RDD.
  */
-private[spark] class RDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
-  extends Logging with Serializable {
+private[spark] abstract class RDDCheckpointData[T: ClassTag](@transient rdd: RDD[T]) {
 
   import CheckpointState._
 
   // The checkpoint state of the associated RDD.
-  private var cpState = Initialized
-
-  // The file to which the associated RDD has been checkpointed to
-  private var cpFile: Option[String] = None
+  protected var cpState = Initialized
 
-  // The CheckpointRDD created from the checkpoint file, that is, the new parent the associated RDD.
-  // This is defined if and only if `cpState` is `Checkpointed`.
+  // The RDD that contains our checkpointed data
   private var cpRDD: Option[CheckpointRDD[T]] = None
 
   // TODO: are we sure we need to use a global lock in the following methods?
 
-  // Is the RDD already checkpointed
+  /**
+   * Return whether the checkpoint data for this RDD is already persisted.
+   */
   def isCheckpointed: Boolean = RDDCheckpointData.synchronized {
     cpState == Checkpointed
   }
 
-  // Get the file to which this RDD was checkpointed to as an Option
-  def getCheckpointFile: Option[String] = RDDCheckpointData.synchronized {
-    cpFile
-  }
-
   /**
-   * Materialize this RDD and write its content to a reliable DFS.
+   * Materialize this RDD and persist its content.
    * This is called immediately after the first action invoked on this RDD has completed.
    */
-  def doCheckpoint(): Unit = {
-
+  final def checkpoint(): Unit = {
     // Guard against multiple threads checkpointing the same RDD by
     // atomically flipping the state of this RDDCheckpointData
     RDDCheckpointData.synchronized {
@@ -82,64 +70,41 @@ private[spark] class RDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
       }
     }
 
-    // Create the output path for the checkpoint
-    val path = RDDCheckpointData.rddCheckpointDataPath(rdd.context, rdd.id).get
-    val fs = path.getFileSystem(rdd.context.hadoopConfiguration)
-    if (!fs.mkdirs(path)) {
-      throw new SparkException(s"Failed to create checkpoint path $path")
-    }
-
-    // Save to file, and reload it as an RDD
-    val broadcastedConf = rdd.context.broadcast(
-      new SerializableConfiguration(rdd.context.hadoopConfiguration))
-    val newRDD = new CheckpointRDD[T](rdd.context, path.toString)
-    if (rdd.conf.getBoolean("spark.cleaner.referenceTracking.cleanCheckpoints", false)) {
-      rdd.context.cleaner.foreach { cleaner =>
-        cleaner.registerRDDCheckpointDataForCleanup(newRDD, rdd.id)
-      }
-    }
-
-    // TODO: This is expensive because it computes the RDD again unnecessarily (SPARK-8582)
-    rdd.context.runJob(rdd, CheckpointRDD.writeToFile[T](path.toString, broadcastedConf) _)
-    if (newRDD.partitions.length != rdd.partitions.length) {
-      throw new SparkException(
-        "Checkpoint RDD " + newRDD + "(" + newRDD.partitions.length + ") has different " +
-          "number of partitions than original RDD " + rdd + "(" + rdd.partitions.length + ")")
-    }
+    val newRDD = doCheckpoint()
 
-    // Change the dependencies and partitions of the RDD
+    // Update our state and mark the RDD as checkpointed
     RDDCheckpointData.synchronized {
-      cpFile = Some(path.toString)
       cpRDD = Some(newRDD)
-      rdd.markCheckpointed(newRDD)   // Update the RDD's dependencies and partitions
       cpState = Checkpointed
+      rdd.markCheckpointed(newRDD)
     }
-    logInfo(s"Done checkpointing RDD ${rdd.id} to $path, new parent is RDD ${newRDD.id}")
-  }
-
-  def getPartitions: Array[Partition] = RDDCheckpointData.synchronized {
-    cpRDD.get.partitions
   }
 
-  def checkpointRDD: Option[CheckpointRDD[T]] = RDDCheckpointData.synchronized {
-    cpRDD
-  }
-}
+  /**
+   * Materialize this RDD and persist its content.
+   *
+   * Subclasses should override this method to define custom checkpointing behavior.
+   * @return the checkpoint RDD created in the process.
+   */
+  protected def doCheckpoint(): CheckpointRDD[T]
 
-private[spark] object RDDCheckpointData {
+  /**
+   * Return the RDD that contains our checkpointed data.
+   * This is only defined if the checkpoint state is `Checkpointed`.
+   */
+  def checkpointRDD: Option[CheckpointRDD[T]] = RDDCheckpointData.synchronized { cpRDD }
 
-  /** Return the path of the directory to which this RDD's checkpoint data is written. */
-  def rddCheckpointDataPath(sc: SparkContext, rddId: Int): Option[Path] = {
-    sc.checkpointDir.map { dir => new Path(dir, s"rdd-$rddId") }
+  /**
+   * Return the partitions of the resulting checkpoint RDD.
+   * For tests only.
+   */
+  def getPartitions: Array[Partition] = RDDCheckpointData.synchronized {
+    cpRDD.map(_.partitions).getOrElse { Array.empty }
   }
 
-  /** Clean up the files associated with the checkpoint data for this RDD. */
-  def clearRDDCheckpointData(sc: SparkContext, rddId: Int): Unit = {
-    rddCheckpointDataPath(sc, rddId).foreach { path =>
-      val fs = path.getFileSystem(sc.hadoopConfiguration)
-      if (fs.exists(path)) {
-        fs.delete(path, true)
-      }
-    }
-  }
 }
+
+/**
+ * Global lock for synchronizing checkpoint operations.
+ */
+private[spark] object RDDCheckpointData
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
new file mode 100644
index 0000000000000..ac2930d015df6
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -0,0 +1,171 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import java.io.IOException
+
+import scala.reflect.ClassTag
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark._
+import org.apache.spark.broadcast.Broadcast
+import org.apache.spark.deploy.SparkHadoopUtil
+import org.apache.spark.util.{SerializableConfiguration, Utils}
+
+/**
+ * An RDD that reads from a checkpoint file previously written to reliable storage.
+ */
+private[spark] class ReliableCheckpointRDD[T: ClassTag](
+    sc: SparkContext,
+    val checkpointPath: String)
+  extends CheckpointRDD[T](sc) {
+
+  private val broadcastedConf = sc.broadcast(new SerializableConfiguration(sc.hadoopConfiguration))
+
+  @transient private val fs = new Path(checkpointPath).getFileSystem(sc.hadoopConfiguration)
+
+  override def getCheckpointFile: Option[String] = Some(checkpointPath)
+
+  override def getPartitions: Array[Partition] = {
+    val cpath = new Path(checkpointPath)
+    val numPartitions =
+      // listStatus can throw exception if path does not exist.
+      if (fs.exists(cpath)) {
+        val dirContents = fs.listStatus(cpath).map(_.getPath)
+        val partitionFiles = dirContents
+          .filter(_.getName.startsWith("part-"))
+          .map(_.toString)
+          .sorted
+        // Verify validity of checkpoint files
+        val numPart = partitionFiles.length
+        val firstPartitionValid =
+          partitionFiles(0).endsWith(ReliableCheckpointRDD.splitIdToFile(0))
+        val lastPartitionValid =
+          partitionFiles(numPart - 1).endsWith(ReliableCheckpointRDD.splitIdToFile(numPart - 1))
+        if (numPart > 0 && (!firstPartitionValid || !lastPartitionValid)) {
+          throw new SparkException(s"Invalid checkpoint directory: $checkpointPath")
+        }
+        numPart
+      } else {
+        0
+      }
+    Array.tabulate(numPartitions)(i => new CheckpointRDDPartition(i))
+  }
+
+  override def getPreferredLocations(split: Partition): Seq[String] = {
+    val status = fs.getFileStatus(
+      new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index)))
+    val locations = fs.getFileBlockLocations(status, 0, status.getLen)
+    locations.headOption.toList.flatMap(_.getHosts).filter(_ != "localhost")
+  }
+
+  override def compute(split: Partition, context: TaskContext): Iterator[T] = {
+    val file = new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index))
+    ReliableCheckpointRDD.readFromFile(file, broadcastedConf, context)
+  }
+}
+
+private[spark] object ReliableCheckpointRDD extends Logging {
+
+  private def splitIdToFile(splitId: Int): String = {
+    "part-%05d".format(splitId)
+  }
+
+  def writeToFile[T: ClassTag](
+      path: String,
+      broadcastedConf: Broadcast[SerializableConfiguration],
+      blockSize: Int = -1)(ctx: TaskContext, iterator: Iterator[T]) {
+    val env = SparkEnv.get
+    val outputDir = new Path(path)
+    val fs = outputDir.getFileSystem(broadcastedConf.value.value)
+
+    val finalOutputName = ReliableCheckpointRDD.splitIdToFile(ctx.partitionId())
+    val finalOutputPath = new Path(outputDir, finalOutputName)
+    val tempOutputPath =
+      new Path(outputDir, s".$finalOutputName-attempt-${ctx.attemptNumber()}")
+
+    if (fs.exists(tempOutputPath)) {
+      throw new IOException(s"Checkpoint failed: temporary path $tempOutputPath already exists")
+    }
+    val bufferSize = env.conf.getInt("spark.buffer.size", 65536)
+
+    val fileOutputStream = if (blockSize < 0) {
+      fs.create(tempOutputPath, false, bufferSize)
+    } else {
+      // This is mainly for testing purpose
+      fs.create(tempOutputPath, false, bufferSize, fs.getDefaultReplication, blockSize)
+    }
+    val serializer = env.serializer.newInstance()
+    val serializeStream = serializer.serializeStream(fileOutputStream)
+    Utils.tryWithSafeFinally {
+      serializeStream.writeAll(iterator)
+    } {
+      serializeStream.close()
+    }
+
+    if (!fs.rename(tempOutputPath, finalOutputPath)) {
+      if (!fs.exists(finalOutputPath)) {
+        logInfo(s"Deleting tempOutputPath $tempOutputPath")
+        fs.delete(tempOutputPath, false)
+        throw new IOException("Checkpoint failed: failed to save output of task: " +
+          s"${ctx.attemptNumber()} and final output path does not exist")
+      } else {
+        // Some other copy of this task must've finished before us and renamed it
+        logInfo(s"Final output path $finalOutputPath already exists; not overwriting it")
+        fs.delete(tempOutputPath, false)
+      }
+    }
+  }
+
+  def readFromFile[T](
+      path: Path,
+      broadcastedConf: Broadcast[SerializableConfiguration],
+      context: TaskContext): Iterator[T] = {
+    val env = SparkEnv.get
+    val fs = path.getFileSystem(broadcastedConf.value.value)
+    val bufferSize = env.conf.getInt("spark.buffer.size", 65536)
+    val fileInputStream = fs.open(path, bufferSize)
+    val serializer = env.serializer.newInstance()
+    val deserializeStream = serializer.deserializeStream(fileInputStream)
+
+    // Register an on-task-completion callback to close the input stream.
+    context.addTaskCompletionListener(context => deserializeStream.close())
+
+    deserializeStream.asIterator.asInstanceOf[Iterator[T]]
+  }
+
+  // Test whether CheckpointRDD generate expected number of partitions despite
+  // each split file having multiple blocks. This needs to be run on a
+  // cluster (mesos or standalone) using HDFS.
+  def main(args: Array[String]) {
+    val Array(cluster, hdfsPath) = args
+    val sc = new SparkContext(cluster, "CheckpointRDD Test")
+    val rdd = sc.makeRDD(1 to 10, 10).flatMap(x => 1 to 10000)
+    val path = new Path(hdfsPath, "temp")
+    val conf = SparkHadoopUtil.get.newConfiguration(new SparkConf())
+    val fs = path.getFileSystem(conf)
+    val broadcastedConf = sc.broadcast(new SerializableConfiguration(conf))
+    sc.runJob(rdd, ReliableCheckpointRDD.writeToFile[Int](path.toString, broadcastedConf, 1024) _)
+    val cpRDD = new ReliableCheckpointRDD[Int](sc, path.toString)
+    assert(cpRDD.partitions.length == rdd.partitions.length, "Number of partitions is not the same")
+    assert(cpRDD.collect().toList == rdd.collect().toList, "Data of partitions not the same")
+    fs.delete(path, true)
+  }
+
+}
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
new file mode 100644
index 0000000000000..9d2c889eff920
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
@@ -0,0 +1,109 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import scala.reflect.ClassTag
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark._
+import org.apache.spark.util.SerializableConfiguration
+
+/**
+ * An implementation of checkpointing that writes the RDD data to reliable storage.
+ * This allows drivers to be restarted on failure with previously computed state.
+ */
+private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
+  extends RDDCheckpointData[T](rdd) with Logging with Serializable {
+
+  // The directory to which the associated RDD has been checkpointed to
+  // This is assumed to be a non-local path that points to some reliable storage
+  private val cpDir: String =
+    ReliableRDDCheckpointData.rddCheckpointDataPath(rdd.context, rdd.id)
+      .map(_.toString)
+      .getOrElse { throw new SparkException("Checkpoint dir must be specified.") }
+
+  /**
+   * Return the directory to which this RDD was checkpointed.
+   * If the RDD is not checkpointed yet, return None.
+   */
+  def getCheckpointDir: Option[String] = RDDCheckpointData.synchronized {
+    if (isCheckpointed) {
+      Some(cpDir.toString)
+    } else {
+      None
+    }
+  }
+
+  /**
+   * Materialize this RDD and write its content to a reliable DFS.
+   * This is called immediately after the first action invoked on this RDD has completed.
+   */
+  protected override def doCheckpoint(): CheckpointRDD[T] = {
+
+    // Create the output path for the checkpoint
+    val path = new Path(cpDir)
+    val fs = path.getFileSystem(rdd.context.hadoopConfiguration)
+    if (!fs.mkdirs(path)) {
+      throw new SparkException(s"Failed to create checkpoint path $cpDir")
+    }
+
+    // Save to file, and reload it as an RDD
+    val broadcastedConf = rdd.context.broadcast(
+      new SerializableConfiguration(rdd.context.hadoopConfiguration))
+    val newRDD = new ReliableCheckpointRDD[T](rdd.context, cpDir)
+
+    // Optionally clean our checkpoint files if the reference is out of scope
+    if (rdd.conf.getBoolean("spark.cleaner.referenceTracking.cleanCheckpoints", false)) {
+      rdd.context.cleaner.foreach { cleaner =>
+        cleaner.registerRDDCheckpointDataForCleanup(newRDD, rdd.id)
+      }
+    }
+
+    // TODO: This is expensive because it computes the RDD again unnecessarily (SPARK-8582)
+    rdd.context.runJob(rdd, ReliableCheckpointRDD.writeToFile[T](cpDir, broadcastedConf) _)
+    if (newRDD.partitions.length != rdd.partitions.length) {
+      throw new SparkException(
+        s"Checkpoint RDD $newRDD(${newRDD.partitions.length}) has different " +
+          s"number of partitions from original RDD $rdd(${rdd.partitions.length})")
+    }
+
+    logInfo(s"Done checkpointing RDD ${rdd.id} to $cpDir, new parent is RDD ${newRDD.id}")
+
+    newRDD
+  }
+
+}
+
+private[spark] object ReliableRDDCheckpointData {
+
+  /** Return the path of the directory to which this RDD's checkpoint data is written. */
+  def rddCheckpointDataPath(sc: SparkContext, rddId: Int): Option[Path] = {
+    sc.checkpointDir.map { dir => new Path(dir, s"rdd-$rddId") }
+  }
+
+  /** Clean up the files associated with the checkpoint data for this RDD. */
+  def clearRDDCheckpointData(sc: SparkContext, rddId: Int): Unit = {
+    rddCheckpointDataPath(sc, rddId).foreach { path =>
+      val fs = path.getFileSystem(sc.hadoopConfiguration)
+      if (fs.exists(path)) {
+        fs.delete(path, true)
+      }
+    }
+  }
+}
diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
index 501fe186bfd7c..40a2efc56d0c8 100644
--- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
@@ -29,7 +29,7 @@ import org.scalatest.concurrent.Eventually._
 import org.scalatest.time.SpanSugar._
 
 import org.apache.spark.SparkContext._
-import org.apache.spark.rdd.{RDDCheckpointData, RDD}
+import org.apache.spark.rdd.{ReliableRDDCheckpointData, RDD}
 import org.apache.spark.storage._
 import org.apache.spark.shuffle.hash.HashShuffleManager
 import org.apache.spark.shuffle.sort.SortShuffleManager
@@ -221,8 +221,8 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     var rddId = rdd.id
 
     // Confirm the checkpoint directory exists
-    assert(RDDCheckpointData.rddCheckpointDataPath(sc, rddId).isDefined)
-    val path = RDDCheckpointData.rddCheckpointDataPath(sc, rddId).get
+    assert(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).isDefined)
+    val path = ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get
     val fs = path.getFileSystem(sc.hadoopConfiguration)
     assert(fs.exists(path))
 
@@ -231,7 +231,7 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rdd = null // Make RDD out of scope, ok if collected earlier
     runGC()
     postGCTester.assertCleanup()
-    assert(fs.exists(RDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
+    assert(fs.exists(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
 
     sc.stop()
     val conf = new SparkConf().setMaster("local[2]").setAppName("cleanupCheckpoint").
@@ -245,7 +245,7 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rddId = rdd.id
 
     // Confirm the checkpoint directory exists
-    assert(fs.exists(RDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
+    assert(fs.exists(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
 
     // Reference rdd to defeat any early collection by the JVM
     rdd.count()
@@ -255,7 +255,7 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rdd = null // Make RDD out of scope
     runGC()
     postGCTester.assertCleanup()
-    assert(!fs.exists(RDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
+    assert(!fs.exists(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
   }
 
   test("automatically cleanup RDD + shuffle + broadcast") {

From 84474547538d42a95ed15eb0915ae1760d975b4a Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Mon, 6 Jul 2015 18:07:16 -0700
Subject: [PATCH 02/27] Fix tests

The parent base class was not serializable while the child is,
causing some java invocation exception. Also, a prior code clean
up caused an array out of bounds exception, which is now fixed
in this commit.
---
 .../org/apache/spark/rdd/CheckpointRDD.scala  |  2 +-
 .../apache/spark/rdd/RDDCheckpointData.scala  |  3 +-
 .../spark/rdd/ReliableCheckpointRDD.scala     | 34 +++++++++++--------
 .../spark/rdd/ReliableRDDCheckpointData.scala |  2 +-
 4 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
index 6681bdd14940e..68a32ada971af 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
@@ -26,7 +26,7 @@ private[spark] class CheckpointRDDPartition(val index: Int) extends Partition
 /**
  * This RDD represents a RDD checkpoint file (similar to HadoopRDD).
  */
-private[spark] abstract class CheckpointRDD[T: ClassTag](sc: SparkContext)
+private[spark] abstract class CheckpointRDD[T: ClassTag](@transient sc: SparkContext)
   extends RDD[T](sc, Nil) {
 
   // CheckpointRDD should not be checkpointed again
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
index 5df420cce1c65..30bcdd95d8a47 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
@@ -36,7 +36,8 @@ private[spark] object CheckpointState extends Enumeration {
  * as well as, manages the post-checkpoint state by providing the updated partitions,
  * iterator and preferred locations of the checkpointed RDD.
  */
-private[spark] abstract class RDDCheckpointData[T: ClassTag](@transient rdd: RDD[T]) {
+private[spark] abstract class RDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
+  extends Serializable {
 
   import CheckpointState._
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index ac2930d015df6..006a25103b013 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -32,7 +32,7 @@ import org.apache.spark.util.{SerializableConfiguration, Utils}
  * An RDD that reads from a checkpoint file previously written to reliable storage.
  */
 private[spark] class ReliableCheckpointRDD[T: ClassTag](
-    sc: SparkContext,
+    @transient sc: SparkContext,
     val checkpointPath: String)
   extends CheckpointRDD[T](sc) {
 
@@ -47,21 +47,12 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
     val numPartitions =
       // listStatus can throw exception if path does not exist.
       if (fs.exists(cpath)) {
-        val dirContents = fs.listStatus(cpath).map(_.getPath)
-        val partitionFiles = dirContents
+        val inputFiles = fs.listStatus(cpath)
+          .map(_.getPath)
           .filter(_.getName.startsWith("part-"))
-          .map(_.toString)
-          .sorted
-        // Verify validity of checkpoint files
-        val numPart = partitionFiles.length
-        val firstPartitionValid =
-          partitionFiles(0).endsWith(ReliableCheckpointRDD.splitIdToFile(0))
-        val lastPartitionValid =
-          partitionFiles(numPart - 1).endsWith(ReliableCheckpointRDD.splitIdToFile(numPart - 1))
-        if (numPart > 0 && (!firstPartitionValid || !lastPartitionValid)) {
-          throw new SparkException(s"Invalid checkpoint directory: $checkpointPath")
-        }
-        numPart
+          .sortBy(_.toString)
+        validateInputFiles(inputFiles)
+        inputFiles.length
       } else {
         0
       }
@@ -79,6 +70,19 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
     val file = new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index))
     ReliableCheckpointRDD.readFromFile(file, broadcastedConf, context)
   }
+
+  /**
+   * Fail fast if our input checkpointed files are invalid.
+   */
+  private def validateInputFiles(files: Array[Path]): Unit = {
+    val valid = files.zipWithIndex.forall { case (f, i) =>
+      f.toString.endsWith(ReliableCheckpointRDD.splitIdToFile(i))
+    }
+    if (!valid) {
+      throw new SparkException(s"Invalid checkpoint directory: $checkpointPath")
+    }
+  }
+
 }
 
 private[spark] object ReliableCheckpointRDD extends Logging {
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
index 9d2c889eff920..30caef1bd4f86 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
@@ -29,7 +29,7 @@ import org.apache.spark.util.SerializableConfiguration
  * This allows drivers to be restarted on failure with previously computed state.
  */
 private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
-  extends RDDCheckpointData[T](rdd) with Logging with Serializable {
+  extends RDDCheckpointData[T](rdd) with Logging {
 
   // The directory to which the associated RDD has been checkpointed to
   // This is assumed to be a non-local path that points to some reliable storage

From 2e902e53a641242b74a2414cd65d403e2273a32d Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 11:02:47 -0700
Subject: [PATCH 03/27] First implementation of local checkpointing

The write path runs a job to put each partition into disk store,
while the read path simply reads these blocks back from the disk
store.
---
 .../apache/spark/rdd/LocalCheckpointRDD.scala | 96 +++++++++++++++++++
 .../spark/rdd/LocalRDDCheckpointData.scala    | 52 ++++++++++
 .../main/scala/org/apache/spark/rdd/RDD.scala |  5 +-
 .../org/apache/spark/storage/BlockId.scala    |  8 ++
 4 files changed, 159 insertions(+), 2 deletions(-)
 create mode 100644 core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
 create mode 100644 core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
new file mode 100644
index 0000000000000..8b417cb64a68a
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -0,0 +1,96 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import scala.reflect.ClassTag
+
+import org.apache.spark.{SparkException, Partition, SparkEnv, TaskContext}
+import org.apache.spark.storage.{LocalCheckpointBlockId, BlockId}
+
+/**
+ * An RDD that reads from a checkpoint file previously written to the local file system.
+ */
+private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
+  extends CheckpointRDD[T](rdd.context) {
+
+  // This is needed because `rdd` is transient and will not be found on the executors
+  private val rddId = rdd.id
+
+  /**
+   * Determine the partitions from the local checkpoint blocks on each executor.
+   */
+  override def getPartitions: Array[Partition] = {
+    val expectedRddId = rddId
+    val blockFilter = (blockId: BlockId) => blockId match {
+      case LocalCheckpointBlockId(rddId, _) => rddId == expectedRddId
+      case _ => false
+    }
+    val inputPartitions: Array[Partition] =
+      SparkEnv.get.blockManager.master
+        .getMatchingBlockIds(blockFilter, askSlaves = true)
+        .collect { case LocalCheckpointBlockId(_, i) => new CheckpointRDDPartition(i) }
+        .toArray
+    validateInputPartitions(inputPartitions)
+    inputPartitions
+  }
+
+  /**
+   * Return the location of the checkpoint block that corresponds to the given partition.
+   */
+  override def getPreferredLocations(partition: Partition): Seq[String] = {
+    val index = partition.asInstanceOf[CheckpointRDDPartition].index
+    val blockId = new LocalCheckpointBlockId(rddId, index)
+    val hosts = SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
+    if (hosts.size != 1) {
+      // We do not replicate the block, so it should be found on exactly one executor
+      logWarning("Expected exactly one host for each local checkpoint block. Instead: " + hosts)
+    }
+    hosts
+  }
+
+  /**
+   * Fetch the local checkpoint block that corresponds to this partition.
+   * This block should be in the disk store of at least one executor.
+   */
+  override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
+    val index = partition.asInstanceOf[CheckpointRDDPartition].index
+    val blockId = new LocalCheckpointBlockId(rddId, index)
+    SparkEnv.get.blockManager.get(blockId) match {
+      case Some(result) =>
+        result.data.asInstanceOf[Iterator[T]]
+      case None => throw new SparkException(s"Checkpoint block $blockId not found.")
+    }
+  }
+
+  /**
+   * Validate that the indices of the input partitions are continuous.
+   */
+  private def validateInputPartitions(partitions: Array[Partition]): Unit = {
+    val sortedIndices = partitions.map(_.index).sorted
+    if (sortedIndices.nonEmpty) {
+      val expectedIndices = (sortedIndices.head to sortedIndices.last).toArray
+      if (!java.util.Arrays.equals(sortedIndices, expectedIndices)) {
+        throw new SparkException(
+          "Local checkpoint partitions are invalid.\n" +
+            s"  Expected indices: ${expectedIndices.mkString(", ")}\n" +
+            s"  Actual indices: ${sortedIndices.mkString(", ")}")
+      }
+    }
+  }
+
+}
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
new file mode 100644
index 0000000000000..ae6aea905e2be
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -0,0 +1,52 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import scala.reflect.ClassTag
+
+import org.apache.spark.{Logging, SparkEnv, TaskContext}
+import org.apache.spark.storage.{LocalCheckpointBlockId, StorageLevel}
+
+/**
+ * An implementation of checkpointing that writes the RDD data to a local file system.
+ *
+ * Local checkpointing trades off fault tolerance for performance by skipping the expensive
+ * step of replicating the checkpointed data in a reliable storage. This is useful for use
+ * cases where RDDs build up long lineages that need to be truncated often (e.g. GraphX).
+ */
+private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
+  extends RDDCheckpointData[T](rdd) with Logging {
+
+  protected override def doCheckpoint(): CheckpointRDD[T] = {
+
+    // Put each partition into the disk store
+    // TODO: if it's already in disk store, just use the existing values
+    val rddId = rdd.id
+    val persistPartition = (taskContext: TaskContext, values: Iterator[T]) => {
+      SparkEnv.get.blockManager.putIterator(
+        LocalCheckpointBlockId(rddId, taskContext.partitionId()),
+        values,
+        StorageLevel.DISK_ONLY)
+    }
+    rdd.context.runJob(rdd, persistPartition)
+
+    // Return an RDD that reads these blocks back
+    new LocalCheckpointRDD[T](rdd)
+  }
+
+}
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index ed3da0a38c0f3..8f07a57e91652 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1479,8 +1479,9 @@ abstract class RDD[T: ClassTag](
    * The actual persisting occurs after the first job involving this RDD has completed.
    * The checkpoint directory set through `SparkContext#setCheckpointDir` is not read.
    */
-  def localCheckpoint(): Unit = {
-
+  def localCheckpoint(): this.type = RDDCheckpointData.synchronized {
+    checkpointData = Some(new LocalRDDCheckpointData(this))
+    this
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockId.scala b/core/src/main/scala/org/apache/spark/storage/BlockId.scala
index 524f6970992a5..0c505129eeb8b 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockId.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockId.scala
@@ -95,6 +95,11 @@ private[spark] case class TempShuffleBlockId(id: UUID) extends BlockId {
   override def name: String = "temp_shuffle_" + id
 }
 
+/** Id associated with locally checkpointed data. */
+private[spark] case class LocalCheckpointBlockId(rddId: Int, partitionIndex: Int) extends BlockId {
+  override def name: String = "local_checkpoint_" + rddId + "_" + partitionIndex
+}
+
 // Intended only for testing purposes
 private[spark] case class TestBlockId(id: String) extends BlockId {
   override def name: String = "test_" + id
@@ -109,6 +114,7 @@ object BlockId {
   val BROADCAST = "broadcast_([0-9]+)([_A-Za-z0-9]*)".r
   val TASKRESULT = "taskresult_([0-9]+)".r
   val STREAM = "input-([0-9]+)-([0-9]+)".r
+  val LOCAL_CHECKPOINT = "local_checkpoint_([0-9]+)_([0-9]+)".r
   val TEST = "test_(.*)".r
 
   /** Converts a BlockId "name" String back into a BlockId. */
@@ -127,6 +133,8 @@ object BlockId {
       TaskResultBlockId(taskId.toLong)
     case STREAM(streamId, uniqueId) =>
       StreamBlockId(streamId.toInt, uniqueId.toLong)
+    case LOCAL_CHECKPOINT(rddId, partitionIndex) =>
+      LocalCheckpointBlockId(rddId.toInt, partitionIndex.toInt)
     case TEST(value) =>
       TestBlockId(value)
     case _ =>

From 0477eec3f27ad37650097661a727c50c740c72c6 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 15:01:29 -0700
Subject: [PATCH 04/27] Rename a few methods with awkward names (minor)

---
 .../main/scala/org/apache/spark/ContextCleaner.scala   |  2 +-
 .../org/apache/spark/rdd/LocalCheckpointRDD.scala      | 10 +++-------
 .../apache/spark/rdd/ReliableRDDCheckpointData.scala   |  8 ++++----
 .../scala/org/apache/spark/ContextCleanerSuite.scala   | 10 +++++-----
 4 files changed, 13 insertions(+), 17 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
index 825081eef2443..1cc39f9966058 100644
--- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala
+++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
@@ -240,7 +240,7 @@ private[spark] class ContextCleaner(sc: SparkContext) extends Logging {
   def doCleanCheckpoint(rddId: Int): Unit = {
     try {
       logDebug("Cleaning rdd checkpoint data " + rddId)
-      ReliableRDDCheckpointData.clearRDDCheckpointData(sc, rddId)
+      ReliableRDDCheckpointData.cleanCheckpoint(sc, rddId)
       listeners.foreach(_.checkpointCleaned(rddId))
       logInfo("Cleaned rdd checkpoint data " + rddId)
     }
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 8b417cb64a68a..93555991b901f 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -28,16 +28,12 @@ import org.apache.spark.storage.{LocalCheckpointBlockId, BlockId}
 private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
   extends CheckpointRDD[T](rdd.context) {
 
-  // This is needed because `rdd` is transient and will not be found on the executors
-  private val rddId = rdd.id
-
   /**
    * Determine the partitions from the local checkpoint blocks on each executor.
    */
   override def getPartitions: Array[Partition] = {
-    val expectedRddId = rddId
     val blockFilter = (blockId: BlockId) => blockId match {
-      case LocalCheckpointBlockId(rddId, _) => rddId == expectedRddId
+      case LocalCheckpointBlockId(x, _) => x == id
       case _ => false
     }
     val inputPartitions: Array[Partition] =
@@ -54,7 +50,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
    */
   override def getPreferredLocations(partition: Partition): Seq[String] = {
     val index = partition.asInstanceOf[CheckpointRDDPartition].index
-    val blockId = new LocalCheckpointBlockId(rddId, index)
+    val blockId = new LocalCheckpointBlockId(id, index)
     val hosts = SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
     if (hosts.size != 1) {
       // We do not replicate the block, so it should be found on exactly one executor
@@ -69,7 +65,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
     val index = partition.asInstanceOf[CheckpointRDDPartition].index
-    val blockId = new LocalCheckpointBlockId(rddId, index)
+    val blockId = new LocalCheckpointBlockId(id, index)
     SparkEnv.get.blockManager.get(blockId) match {
       case Some(result) =>
         result.data.asInstanceOf[Iterator[T]]
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
index 30caef1bd4f86..7387875d1c743 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
@@ -34,7 +34,7 @@ private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[
   // The directory to which the associated RDD has been checkpointed to
   // This is assumed to be a non-local path that points to some reliable storage
   private val cpDir: String =
-    ReliableRDDCheckpointData.rddCheckpointDataPath(rdd.context, rdd.id)
+    ReliableRDDCheckpointData.checkpointPath(rdd.context, rdd.id)
       .map(_.toString)
       .getOrElse { throw new SparkException("Checkpoint dir must be specified.") }
 
@@ -93,13 +93,13 @@ private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[
 private[spark] object ReliableRDDCheckpointData {
 
   /** Return the path of the directory to which this RDD's checkpoint data is written. */
-  def rddCheckpointDataPath(sc: SparkContext, rddId: Int): Option[Path] = {
+  def checkpointPath(sc: SparkContext, rddId: Int): Option[Path] = {
     sc.checkpointDir.map { dir => new Path(dir, s"rdd-$rddId") }
   }
 
   /** Clean up the files associated with the checkpoint data for this RDD. */
-  def clearRDDCheckpointData(sc: SparkContext, rddId: Int): Unit = {
-    rddCheckpointDataPath(sc, rddId).foreach { path =>
+  def cleanCheckpoint(sc: SparkContext, rddId: Int): Unit = {
+    checkpointPath(sc, rddId).foreach { path =>
       val fs = path.getFileSystem(sc.hadoopConfiguration)
       if (fs.exists(path)) {
         fs.delete(path, true)
diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
index 40a2efc56d0c8..05c3d1245d41a 100644
--- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
@@ -221,8 +221,8 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     var rddId = rdd.id
 
     // Confirm the checkpoint directory exists
-    assert(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).isDefined)
-    val path = ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get
+    assert(ReliableRDDCheckpointData.checkpointPath(sc, rddId).isDefined)
+    val path = ReliableRDDCheckpointData.checkpointPath(sc, rddId).get
     val fs = path.getFileSystem(sc.hadoopConfiguration)
     assert(fs.exists(path))
 
@@ -231,7 +231,7 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rdd = null // Make RDD out of scope, ok if collected earlier
     runGC()
     postGCTester.assertCleanup()
-    assert(fs.exists(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
+    assert(fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
 
     sc.stop()
     val conf = new SparkConf().setMaster("local[2]").setAppName("cleanupCheckpoint").
@@ -245,7 +245,7 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rddId = rdd.id
 
     // Confirm the checkpoint directory exists
-    assert(fs.exists(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
+    assert(fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
 
     // Reference rdd to defeat any early collection by the JVM
     rdd.count()
@@ -255,7 +255,7 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rdd = null // Make RDD out of scope
     runGC()
     postGCTester.assertCleanup()
-    assert(!fs.exists(ReliableRDDCheckpointData.rddCheckpointDataPath(sc, rddId).get))
+    assert(!fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
   }
 
   test("automatically cleanup RDD + shuffle + broadcast") {

From 4514dc9a1b743d04a1ffd3fab741312f601eb6f8 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 15:30:55 -0700
Subject: [PATCH 05/27] Clean local checkpoint files through RDD cleanups

This commit simplifies the previous one in removing the special
`LocalCheckpointBlockId`, which is not needed if we use the
checkpoint RDD's ID instead of the parent RDD's. This allows us
to simply reuse the RDD cleanup code path, which is nice.
---
 .../org/apache/spark/ContextCleaner.scala     |  8 ++---
 .../apache/spark/rdd/LocalCheckpointRDD.scala | 27 +++++++---------
 .../spark/rdd/LocalRDDCheckpointData.scala    | 32 ++++++++++++-------
 .../org/apache/spark/storage/BlockId.scala    |  8 -----
 4 files changed, 35 insertions(+), 40 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
index 1cc39f9966058..fabf5a619b1ac 100644
--- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala
+++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
@@ -22,7 +22,7 @@ import java.lang.ref.{ReferenceQueue, WeakReference}
 import scala.collection.mutable.{ArrayBuffer, SynchronizedBuffer}
 
 import org.apache.spark.broadcast.Broadcast
-import org.apache.spark.rdd.{ReliableRDDCheckpointData, RDD}
+import org.apache.spark.rdd.{RDD, ReliableRDDCheckpointData}
 import org.apache.spark.util.Utils
 
 /**
@@ -232,10 +232,8 @@ private[spark] class ContextCleaner(sc: SparkContext) extends Logging {
   }
 
   /**
-   * Perform checkpoint cleanup.
-   *
-   * Note that this only cleans up reliable checkpoint files.
-   * Unsafe checkpoint files are just cached RDD blocks and are cleaned up separately.
+   * Perform checkpoint cleanup on checkpoint files in reliable storage.
+   * Locally checkpointed files are cleaned up separately through RDD clean ups.
    */
   def doCleanCheckpoint(rddId: Int): Unit = {
     try {
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 93555991b901f..5aaade36d4822 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -19,27 +19,27 @@ package org.apache.spark.rdd
 
 import scala.reflect.ClassTag
 
-import org.apache.spark.{SparkException, Partition, SparkEnv, TaskContext}
-import org.apache.spark.storage.{LocalCheckpointBlockId, BlockId}
+import org.apache.spark.{Partition, SparkEnv, SparkException, SparkContext, TaskContext}
+import org.apache.spark.storage.{BlockId, RDDBlockId}
 
 /**
  * An RDD that reads from a checkpoint file previously written to the local file system.
  */
-private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
-  extends CheckpointRDD[T](rdd.context) {
+private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext)
+  extends CheckpointRDD[T](sc) {
 
   /**
    * Determine the partitions from the local checkpoint blocks on each executor.
    */
   override def getPartitions: Array[Partition] = {
-    val blockFilter = (blockId: BlockId) => blockId match {
-      case LocalCheckpointBlockId(x, _) => x == id
-      case _ => false
+    val ourId = id // define this locally for serialization purposes
+    val blockFilter = (blockId: BlockId) => {
+      blockId.asRDDId.filter(_.rddId == ourId).isDefined
     }
     val inputPartitions: Array[Partition] =
       SparkEnv.get.blockManager.master
         .getMatchingBlockIds(blockFilter, askSlaves = true)
-        .collect { case LocalCheckpointBlockId(_, i) => new CheckpointRDDPartition(i) }
+        .collect { case RDDBlockId(_, i) => new CheckpointRDDPartition(i) }
         .toArray
     validateInputPartitions(inputPartitions)
     inputPartitions
@@ -50,13 +50,8 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
    */
   override def getPreferredLocations(partition: Partition): Seq[String] = {
     val index = partition.asInstanceOf[CheckpointRDDPartition].index
-    val blockId = new LocalCheckpointBlockId(id, index)
-    val hosts = SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
-    if (hosts.size != 1) {
-      // We do not replicate the block, so it should be found on exactly one executor
-      logWarning("Expected exactly one host for each local checkpoint block. Instead: " + hosts)
-    }
-    hosts
+    val blockId = RDDBlockId(id, index)
+    SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
   }
 
   /**
@@ -65,7 +60,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient rdd: RDD[T])
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
     val index = partition.asInstanceOf[CheckpointRDDPartition].index
-    val blockId = new LocalCheckpointBlockId(id, index)
+    val blockId = RDDBlockId(id, index)
     SparkEnv.get.blockManager.get(blockId) match {
       case Some(result) =>
         result.data.asInstanceOf[Iterator[T]]
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index ae6aea905e2be..8fbb9fbf46c89 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -20,7 +20,7 @@ package org.apache.spark.rdd
 import scala.reflect.ClassTag
 
 import org.apache.spark.{Logging, SparkEnv, TaskContext}
-import org.apache.spark.storage.{LocalCheckpointBlockId, StorageLevel}
+import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 
 /**
  * An implementation of checkpointing that writes the RDD data to a local file system.
@@ -32,21 +32,31 @@ import org.apache.spark.storage.{LocalCheckpointBlockId, StorageLevel}
 private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
   extends RDDCheckpointData[T](rdd) with Logging {
 
+  /**
+   * Write the content of each partition on a local disk store.
+   */
   protected override def doCheckpoint(): CheckpointRDD[T] = {
+    val checkpointRdd = new LocalCheckpointRDD[T](rdd.context)
 
-    // Put each partition into the disk store
-    // TODO: if it's already in disk store, just use the existing values
-    val rddId = rdd.id
+    // Note: we persist the partitions using the checkpoint RDD's ID rather than the
+    // parent RDD's ID. This is not fundamental in design, but allows us to simply
+    // reuse the RDD clean up code path to clean the checkpointed files.
+    val checkpointRddId = checkpointRdd.id
     val persistPartition = (taskContext: TaskContext, values: Iterator[T]) => {
-      SparkEnv.get.blockManager.putIterator(
-        LocalCheckpointBlockId(rddId, taskContext.partitionId()),
-        values,
-        StorageLevel.DISK_ONLY)
+      // TODO: if it's already in disk store, just use the existing values
+      val blockId = RDDBlockId(checkpointRddId, taskContext.partitionId())
+      SparkEnv.get.blockManager.putIterator(blockId, values, StorageLevel.DISK_ONLY)
+    }
+
+    // Optionally clean our checkpoint files if the reference is out of scope
+    if (rdd.conf.getBoolean("spark.cleaner.referenceTracking.cleanCheckpoints", false)) {
+      rdd.context.cleaner.foreach { cleaner =>
+        cleaner.registerRDDForCleanup(checkpointRdd)
+      }
     }
-    rdd.context.runJob(rdd, persistPartition)
 
-    // Return an RDD that reads these blocks back
-    new LocalCheckpointRDD[T](rdd)
+    rdd.context.runJob(rdd, persistPartition)
+    checkpointRdd
   }
 
 }
diff --git a/core/src/main/scala/org/apache/spark/storage/BlockId.scala b/core/src/main/scala/org/apache/spark/storage/BlockId.scala
index 0c505129eeb8b..524f6970992a5 100644
--- a/core/src/main/scala/org/apache/spark/storage/BlockId.scala
+++ b/core/src/main/scala/org/apache/spark/storage/BlockId.scala
@@ -95,11 +95,6 @@ private[spark] case class TempShuffleBlockId(id: UUID) extends BlockId {
   override def name: String = "temp_shuffle_" + id
 }
 
-/** Id associated with locally checkpointed data. */
-private[spark] case class LocalCheckpointBlockId(rddId: Int, partitionIndex: Int) extends BlockId {
-  override def name: String = "local_checkpoint_" + rddId + "_" + partitionIndex
-}
-
 // Intended only for testing purposes
 private[spark] case class TestBlockId(id: String) extends BlockId {
   override def name: String = "test_" + id
@@ -114,7 +109,6 @@ object BlockId {
   val BROADCAST = "broadcast_([0-9]+)([_A-Za-z0-9]*)".r
   val TASKRESULT = "taskresult_([0-9]+)".r
   val STREAM = "input-([0-9]+)-([0-9]+)".r
-  val LOCAL_CHECKPOINT = "local_checkpoint_([0-9]+)_([0-9]+)".r
   val TEST = "test_(.*)".r
 
   /** Converts a BlockId "name" String back into a BlockId. */
@@ -133,8 +127,6 @@ object BlockId {
       TaskResultBlockId(taskId.toLong)
     case STREAM(streamId, uniqueId) =>
       StreamBlockId(streamId.toInt, uniqueId.toLong)
-    case LOCAL_CHECKPOINT(rddId, partitionIndex) =>
-      LocalCheckpointBlockId(rddId.toInt, partitionIndex.toInt)
     case TEST(value) =>
       TestBlockId(value)
     case _ =>

From 4dbbab19031b4e05caf787c5317f6294b79297a1 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 16:51:13 -0700
Subject: [PATCH 06/27] Refactor CheckpointSuite to test local checkpointing

This commit makes each test in CheckpointSuite run twice, once
for normal checkpointing and another time for local checkpointing.
This commit also fixes legitimate test failures after the
refactoring.
---
 .../apache/spark/rdd/LocalCheckpointRDD.scala |   7 +-
 .../org/apache/spark/CheckpointSuite.scala    | 162 +++++++++++-------
 2 files changed, 102 insertions(+), 67 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 5aaade36d4822..755c81cd90eef 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -40,6 +40,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
       SparkEnv.get.blockManager.master
         .getMatchingBlockIds(blockFilter, askSlaves = true)
         .collect { case RDDBlockId(_, i) => new CheckpointRDDPartition(i) }
+        .sortBy(_.index)
         .toArray
     validateInputPartitions(inputPartitions)
     inputPartitions
@@ -49,8 +50,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
    * Return the location of the checkpoint block that corresponds to the given partition.
    */
   override def getPreferredLocations(partition: Partition): Seq[String] = {
-    val index = partition.asInstanceOf[CheckpointRDDPartition].index
-    val blockId = RDDBlockId(id, index)
+    val blockId = RDDBlockId(id, partition.index)
     SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
   }
 
@@ -59,8 +59,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
    * This block should be in the disk store of at least one executor.
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
-    val index = partition.asInstanceOf[CheckpointRDDPartition].index
-    val blockId = RDDBlockId(id, index)
+    val blockId = RDDBlockId(id, partition.index)
     SparkEnv.get.blockManager.get(blockId) match {
       case Some(result) =>
         result.data.asInstanceOf[Iterator[T]]
diff --git a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
index cc50e6d79a3e2..c6bbfe8725bcd 100644
--- a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
@@ -25,11 +25,15 @@ import org.apache.spark.rdd._
 import org.apache.spark.storage.{BlockId, StorageLevel, TestBlockId}
 import org.apache.spark.util.Utils
 
+/**
+ * Test suite for end-to-end checkpointing functionality.
+ * This tests both normal checkpoints and local checkpoints.
+ */
 class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging {
-  var checkpointDir: File = _
-  val partitioner = new HashPartitioner(2)
+  private var checkpointDir: File = _
+  private val partitioner = new HashPartitioner(2)
 
-  override def beforeEach() {
+  override def beforeEach(): Unit = {
     super.beforeEach()
     checkpointDir = File.createTempFile("temp", "", Utils.createTempDir())
     checkpointDir.delete()
@@ -37,40 +41,42 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     sc.setCheckpointDir(checkpointDir.toString)
   }
 
-  override def afterEach() {
+  override def afterEach(): Unit = {
     super.afterEach()
     Utils.deleteRecursively(checkpointDir)
   }
 
-  test("basic checkpointing") {
+  runTest("basic checkpointing") { normal: Boolean =>
     val parCollection = sc.makeRDD(1 to 4)
     val flatMappedRDD = parCollection.flatMap(x => 1 to x)
-    flatMappedRDD.checkpoint()
+    checkpoint(flatMappedRDD, normal)
     assert(flatMappedRDD.dependencies.head.rdd === parCollection)
     val result = flatMappedRDD.collect()
     assert(flatMappedRDD.dependencies.head.rdd != parCollection)
     assert(flatMappedRDD.collect() === result)
   }
 
-  test("RDDs with one-to-one dependencies") {
-    testRDD(_.map(x => x.toString))
-    testRDD(_.flatMap(x => 1 to x))
-    testRDD(_.filter(_ % 2 == 0))
-    testRDD(_.sample(false, 0.5, 0))
-    testRDD(_.glom())
-    testRDD(_.mapPartitions(_.map(_.toString)))
-    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).mapValues(_.toString))
-    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).flatMapValues(x => 1 to x))
-    testRDD(_.pipe(Seq("cat")))
+  runTest("RDDs with one-to-one dependencies") { normal: Boolean =>
+    testRDD(_.map(x => x.toString), normal)
+    testRDD(_.flatMap(x => 1 to x), normal)
+    testRDD(_.filter(_ % 2 == 0), normal)
+    testRDD(_.sample(false, 0.5, 0), normal)
+    testRDD(_.glom(), normal)
+    testRDD(_.mapPartitions(_.map(_.toString)), normal)
+    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).mapValues(_.toString), normal)
+    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).flatMapValues(x => 1 to x), normal)
+    testRDD(_.pipe(Seq("cat")), normal)
   }
 
-  test("ParallelCollection") {
+  runTest("ParallelCollectionRDD") { normal: Boolean =>
     val parCollection = sc.makeRDD(1 to 4, 2)
     val numPartitions = parCollection.partitions.size
-    parCollection.checkpoint()
+    checkpoint(parCollection, normal)
     assert(parCollection.dependencies === Nil)
     val result = parCollection.collect()
-    assert(sc.checkpointFile[Int](parCollection.getCheckpointFile.get).collect() === result)
+    if (normal) {
+      assert(sc.checkpointFile[Int](parCollection.getCheckpointFile.get).collect() === result)
+    }
     assert(parCollection.dependencies != Nil)
     assert(parCollection.partitions.length === numPartitions)
     assert(parCollection.partitions.toList ===
@@ -78,44 +84,46 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     assert(parCollection.collect() === result)
   }
 
-  test("BlockRDD") {
+  runTest("BlockRDD") { normal: Boolean =>
     val blockId = TestBlockId("id")
     val blockManager = SparkEnv.get.blockManager
     blockManager.putSingle(blockId, "test", StorageLevel.MEMORY_ONLY)
     val blockRDD = new BlockRDD[String](sc, Array(blockId))
     val numPartitions = blockRDD.partitions.size
-    blockRDD.checkpoint()
+    checkpoint(blockRDD, normal)
     val result = blockRDD.collect()
-    assert(sc.checkpointFile[String](blockRDD.getCheckpointFile.get).collect() === result)
+    if (normal) {
+      assert(sc.checkpointFile[String](blockRDD.getCheckpointFile.get).collect() === result)
+    }
     assert(blockRDD.dependencies != Nil)
     assert(blockRDD.partitions.length === numPartitions)
     assert(blockRDD.partitions.toList === blockRDD.checkpointData.get.getPartitions.toList)
     assert(blockRDD.collect() === result)
   }
 
-  test("ShuffledRDD") {
+  runTest("ShuffleRDD") { normal: Boolean =>
     testRDD(rdd => {
       // Creating ShuffledRDD directly as PairRDDFunctions.combineByKey produces a MapPartitionedRDD
       new ShuffledRDD[Int, Int, Int](rdd.map(x => (x % 2, 1)), partitioner)
-    })
+    }, normal)
   }
 
-  test("UnionRDD") {
+  runTest("UnionRDD") { normal: Boolean =>
     def otherRDD: RDD[Int] = sc.makeRDD(1 to 10, 1)
-    testRDD(_.union(otherRDD))
-    testRDDPartitions(_.union(otherRDD))
+    testRDD(_.union(otherRDD), normal)
+    testRDDPartitions(_.union(otherRDD), normal)
   }
 
-  test("CartesianRDD") {
+  runTest("CartesianRDD") { normal: Boolean =>
     def otherRDD: RDD[Int] = sc.makeRDD(1 to 10, 1)
-    testRDD(new CartesianRDD(sc, _, otherRDD))
-    testRDDPartitions(new CartesianRDD(sc, _, otherRDD))
+    testRDD(new CartesianRDD(sc, _, otherRDD), normal)
+    testRDDPartitions(new CartesianRDD(sc, _, otherRDD), normal)
 
     // Test that the CartesianRDD updates parent partitions (CartesianRDD.s1/s2) after
     // the parent RDD has been checkpointed and parent partitions have been changed.
     // Note that this test is very specific to the current implementation of CartesianRDD.
     val ones = sc.makeRDD(1 to 100, 10).map(x => x)
-    ones.checkpoint() // checkpoint that MappedRDD
+    checkpoint(ones, normal) // checkpoint that MappedRDD
     val cartesian = new CartesianRDD(sc, ones, ones)
     val splitBeforeCheckpoint =
       serializeDeserialize(cartesian.partitions.head.asInstanceOf[CartesianPartition])
@@ -129,16 +137,16 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  test("CoalescedRDD") {
-    testRDD(_.coalesce(2))
-    testRDDPartitions(_.coalesce(2))
+  runTest("CoalescedRDD") { normal: Boolean =>
+    testRDD(_.coalesce(2), normal)
+    testRDDPartitions(_.coalesce(2), normal)
 
     // Test that the CoalescedRDDPartition updates parent partitions (CoalescedRDDPartition.parents)
     // after the parent RDD has been checkpointed and parent partitions have been changed.
     // Note that this test is very specific to the current implementation of
     // CoalescedRDDPartitions.
     val ones = sc.makeRDD(1 to 100, 10).map(x => x)
-    ones.checkpoint() // checkpoint that MappedRDD
+    checkpoint(ones, normal) // checkpoint that MappedRDD
     val coalesced = new CoalescedRDD(ones, 2)
     val splitBeforeCheckpoint =
       serializeDeserialize(coalesced.partitions.head.asInstanceOf[CoalescedRDDPartition])
@@ -151,7 +159,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  test("CoGroupedRDD") {
+  runTest("CoGroupedRDD") { normal: Boolean =>
     val longLineageRDD1 = generateFatPairRDD()
 
     // Collect the RDD as sequences instead of arrays to enable equality tests in testRDD
@@ -160,26 +168,26 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
 
     testRDD(rdd => {
       CheckpointSuite.cogroup(longLineageRDD1, rdd.map(x => (x % 2, 1)), partitioner)
-    }, seqCollectFunc)
+    }, normal, seqCollectFunc)
 
     val longLineageRDD2 = generateFatPairRDD()
     testRDDPartitions(rdd => {
       CheckpointSuite.cogroup(
         longLineageRDD2, sc.makeRDD(1 to 2, 2).map(x => (x % 2, 1)), partitioner)
-    }, seqCollectFunc)
+    }, normal, seqCollectFunc)
   }
 
-  test("ZippedPartitionsRDD") {
-    testRDD(rdd => rdd.zip(rdd.map(x => x)))
-    testRDDPartitions(rdd => rdd.zip(rdd.map(x => x)))
+  runTest("ZippedPartitionsRDD") { normal: Boolean =>
+    testRDD(rdd => rdd.zip(rdd.map(x => x)), normal)
+    testRDDPartitions(rdd => rdd.zip(rdd.map(x => x)), normal)
 
     // Test that ZippedPartitionsRDD updates parent partitions after parent RDDs have
     // been checkpointed and parent partitions have been changed.
     // Note that this test is very specific to the implementation of ZippedPartitionsRDD.
     val rdd = generateFatRDD()
     val zippedRDD = rdd.zip(rdd.map(x => x)).asInstanceOf[ZippedPartitionsRDD2[_, _, _]]
-    zippedRDD.rdd1.checkpoint()
-    zippedRDD.rdd2.checkpoint()
+    checkpoint(zippedRDD.rdd1, normal)
+    checkpoint(zippedRDD.rdd2, normal)
     val partitionBeforeCheckpoint =
       serializeDeserialize(zippedRDD.partitions.head.asInstanceOf[ZippedPartitionsPartition])
     zippedRDD.count()
@@ -194,27 +202,27 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  test("PartitionerAwareUnionRDD") {
+  runTest("PartitionerAwareUnionRDD") { normal: Boolean =>
     testRDD(rdd => {
       new PartitionerAwareUnionRDD[(Int, Int)](sc, Array(
         generateFatPairRDD(),
         rdd.map(x => (x % 2, 1)).reduceByKey(partitioner, _ + _)
       ))
-    })
+    }, normal)
 
     testRDDPartitions(rdd => {
       new PartitionerAwareUnionRDD[(Int, Int)](sc, Array(
         generateFatPairRDD(),
         rdd.map(x => (x % 2, 1)).reduceByKey(partitioner, _ + _)
       ))
-    })
+    }, normal)
 
     // Test that the PartitionerAwareUnionRDD updates parent partitions
     // (PartitionerAwareUnionRDD.parents) after the parent RDD has been checkpointed and parent
     // partitions have been changed. Note that this test is very specific to the current
     // implementation of PartitionerAwareUnionRDD.
     val pairRDD = generateFatPairRDD()
-    pairRDD.checkpoint()
+    checkpoint(pairRDD, normal)
     val unionRDD = new PartitionerAwareUnionRDD(sc, Array(pairRDD))
     val partitionBeforeCheckpoint = serializeDeserialize(
       unionRDD.partitions.head.asInstanceOf[PartitionerAwareUnionRDDPartition])
@@ -228,17 +236,34 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  test("CheckpointRDD with zero partitions") {
+  runTest("CheckpointRDD with zero partitions") { normal: Boolean =>
     val rdd = new BlockRDD[Int](sc, Array[BlockId]())
     assert(rdd.partitions.size === 0)
     assert(rdd.isCheckpointed === false)
-    rdd.checkpoint()
+    checkpoint(rdd, normal)
     assert(rdd.count() === 0)
     assert(rdd.isCheckpointed === true)
     assert(rdd.partitions.size === 0)
   }
 
-  def defaultCollectFunc[T](rdd: RDD[T]): Any = rdd.collect()
+  // Utility test methods
+
+  /** Checkpoint the RDD either locally or normally. */
+  private def checkpoint(rdd: RDD[_], normal: Boolean): Unit = {
+    if (normal) {
+      rdd.checkpoint()
+    } else {
+      rdd.localCheckpoint()
+    }
+  }
+
+  /** Run a test twice, once for local checkpointing and once for normal checkpointing. */
+  private def runTest(name: String)(body: Boolean => Unit): Unit = {
+    test(name + " [normal checkpointing]")(body(true))
+    test(name + " [local checkpointing]")(body(false))
+  }
+
+  private def defaultCollectFunc[T](rdd: RDD[T]): Any = rdd.collect()
 
   /**
    * Test checkpointing of the RDD generated by the given operation. It tests whether the
@@ -248,9 +273,12 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * @param op an operation to run on the RDD
    * @param collectFunc a function for collecting the values in the RDD, in case there are
    *   non-comparable types like arrays that we want to convert to something that supports ==
+   * @param normal if true, use normal checkpoints, otherwise use local checkpoints
    */
-  def testRDD[U: ClassTag](op: (RDD[Int]) => RDD[U],
-      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _) {
+  private def testRDD[U: ClassTag](
+      op: (RDD[Int]) => RDD[U],
+      normal: Boolean,
+      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
     // Generate the final RDD using given RDD operation
     val baseRDD = generateFatRDD()
     val operatedRDD = op(baseRDD)
@@ -267,14 +295,16 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     // Find serialized sizes before and after the checkpoint
     logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
     val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
-    operatedRDD.checkpoint()
+    checkpoint(operatedRDD, normal)
     val result = collectFunc(operatedRDD)
     operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
     val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
     logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
 
     // Test whether the checkpoint file has been created
-    assert(collectFunc(sc.checkpointFile[U](operatedRDD.getCheckpointFile.get)) === result)
+    if (normal) {
+      assert(collectFunc(sc.checkpointFile[U](operatedRDD.getCheckpointFile.get)) === result)
+    }
 
     // Test whether dependencies have been changed from its earlier parent RDD
     assert(operatedRDD.dependencies.head.rdd != parentRDD)
@@ -313,8 +343,10 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * @param collectFunc a function for collecting the values in the RDD, in case there are
    *   non-comparable types like arrays that we want to convert to something that supports ==
    */
-  def testRDDPartitions[U: ClassTag](op: (RDD[Int]) => RDD[U],
-       collectFunc: RDD[U] => Any = defaultCollectFunc[U] _) {
+  private def testRDDPartitions[U: ClassTag](
+      op: (RDD[Int]) => RDD[U],
+      normal: Boolean,
+      collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
     // Generate the final RDD using given RDD operation
     val baseRDD = generateFatRDD()
     val operatedRDD = op(baseRDD)
@@ -328,7 +360,10 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     // Find serialized sizes before and after the checkpoint
     logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
     val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
-    parentRDDs.foreach(_.checkpoint())  // checkpoint the parent RDD, not the generated one
+    // checkpoint the parent RDD, not the generated one
+    parentRDDs.foreach { rdd =>
+      checkpoint(rdd, normal)
+    }
     val result = collectFunc(operatedRDD)  // force checkpointing
     operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
     val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
@@ -350,7 +385,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
   /**
    * Generate an RDD such that both the RDD and its partitions have large size.
    */
-  def generateFatRDD(): RDD[Int] = {
+  private def generateFatRDD(): RDD[Int] = {
     new FatRDD(sc.makeRDD(1 to 100, 4)).map(x => x)
   }
 
@@ -358,7 +393,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * Generate an pair RDD (with partitioner) such that both the RDD and its partitions
    * have large size.
    */
-  def generateFatPairRDD(): RDD[(Int, Int)] = {
+  private def generateFatPairRDD(): RDD[(Int, Int)] = {
     new FatPairRDD(sc.makeRDD(1 to 100, 4), partitioner).mapValues(x => x)
   }
 
@@ -366,7 +401,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * Get serialized sizes of the RDD and its partitions, in order to test whether the size shrinks
    * upon checkpointing. Ignores the checkpointData field, which may grow when we checkpoint.
    */
-  def getSerializedSizes(rdd: RDD[_]): (Int, Int) = {
+  private def getSerializedSizes(rdd: RDD[_]): (Int, Int) = {
     val rddSize = Utils.serialize(rdd).size
     val rddCpDataSize = Utils.serialize(rdd.checkpointData).size
     val rddPartitionSize = Utils.serialize(rdd.partitions).size
@@ -394,7 +429,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * contents after deserialization (e.g., the contents of an RDD split after
    * it is sent to a slave along with a task)
    */
-  def serializeDeserialize[T](obj: T): T = {
+  private def serializeDeserialize[T](obj: T): T = {
     val bytes = Utils.serialize(obj)
     Utils.deserialize[T](bytes)
   }
@@ -402,10 +437,11 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
   /**
    * Recursively force the initialization of the all members of an RDD and it parents.
    */
-  def initializeRdd(rdd: RDD[_]) {
+  private def initializeRdd(rdd: RDD[_]): Unit = {
     rdd.partitions // forces the
-    rdd.dependencies.map(_.rdd).foreach(initializeRdd(_))
+    rdd.dependencies.map(_.rdd).foreach(initializeRdd)
   }
+
 }
 
 /** RDD partition that has large serialized size. */

From 2e59646d32b0e16161d434ba05395d84fe2b1f70 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 18:39:38 -0700
Subject: [PATCH 07/27] Add local checkpoint clean up tests

---
 .../apache/spark/ContextCleanerSuite.scala    | 58 +++++++++++++++----
 1 file changed, 46 insertions(+), 12 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
index 05c3d1245d41a..0f2eb847818eb 100644
--- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
@@ -24,11 +24,10 @@ import scala.language.existentials
 import scala.util.Random
 
 import org.scalatest.BeforeAndAfter
-import org.scalatest.concurrent.{PatienceConfiguration, Eventually}
+import org.scalatest.concurrent.PatienceConfiguration
 import org.scalatest.concurrent.Eventually._
 import org.scalatest.time.SpanSugar._
 
-import org.apache.spark.SparkContext._
 import org.apache.spark.rdd.{ReliableRDDCheckpointData, RDD}
 import org.apache.spark.storage._
 import org.apache.spark.shuffle.hash.HashShuffleManager
@@ -52,6 +51,7 @@ abstract class ContextCleanerSuiteBase(val shuffleManager: Class[_] = classOf[Ha
     .setAppName("ContextCleanerSuite")
     .set("spark.cleaner.referenceTracking.blocking", "true")
     .set("spark.cleaner.referenceTracking.blocking.shuffle", "true")
+    .set("spark.cleaner.referenceTracking.cleanCheckpoints", "true")
     .set("spark.shuffle.manager", shuffleManager.getName)
 
   before {
@@ -209,11 +209,11 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     postGCTester.assertCleanup()
   }
 
-  test("automatically cleanup checkpoint") {
+  test("automatically cleanup normal checkpoint") {
     val checkpointDir = java.io.File.createTempFile("temp", "")
     checkpointDir.deleteOnExit()
     checkpointDir.delete()
-    var rdd = newPairRDD
+    var rdd = newPairRDD()
     sc.setCheckpointDir(checkpointDir.toString)
     rdd.checkpoint()
     rdd.cache()
@@ -227,17 +227,20 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     assert(fs.exists(path))
 
     // the checkpoint is not cleaned by default (without the configuration set)
-    var postGCTester = new CleanerTester(sc, Seq(rddId), Nil, Nil, Nil)
+    var postGCTester = new CleanerTester(sc, Seq(rddId), Nil, Nil, Seq(rddId))
     rdd = null // Make RDD out of scope, ok if collected earlier
     runGC()
     postGCTester.assertCleanup()
-    assert(fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
+    assert(!fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
 
+    // Verify that checkpoints are NOT cleaned up if the config is not enabled
     sc.stop()
-    val conf = new SparkConf().setMaster("local[2]").setAppName("cleanupCheckpoint").
-      set("spark.cleaner.referenceTracking.cleanCheckpoints", "true")
+    val conf = new SparkConf()
+      .setMaster("local[2]")
+      .setAppName("cleanupCheckpoint")
+      .set("spark.cleaner.referenceTracking.cleanCheckpoints", "false")
     sc = new SparkContext(conf)
-    rdd = newPairRDD
+    rdd = newPairRDD()
     sc.setCheckpointDir(checkpointDir.toString)
     rdd.checkpoint()
     rdd.cache()
@@ -251,11 +254,39 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
     rdd.count()
 
     // Test that GC causes checkpoint data cleanup after dereferencing the RDD
-    postGCTester = new CleanerTester(sc, Seq(rddId), Nil, Nil, Seq(rddId))
+    postGCTester = new CleanerTester(sc, Seq(rddId))
     rdd = null // Make RDD out of scope
     runGC()
     postGCTester.assertCleanup()
-    assert(!fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
+    assert(fs.exists(ReliableRDDCheckpointData.checkpointPath(sc, rddId).get))
+  }
+
+  test("automatically clean up local checkpoint") {
+    var rdd = newPairRDD().localCheckpoint()
+    assert(rdd.checkpointData.isDefined)
+    assert(rdd.checkpointData.get.checkpointRDD.isEmpty)
+
+    // The checkpoint RDD should be set after an action
+    rdd.count()
+    assert(rdd.checkpointData.get.checkpointRDD.isDefined)
+    var checkpointRdd = rdd.checkpointData.flatMap(_.checkpointRDD).get
+
+    // Do this to make CleanerTester happy (should not affect anything)
+    checkpointRdd.persist()
+
+    // Test that GC does not cause checkpoint cleanup due to a strong reference
+    val preGCTester = new CleanerTester(sc, rddIds = Seq(checkpointRdd.id))
+    runGC()
+    intercept[Exception] {
+      preGCTester.assertCleanup()(timeout(1000 millis))
+    }
+
+    // Test that RDD going out of scope does cause the checkpoint files to be cleaned up
+    val postGCTester = new CleanerTester(sc, rddIds = Seq(checkpointRdd.id))
+    rdd = null
+    checkpointRdd = null
+    runGC()
+    postGCTester.assertCleanup()
   }
 
   test("automatically cleanup RDD + shuffle + broadcast") {
@@ -408,7 +439,10 @@ class SortShuffleContextCleanerSuite extends ContextCleanerSuiteBase(classOf[Sor
 }
 
 
-/** Class to test whether RDDs, shuffles, etc. have been successfully cleaned. */
+/**
+ * Class to test whether RDDs, shuffles, etc. have been successfully cleaned.
+ * The checkpoint here refers only to normal (reliable) checkpoints, not local checkpoints.
+ */
 class CleanerTester(
     sc: SparkContext,
     rddIds: Seq[Int] = Seq.empty,

From 56831c551bab95f46649b9932fc46ea4fc9eaf95 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 18:48:22 -0700
Subject: [PATCH 08/27] Add a few warnings and clear exception messages

---
 .../org/apache/spark/rdd/LocalCheckpointRDD.scala      |  3 +++
 core/src/main/scala/org/apache/spark/rdd/RDD.scala     | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 755c81cd90eef..df464f59db57c 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -80,6 +80,9 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
             s"  Expected indices: ${expectedIndices.mkString(", ")}\n" +
             s"  Actual indices: ${sortedIndices.mkString(", ")}")
       }
+    } else {
+      // An exception with a clear message here is better than a wrong answer
+      throw new SparkException("No checkpointed partitions found when reloading an RDD.")
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 8f07a57e91652..d17cd01185580 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1478,8 +1478,18 @@ abstract class RDD[T: ClassTag](
    *
    * The actual persisting occurs after the first job involving this RDD has completed.
    * The checkpoint directory set through `SparkContext#setCheckpointDir` is not read.
+   *
+   * This is NOT safe to use with dynamic allocation, which removes cached blocks with
+   * executors that it removes. If you must use both features, you are advised to set
+   * `spark.dynamicAllocation.cachedExecutorIdleTimeout` to a high value.
    */
   def localCheckpoint(): this.type = RDDCheckpointData.synchronized {
+    if (conf.getBoolean("spark.dynamicAllocation.enabled", false)) {
+      logWarning("Local checkpointing is NOT safe to use with dynamic allocation, " +
+        "removes cached blocks with executors that it removes. If you must use both " +
+        "features, you are advised to set `spark.dynamicAllocation.cachedExecutorIdleTimout`" +
+        "to a high value.")
+    }
     checkpointData = Some(new LocalRDDCheckpointData(this))
     this
   }

From e53d964829eff58d16a3969aea02eab7ed92f31f Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 21:31:03 -0700
Subject: [PATCH 09/27] Fix style

---
 .../main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala    | 2 +-
 .../scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index df464f59db57c..dbf2828d54d04 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -1,4 +1,4 @@
-/**
+/*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 8fbb9fbf46c89..8db4802584f28 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -1,4 +1,4 @@
-/**
+/*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.

From 172cb6671e417fc17fdda211ab220800a149fc9b Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 7 Jul 2015 23:41:17 -0700
Subject: [PATCH 10/27] Fix mima?

---
 .../scala/org/apache/spark/rdd/LocalCheckpointRDD.scala  | 4 ++--
 project/MimaExcludes.scala                               | 9 +++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index dbf2828d54d04..69c50294cdff8 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -31,7 +31,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
   /**
    * Determine the partitions from the local checkpoint blocks on each executor.
    */
-  override def getPartitions: Array[Partition] = {
+  protected override def getPartitions: Array[Partition] = {
     val ourId = id // define this locally for serialization purposes
     val blockFilter = (blockId: BlockId) => {
       blockId.asRDDId.filter(_.rddId == ourId).isDefined
@@ -49,7 +49,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
   /**
    * Return the location of the checkpoint block that corresponds to the given partition.
    */
-  override def getPreferredLocations(partition: Partition): Seq[String] = {
+  protected override def getPreferredLocations(partition: Partition): Seq[String] = {
     val blockId = RDDBlockId(id, partition.index)
     SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
   }
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 821aadd477ef3..c1eb5dbfec72f 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -75,8 +75,13 @@ object MimaExcludes {
               "org.apache.spark.mllib.linalg.Matrix.numActives")
           ) ++ Seq(
             // SPARK-8914 Remove RDDApi
-            ProblemFilters.exclude[MissingClassProblem](
-            "org.apache.spark.sql.RDDApi")
+            ProblemFilters.exclude[MissingClassProblem]("org.apache.spark.sql.RDDApi")
+          ) ++ Seq(
+            // SPARK-7292 Provide operator to truncate lineage cheaply
+            ProblemFilters.exclude[AbstractClassProblem](
+              "org.apache.spark.rdd.RDDCheckpointData"),
+            ProblemFilters.exclude[AbstractClassProblem](
+              "org.apache.spark.rdd.CheckpointRDD")
           )
 
         case v if v.startsWith("1.4") =>

From d096c674b9c8b292f7fc8fbd1ff37d86a9ee5ac8 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 8 Jul 2015 11:59:29 -0700
Subject: [PATCH 11/27] Fix mima

---
 .../main/scala/org/apache/spark/rdd/CheckpointRDD.scala    | 7 ++++++-
 .../scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala | 4 ++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
index 68a32ada971af..385ec8751a0ee 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
@@ -19,7 +19,7 @@ package org.apache.spark.rdd
 
 import scala.reflect.ClassTag
 
-import org.apache.spark.{Partition, SparkContext}
+import org.apache.spark.{Partition, SparkContext, TaskContext}
 
 private[spark] class CheckpointRDDPartition(val index: Int) extends Partition
 
@@ -29,6 +29,11 @@ private[spark] class CheckpointRDDPartition(val index: Int) extends Partition
 private[spark] abstract class CheckpointRDD[T: ClassTag](@transient sc: SparkContext)
   extends RDD[T](sc, Nil) {
 
+  // Note: override these here to work around a MiMa bug that complains
+  // about `AbstractMethodProblem`s in the RDD class if these are missing
+  protected override def getPartitions: Array[Partition] = ???
+  override def compute(p: Partition, tc: TaskContext): Iterator[T] = ???
+
   // CheckpointRDD should not be checkpointed again
   override def checkpoint(): Unit = { }
   override def doCheckpoint(): Unit = { }
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index 006a25103b013..1940f8b8b9202 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -42,7 +42,7 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
 
   override def getCheckpointFile: Option[String] = Some(checkpointPath)
 
-  override def getPartitions: Array[Partition] = {
+  protected override def getPartitions: Array[Partition] = {
     val cpath = new Path(checkpointPath)
     val numPartitions =
       // listStatus can throw exception if path does not exist.
@@ -59,7 +59,7 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
     Array.tabulate(numPartitions)(i => new CheckpointRDDPartition(i))
   }
 
-  override def getPreferredLocations(split: Partition): Seq[String] = {
+  protected override def getPreferredLocations(split: Partition): Seq[String] = {
     val status = fs.getFileStatus(
       new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index)))
     val locations = fs.getFileBlockLocations(status, 0, status.getLen)

From 4880deb750d7d9362c90742a3592720ef887a52d Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 8 Jul 2015 12:17:04 -0700
Subject: [PATCH 12/27] Fix style

---
 core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
index 385ec8751a0ee..6e141d05f35f2 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
@@ -31,8 +31,10 @@ private[spark] abstract class CheckpointRDD[T: ClassTag](@transient sc: SparkCon
 
   // Note: override these here to work around a MiMa bug that complains
   // about `AbstractMethodProblem`s in the RDD class if these are missing
+  // scalastyle:off
   protected override def getPartitions: Array[Partition] = ???
   override def compute(p: Partition, tc: TaskContext): Iterator[T] = ???
+  // scalastyle:on
 
   // CheckpointRDD should not be checkpointed again
   override def checkpoint(): Unit = { }

From e4cf071ff29939811739fd950c8886e4dd383a40 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 8 Jul 2015 17:34:29 -0700
Subject: [PATCH 13/27] Simplify LocalCheckpointRDD + docs + clean ups

This commit does several things:

(1) First, LocalCheckpointRDD is made significantly simplier.
Instead of fetching block IDs from everyone and verifying whether
the partition indices are continuous, we simply use the original
RDD's partition indices.

(2) Many checkpoint-related methods are now documented, and failure
conditions in local checkpointing now present more informative
error messages.

(3) General code clean ups (reordering things for readability)
---
 .../org/apache/spark/rdd/CheckpointRDD.scala  | 17 +++--
 .../apache/spark/rdd/LocalCheckpointRDD.scala | 72 ++++++++-----------
 .../spark/rdd/LocalRDDCheckpointData.scala    |  9 +--
 .../spark/rdd/ReliableCheckpointRDD.scala     | 43 ++++++-----
 .../spark/rdd/ReliableRDDCheckpointData.scala | 15 ++--
 5 files changed, 79 insertions(+), 77 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
index 6e141d05f35f2..9e24b58ebcd54 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
@@ -21,22 +21,27 @@ import scala.reflect.ClassTag
 
 import org.apache.spark.{Partition, SparkContext, TaskContext}
 
+/**
+ * An RDD partition that maps to a checkpoint file.
+ */
 private[spark] class CheckpointRDDPartition(val index: Int) extends Partition
 
 /**
- * This RDD represents a RDD checkpoint file (similar to HadoopRDD).
+ * An RDD that recovers checkpointed data from persisted files.
  */
 private[spark] abstract class CheckpointRDD[T: ClassTag](@transient sc: SparkContext)
   extends RDD[T](sc, Nil) {
 
-  // Note: override these here to work around a MiMa bug that complains
-  // about `AbstractMethodProblem`s in the RDD class if these are missing
+  // CheckpointRDD should not be checkpointed again
+  override def doCheckpoint(): Unit = { }
+  override def checkpoint(): Unit = { }
+  override def localCheckpoint(): this.type = this
+
+  // Note: There is a bug in MiMa that complains about `AbstractMethodProblem`s in the
+  // base [[org.apache.spark.rdd.RDD]] class if we do not override the following methods.
   // scalastyle:off
   protected override def getPartitions: Array[Partition] = ???
   override def compute(p: Partition, tc: TaskContext): Iterator[T] = ???
   // scalastyle:on
 
-  // CheckpointRDD should not be checkpointed again
-  override def checkpoint(): Unit = { }
-  override def doCheckpoint(): Unit = { }
 }
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 69c50294cdff8..fde890639479e 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -19,35 +19,36 @@ package org.apache.spark.rdd
 
 import scala.reflect.ClassTag
 
-import org.apache.spark.{Partition, SparkEnv, SparkException, SparkContext, TaskContext}
-import org.apache.spark.storage.{BlockId, RDDBlockId}
+import org.apache.spark.{Partition, SparkContext, SparkEnv, SparkException, TaskContext}
+import org.apache.spark.storage.RDDBlockId
 
 /**
- * An RDD that reads from a checkpoint file previously written to the local file system.
+ * An RDD that reads from checkpoint files previously written to the local file system.
+ *
+ * Since local checkpointing is not intended for recovery across applications, it is possible
+ * to always know a priori the exact partitions to compute. There are no guarantees, however,
+ * that the checkpoint files backing these partitions still exist when we try to read them
+ * later. This is because the lifecycle of local checkpoint files is tied to that of executors,
+ * whose failures are conducive to irrecoverable calamity.
  */
-private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext)
+private[spark] class LocalCheckpointRDD[T: ClassTag](
+    @transient sc: SparkContext,
+    originalPartitionIndices: Array[Int])
   extends CheckpointRDD[T](sc) {
 
+  def this(rdd: RDD[T]) {
+    this(rdd.context, rdd.partitions.map(_.index))
+  }
+
   /**
-   * Determine the partitions from the local checkpoint blocks on each executor.
+   * Return partitions that describe how to recover the checkpointed data.
    */
   protected override def getPartitions: Array[Partition] = {
-    val ourId = id // define this locally for serialization purposes
-    val blockFilter = (blockId: BlockId) => {
-      blockId.asRDDId.filter(_.rddId == ourId).isDefined
-    }
-    val inputPartitions: Array[Partition] =
-      SparkEnv.get.blockManager.master
-        .getMatchingBlockIds(blockFilter, askSlaves = true)
-        .collect { case RDDBlockId(_, i) => new CheckpointRDDPartition(i) }
-        .sortBy(_.index)
-        .toArray
-    validateInputPartitions(inputPartitions)
-    inputPartitions
+    originalPartitionIndices.map { i => new CheckpointRDDPartition(i) }
   }
 
   /**
-   * Return the location of the checkpoint block that corresponds to the given partition.
+   * Return the location of the checkpoint block associated with the given partition.
    */
   protected override def getPreferredLocations(partition: Partition): Seq[String] = {
     val blockId = RDDBlockId(id, partition.index)
@@ -55,34 +56,21 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](@transient sc: SparkContext
   }
 
   /**
-   * Fetch the local checkpoint block that corresponds to this partition.
-   * This block should be in the disk store of at least one executor.
+   * Read the content of the checkpoint block associated with this partition.
+   *
+   * Note that the block may not exist if the executor that wrote it is no longer alive.
+   * This is an irrecoverable failure and we should convey this to the user. In normal
+   * cases, however, this block should already be local in this executor's disk store.
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
     val blockId = RDDBlockId(id, partition.index)
     SparkEnv.get.blockManager.get(blockId) match {
-      case Some(result) =>
-        result.data.asInstanceOf[Iterator[T]]
-      case None => throw new SparkException(s"Checkpoint block $blockId not found.")
-    }
-  }
-
-  /**
-   * Validate that the indices of the input partitions are continuous.
-   */
-  private def validateInputPartitions(partitions: Array[Partition]): Unit = {
-    val sortedIndices = partitions.map(_.index).sorted
-    if (sortedIndices.nonEmpty) {
-      val expectedIndices = (sortedIndices.head to sortedIndices.last).toArray
-      if (!java.util.Arrays.equals(sortedIndices, expectedIndices)) {
-        throw new SparkException(
-          "Local checkpoint partitions are invalid.\n" +
-            s"  Expected indices: ${expectedIndices.mkString(", ")}\n" +
-            s"  Actual indices: ${sortedIndices.mkString(", ")}")
-      }
-    } else {
-      // An exception with a clear message here is better than a wrong answer
-      throw new SparkException("No checkpointed partitions found when reloading an RDD.")
+      case Some(result) => result.data.asInstanceOf[Iterator[T]]
+      case None => throw new SparkException(
+        s"Checkpoint block $blockId not found! It is likely that the executor that " +
+        "originally checkpointed this block is no longer alive. If this problem persists, " +
+        "you may consider using `rdd.checkpoint()` instead, which is slower than local " +
+        "checkpointing but more fault-tolerant.")
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 8db4802584f28..086dfbd7cc102 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -33,29 +33,30 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
   extends RDDCheckpointData[T](rdd) with Logging {
 
   /**
-   * Write the content of each partition on a local disk store.
+   * Write the content of each partition to a local disk store.
    */
   protected override def doCheckpoint(): CheckpointRDD[T] = {
-    val checkpointRdd = new LocalCheckpointRDD[T](rdd.context)
 
     // Note: we persist the partitions using the checkpoint RDD's ID rather than the
     // parent RDD's ID. This is not fundamental in design, but allows us to simply
     // reuse the RDD clean up code path to clean the checkpointed files.
+    val checkpointRdd = new LocalCheckpointRDD[T](rdd)
     val checkpointRddId = checkpointRdd.id
     val persistPartition = (taskContext: TaskContext, values: Iterator[T]) => {
       // TODO: if it's already in disk store, just use the existing values
       val blockId = RDDBlockId(checkpointRddId, taskContext.partitionId())
       SparkEnv.get.blockManager.putIterator(blockId, values, StorageLevel.DISK_ONLY)
     }
+    rdd.context.runJob(rdd, persistPartition)
 
-    // Optionally clean our checkpoint files if the reference is out of scope
+    // Since checkpoint files are just cached blocks that belong to the checkpoint RDD,
+    // we can just register it for clean up like any other RDD.
     if (rdd.conf.getBoolean("spark.cleaner.referenceTracking.cleanCheckpoints", false)) {
       rdd.context.cleaner.foreach { cleaner =>
         cleaner.registerRDDForCleanup(checkpointRdd)
       }
     }
 
-    rdd.context.runJob(rdd, persistPartition)
     checkpointRdd
   }
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index 1940f8b8b9202..bb3420342c870 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -29,19 +29,29 @@ import org.apache.spark.deploy.SparkHadoopUtil
 import org.apache.spark.util.{SerializableConfiguration, Utils}
 
 /**
- * An RDD that reads from a checkpoint file previously written to reliable storage.
+ * An RDD that reads from checkpoint files previously written to reliable storage.
  */
 private[spark] class ReliableCheckpointRDD[T: ClassTag](
     @transient sc: SparkContext,
     val checkpointPath: String)
   extends CheckpointRDD[T](sc) {
 
-  private val broadcastedConf = sc.broadcast(new SerializableConfiguration(sc.hadoopConfiguration))
-
-  @transient private val fs = new Path(checkpointPath).getFileSystem(sc.hadoopConfiguration)
+  @transient private val hadoopConf = sc.hadoopConfiguration
+  @transient private val fs = new Path(checkpointPath).getFileSystem(hadoopConf)
+  private val broadcastedConf = sc.broadcast(new SerializableConfiguration(hadoopConf))
 
+  /**
+   * Return the path of the checkpoint directory this RDD reads data from.
+   */
   override def getCheckpointFile: Option[String] = Some(checkpointPath)
 
+  /**
+   * Return partitions described by the files in the checkpoint directory.
+   *
+   * Since the original RDD may belong to a prior application, there is no way to know a
+   * priori the number of partitions to expect. This method assumes that the original set of
+   * checkpoint files are fully preserved in a reliable storage across application lifespans.
+   */
   protected override def getPartitions: Array[Partition] = {
     val cpath = new Path(checkpointPath)
     val numPartitions =
@@ -51,7 +61,12 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
           .map(_.getPath)
           .filter(_.getName.startsWith("part-"))
           .sortBy(_.toString)
-        validateInputFiles(inputFiles)
+        // Fail fast if input files are invalid
+        inputFiles.zipWithIndex.foreach { case (path, i) =>
+          if (!path.toString.endsWith(ReliableCheckpointRDD.splitIdToFile(i))) {
+            throw new SparkException(s"Invalid checkpoint file: $path")
+          }
+        }
         inputFiles.length
       } else {
         0
@@ -59,6 +74,9 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
     Array.tabulate(numPartitions)(i => new CheckpointRDDPartition(i))
   }
 
+  /**
+   * Return the locations of the checkpoint file associated with the given partition.
+   */
   protected override def getPreferredLocations(split: Partition): Seq[String] = {
     val status = fs.getFileStatus(
       new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index)))
@@ -66,23 +84,14 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
     locations.headOption.toList.flatMap(_.getHosts).filter(_ != "localhost")
   }
 
+  /**
+   * Read the content of the checkpoint file associated with the given partition.
+   */
   override def compute(split: Partition, context: TaskContext): Iterator[T] = {
     val file = new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index))
     ReliableCheckpointRDD.readFromFile(file, broadcastedConf, context)
   }
 
-  /**
-   * Fail fast if our input checkpointed files are invalid.
-   */
-  private def validateInputFiles(files: Array[Path]): Unit = {
-    val valid = files.zipWithIndex.forall { case (f, i) =>
-      f.toString.endsWith(ReliableCheckpointRDD.splitIdToFile(i))
-    }
-    if (!valid) {
-      throw new SparkException(s"Invalid checkpoint directory: $checkpointPath")
-    }
-  }
-
 }
 
 private[spark] object ReliableCheckpointRDD extends Logging {
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
index 7387875d1c743..64f276ee6ef17 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
@@ -66,7 +66,14 @@ private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[
     // Save to file, and reload it as an RDD
     val broadcastedConf = rdd.context.broadcast(
       new SerializableConfiguration(rdd.context.hadoopConfiguration))
+    // TODO: This is expensive because it computes the RDD again unnecessarily (SPARK-8582)
+    rdd.context.runJob(rdd, ReliableCheckpointRDD.writeToFile[T](cpDir, broadcastedConf) _)
     val newRDD = new ReliableCheckpointRDD[T](rdd.context, cpDir)
+    if (newRDD.partitions.length != rdd.partitions.length) {
+      throw new SparkException(
+        s"Checkpoint RDD $newRDD(${newRDD.partitions.length}) has different " +
+          s"number of partitions from original RDD $rdd(${rdd.partitions.length})")
+    }
 
     // Optionally clean our checkpoint files if the reference is out of scope
     if (rdd.conf.getBoolean("spark.cleaner.referenceTracking.cleanCheckpoints", false)) {
@@ -75,14 +82,6 @@ private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[
       }
     }
 
-    // TODO: This is expensive because it computes the RDD again unnecessarily (SPARK-8582)
-    rdd.context.runJob(rdd, ReliableCheckpointRDD.writeToFile[T](cpDir, broadcastedConf) _)
-    if (newRDD.partitions.length != rdd.partitions.length) {
-      throw new SparkException(
-        s"Checkpoint RDD $newRDD(${newRDD.partitions.length}) has different " +
-          s"number of partitions from original RDD $rdd(${rdd.partitions.length})")
-    }
-
     logInfo(s"Done checkpointing RDD ${rdd.id} to $cpDir, new parent is RDD ${newRDD.id}")
 
     newRDD

From 53b363ba27166a608039ca73dfc78e4ae193eefd Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 8 Jul 2015 17:41:09 -0700
Subject: [PATCH 14/27] Rename a few more awkwardly named methods (minor)

---
 .../spark/rdd/LocalRDDCheckpointData.scala    | 15 ++++++----
 .../spark/rdd/ReliableCheckpointRDD.scala     | 29 ++++++++++++-------
 .../spark/rdd/ReliableRDDCheckpointData.scala |  2 +-
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 086dfbd7cc102..88dfa84ac1ffa 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -36,14 +36,19 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
    * Write the content of each partition to a local disk store.
    */
   protected override def doCheckpoint(): CheckpointRDD[T] = {
-
-    // Note: we persist the partitions using the checkpoint RDD's ID rather than the
-    // parent RDD's ID. This is not fundamental in design, but allows us to simply
-    // reuse the RDD clean up code path to clean the checkpointed files.
     val checkpointRdd = new LocalCheckpointRDD[T](rdd)
     val checkpointRddId = checkpointRdd.id
     val persistPartition = (taskContext: TaskContext, values: Iterator[T]) => {
-      // TODO: if it's already in disk store, just use the existing values
+
+      // This uses the existing caching interface to write the checkpoint files.
+      // Each partition is cached on disk without replication using the checkpoint RDD's ID.
+      //
+      // The reason why the original RDD's ID is not used is because the original RDD may
+      // already be cached with a different storage level. The alternative of modifying the
+      // original storage level is significantly more complicated downstream especially if
+      // replication is involved.
+      // TODO: if a partition is already in disk store, do not write it again
+
       val blockId = RDDBlockId(checkpointRddId, taskContext.partitionId())
       SparkEnv.get.blockManager.putIterator(blockId, values, StorageLevel.DISK_ONLY)
     }
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index bb3420342c870..05fc069b375f4 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -63,7 +63,7 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
           .sortBy(_.toString)
         // Fail fast if input files are invalid
         inputFiles.zipWithIndex.foreach { case (path, i) =>
-          if (!path.toString.endsWith(ReliableCheckpointRDD.splitIdToFile(i))) {
+          if (!path.toString.endsWith(ReliableCheckpointRDD.checkpointFileName(i))) {
             throw new SparkException(s"Invalid checkpoint file: $path")
           }
         }
@@ -79,7 +79,7 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
    */
   protected override def getPreferredLocations(split: Partition): Seq[String] = {
     val status = fs.getFileStatus(
-      new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index)))
+      new Path(checkpointPath, ReliableCheckpointRDD.checkpointFileName(split.index)))
     val locations = fs.getFileBlockLocations(status, 0, status.getLen)
     locations.headOption.toList.flatMap(_.getHosts).filter(_ != "localhost")
   }
@@ -88,19 +88,25 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
    * Read the content of the checkpoint file associated with the given partition.
    */
   override def compute(split: Partition, context: TaskContext): Iterator[T] = {
-    val file = new Path(checkpointPath, ReliableCheckpointRDD.splitIdToFile(split.index))
-    ReliableCheckpointRDD.readFromFile(file, broadcastedConf, context)
+    val file = new Path(checkpointPath, ReliableCheckpointRDD.checkpointFileName(split.index))
+    ReliableCheckpointRDD.readCheckpointFile(file, broadcastedConf, context)
   }
 
 }
 
 private[spark] object ReliableCheckpointRDD extends Logging {
 
-  private def splitIdToFile(splitId: Int): String = {
-    "part-%05d".format(splitId)
+  /**
+   * Return the checkpoint file name for the given partition.
+   */
+  private def checkpointFileName(partitionIndex: Int): String = {
+    "part-%05d".format(partitionIndex)
   }
 
-  def writeToFile[T: ClassTag](
+  /**
+   * Write this partition's values to a checkpoint file.
+   */
+  def writeCheckpointFile[T: ClassTag](
       path: String,
       broadcastedConf: Broadcast[SerializableConfiguration],
       blockSize: Int = -1)(ctx: TaskContext, iterator: Iterator[T]) {
@@ -108,7 +114,7 @@ private[spark] object ReliableCheckpointRDD extends Logging {
     val outputDir = new Path(path)
     val fs = outputDir.getFileSystem(broadcastedConf.value.value)
 
-    val finalOutputName = ReliableCheckpointRDD.splitIdToFile(ctx.partitionId())
+    val finalOutputName = ReliableCheckpointRDD.checkpointFileName(ctx.partitionId())
     val finalOutputPath = new Path(outputDir, finalOutputName)
     val tempOutputPath =
       new Path(outputDir, s".$finalOutputName-attempt-${ctx.attemptNumber()}")
@@ -146,7 +152,10 @@ private[spark] object ReliableCheckpointRDD extends Logging {
     }
   }
 
-  def readFromFile[T](
+  /**
+   * Read the content of the specified checkpoint file.
+   */
+  def readCheckpointFile[T](
       path: Path,
       broadcastedConf: Broadcast[SerializableConfiguration],
       context: TaskContext): Iterator[T] = {
@@ -174,7 +183,7 @@ private[spark] object ReliableCheckpointRDD extends Logging {
     val conf = SparkHadoopUtil.get.newConfiguration(new SparkConf())
     val fs = path.getFileSystem(conf)
     val broadcastedConf = sc.broadcast(new SerializableConfiguration(conf))
-    sc.runJob(rdd, ReliableCheckpointRDD.writeToFile[Int](path.toString, broadcastedConf, 1024) _)
+    sc.runJob(rdd, ReliableCheckpointRDD.writeCheckpointFile[Int](path.toString, broadcastedConf, 1024) _)
     val cpRDD = new ReliableCheckpointRDD[Int](sc, path.toString)
     assert(cpRDD.partitions.length == rdd.partitions.length, "Number of partitions is not the same")
     assert(cpRDD.collect().toList == rdd.collect().toList, "Data of partitions not the same")
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
index 64f276ee6ef17..1df8eef5ff2b9 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableRDDCheckpointData.scala
@@ -67,7 +67,7 @@ private[spark] class ReliableRDDCheckpointData[T: ClassTag](@transient rdd: RDD[
     val broadcastedConf = rdd.context.broadcast(
       new SerializableConfiguration(rdd.context.hadoopConfiguration))
     // TODO: This is expensive because it computes the RDD again unnecessarily (SPARK-8582)
-    rdd.context.runJob(rdd, ReliableCheckpointRDD.writeToFile[T](cpDir, broadcastedConf) _)
+    rdd.context.runJob(rdd, ReliableCheckpointRDD.writeCheckpointFile[T](cpDir, broadcastedConf) _)
     val newRDD = new ReliableCheckpointRDD[T](rdd.context, cpDir)
     if (newRDD.partitions.length != rdd.partitions.length) {
       throw new SparkException(

From 4a182f3ed23b2f76853ec845c6d190f341eb975c Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 8 Jul 2015 19:57:06 -0700
Subject: [PATCH 15/27] Add fine-grained tests for local checkpointing

This augments the existing end-to-end tests in CheckpointSuite.
---
 .../org/apache/spark/ContextCleaner.scala     |   4 +-
 .../spark/rdd/LocalCheckpointSuite.scala      | 178 ++++++++++++++++++
 2 files changed, 180 insertions(+), 2 deletions(-)
 create mode 100644 core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala

diff --git a/core/src/main/scala/org/apache/spark/ContextCleaner.scala b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
index fabf5a619b1ac..d23c1533db758 100644
--- a/core/src/main/scala/org/apache/spark/ContextCleaner.scala
+++ b/core/src/main/scala/org/apache/spark/ContextCleaner.scala
@@ -232,8 +232,8 @@ private[spark] class ContextCleaner(sc: SparkContext) extends Logging {
   }
 
   /**
-   * Perform checkpoint cleanup on checkpoint files in reliable storage.
-   * Locally checkpointed files are cleaned up separately through RDD clean ups.
+   * Clean up checkpoint files written to a reliable storage.
+   * Locally checkpointed files are cleaned up separately through RDD cleanups.
    */
   def doCleanCheckpoint(rddId: Int): Unit = {
     try {
diff --git a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
new file mode 100644
index 0000000000000..4439283fe774f
--- /dev/null
+++ b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.rdd
+
+import org.apache.spark.{SparkException, SparkContext, LocalSparkContext, SparkFunSuite}
+
+import org.mockito.Mockito.spy
+import org.apache.spark.storage.{RDDBlockId, StorageLevel}
+
+/**
+ * Fine-grained tests for local checkpointing.
+ * For end-to-end tests, see CheckpointSuite.
+ */
+class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
+
+  override def beforeEach(): Unit = {
+    sc = spy(new SparkContext("local[2]", "test"))
+  }
+
+  test("basic lineage truncation") {
+    val numPartitions = 4
+    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
+    val mappedRdd = parallelRdd.map { i => i + 1 }
+    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }
+    val expectedPartitionIndices = (0 until numPartitions).toArray
+    assert(filteredRdd.dependencies.size === 1)
+    assert(filteredRdd.dependencies.head.rdd === mappedRdd)
+    assert(filteredRdd.partitions.map(_.index) === expectedPartitionIndices)
+    assert(filteredRdd.checkpointData.isEmpty)
+    assert(mappedRdd.dependencies.size === 1)
+    assert(parallelRdd.dependencies.size === 0)
+    assert(mappedRdd.dependencies.head.rdd === parallelRdd)
+    filteredRdd.localCheckpoint()
+    assert(filteredRdd.checkpointData.isDefined)
+    assert(!filteredRdd.checkpointData.get.isCheckpointed)
+    assert(!filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+
+    // After an action, the lineage is truncated
+    val result = filteredRdd.collect()
+    assert(filteredRdd.checkpointData.get.isCheckpointed)
+    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+    val checkpointRdd = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get
+    assert(filteredRdd.dependencies.size === 1)
+    assert(filteredRdd.dependencies.head.rdd === checkpointRdd)
+    assert(filteredRdd.partitions.map(_.index) === expectedPartitionIndices)
+    assert(checkpointRdd.partitions.map(_.index) === expectedPartitionIndices)
+
+    // Recomputation should yield the same result
+    assert(filteredRdd.collect() === result)
+    assert(filteredRdd.collect() === result)
+  }
+
+  test("indirect lineage truncation") {
+    val numPartitions = 4
+    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
+    val mappedRdd = parallelRdd.map { i => i + 1 }
+    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }.localCheckpoint()
+    val coalescedRdd = filteredRdd.repartition(10)
+
+    // After an action, only the dependencies of the checkpointed RDD changes
+    val filteredDependencies = filteredRdd.dependencies
+    val coalescedDependencies = coalescedRdd.dependencies
+    val result = coalescedRdd.collect()
+    assert(filteredRdd.dependencies !== filteredDependencies)
+    assert(coalescedRdd.dependencies === coalescedDependencies)
+
+    // Recomputation should yield the same result
+    assert(coalescedRdd.collect() === result)
+    assert(coalescedRdd.collect() === result)
+  }
+
+  test("checkpoint files are in disk store") {
+    val numPartitions = 4
+    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
+    val mappedRdd = parallelRdd.map { i => i + 1 }
+    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }.localCheckpoint()
+    val bmm = sc.env.blockManager.master
+
+    // After an action, the blocks should be found somewhere on disk
+    assert(bmm.getStorageStatus.forall(_.diskUsed == 0))
+    filteredRdd.collect()
+    assert(bmm.getStorageStatus.forall(_.diskUsed > 0))
+    assert(filteredRdd.checkpointData.isDefined)
+    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+
+    // These blocks should be cached using the checkpoint RDD's ID
+    val checkpointRddId = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get.id
+    (0 until numPartitions).foreach { i =>
+      val blockId = RDDBlockId(checkpointRddId, i)
+      val status = bmm.getBlockStatus(blockId).values.head
+      assert(status.storageLevel === StorageLevel.DISK_ONLY)
+      assert(status.diskSize > 0)
+    }
+  }
+
+  test("checkpoint files do not interfere with caching") {
+    val numPartitions = 4
+    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
+    val mappedRdd = parallelRdd.map { i => i + 1 }
+    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }
+      .persist(StorageLevel.MEMORY_ONLY_2) // also cache it in memory
+      .localCheckpoint()
+    val bmm = sc.env.blockManager.master
+
+    // After an action, the blocks should be found somewhere on disk
+    assert(bmm.getStorageStatus.forall(_.memUsed == 0))
+    assert(bmm.getStorageStatus.forall(_.diskUsed == 0))
+    filteredRdd.collect()
+    assert(bmm.getStorageStatus.forall(_.memUsed > 0))
+    assert(bmm.getStorageStatus.forall(_.diskUsed > 0))
+    assert(filteredRdd.checkpointData.isDefined)
+    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+
+    /** Return whether blocks of the specified RDD are cached with a particular storage level. */
+    def areBlocksCached(rddId: Int, level: StorageLevel): Boolean = {
+      (0 until numPartitions).forall { i =>
+        val blockId = RDDBlockId(rddId, i)
+        val status = bmm.getBlockStatus(blockId).values
+        status.nonEmpty && status.head.storageLevel == level
+      }
+    }
+
+    // Checkpoint files should remain even if we unpersist the RDD
+    val checkpointRddId = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get.id
+    assert(areBlocksCached(filteredRdd.id, StorageLevel.MEMORY_ONLY_2))
+    assert(areBlocksCached(checkpointRddId, StorageLevel.DISK_ONLY))
+    filteredRdd.unpersist(blocking = true)
+    assert(!areBlocksCached(filteredRdd.id, StorageLevel.MEMORY_ONLY_2))
+    assert(areBlocksCached(checkpointRddId, StorageLevel.DISK_ONLY))
+  }
+
+  test("missing checkpoint file fails with informative message") {
+    val numPartitions = 4
+    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
+    val mappedRdd = parallelRdd.map { i => i + 1 }
+    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }.localCheckpoint()
+    val bmm = sc.env.blockManager.master
+
+    // After an action, the blocks should be found somewhere on disk
+    filteredRdd.collect()
+    assert(filteredRdd.checkpointData.isDefined)
+    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+    val checkpointRddId = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get.id
+    (0 until numPartitions).foreach { i =>
+      assert(bmm.contains(RDDBlockId(checkpointRddId, i)))
+    }
+
+    // Remove one of the blocks to simulate executor failure
+    // Collecting the RDD should now fail with an informative exception
+    val blockId = RDDBlockId(checkpointRddId, numPartitions - 1)
+    bmm.removeBlock(blockId)
+    try {
+      filteredRdd.collect()
+      fail("Collect should have failed if local checkpoint block is removed...")
+    } catch {
+      case se: SparkException =>
+        assert(se.getMessage.contains(s"Checkpoint block $blockId not found"))
+        assert(se.getMessage.contains("rdd.checkpoint()")) // suggest an alternative
+        assert(se.getMessage.contains("fault-tolerant")) // justify the alternative
+    }
+  }
+
+}

From c449b38f420e07e581541933c53060f319f948ec Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 8 Jul 2015 20:08:19 -0700
Subject: [PATCH 16/27] Fix style

---
 .../scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala     | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index 05fc069b375f4..2dc5cd85a1111 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -183,7 +183,8 @@ private[spark] object ReliableCheckpointRDD extends Logging {
     val conf = SparkHadoopUtil.get.newConfiguration(new SparkConf())
     val fs = path.getFileSystem(conf)
     val broadcastedConf = sc.broadcast(new SerializableConfiguration(conf))
-    sc.runJob(rdd, ReliableCheckpointRDD.writeCheckpointFile[Int](path.toString, broadcastedConf, 1024) _)
+    sc.runJob(
+      rdd, ReliableCheckpointRDD.writeCheckpointFile[Int](path.toString, broadcastedConf, 1024) _)
     val cpRDD = new ReliableCheckpointRDD[Int](sc, path.toString)
     assert(cpRDD.partitions.length == rdd.partitions.length, "Number of partitions is not the same")
     assert(cpRDD.collect().toList == rdd.collect().toList, "Data of partitions not the same")

From db70dc23262bc7baab8d4eacbbf6df0ef2187306 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Mon, 13 Jul 2015 17:52:04 -0700
Subject: [PATCH 17/27] Express local checkpointing through caching the
 original RDD

This allows us to run local checkpointing using MEMORY_AND_DISK
instead of just DISK_ONLY. The motivation is primarily for
performance.

One potential complexity is that it adds more things into
memory store, which may change the LRU behavior for certain
workloads. This happens, for instance, if the RDD action did
not fully drain the iterator (e.g. take).
---
 .../scala/org/apache/spark/SparkContext.scala |   1 +
 .../scala/org/apache/spark/TaskContext.scala  |   8 ++
 .../apache/spark/rdd/LocalCheckpointRDD.scala |   9 +-
 .../spark/rdd/LocalRDDCheckpointData.scala    |  66 +++++----
 .../main/scala/org/apache/spark/rdd/RDD.scala |  20 +++
 .../spark/rdd/LocalCheckpointSuite.scala      | 132 ++++++++++++------
 6 files changed, 160 insertions(+), 76 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 53c6ebb73ed3b..62e2baea170d2 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1741,6 +1741,7 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
     if (conf.getBoolean("spark.logLineage", false)) {
       logInfo("RDD's recursive dependencies:\n" + rdd.toDebugString)
     }
+    rdd.beforeRunJob()
     dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, allowLocal,
       resultHandler, localProperties.get)
     progressBar.foreach(_.finishAll())
diff --git a/core/src/main/scala/org/apache/spark/TaskContext.scala b/core/src/main/scala/org/apache/spark/TaskContext.scala
index d09e17dea0911..7c2238b38197f 100644
--- a/core/src/main/scala/org/apache/spark/TaskContext.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContext.scala
@@ -45,6 +45,14 @@ object TaskContext {
    * Unset the thread local TaskContext. Internal to Spark.
    */
   protected[spark] def unset(): Unit = taskContext.remove()
+
+  /**
+   * Return an empty task context that is not actually used.
+   * Internal use only.
+   */
+  private[spark] def empty(): TaskContext = {
+    new TaskContextImpl(0, 0, 0, 0, null, false)
+  }
 }
 
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index fde890639479e..16c8fc102d246 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -23,7 +23,7 @@ import org.apache.spark.{Partition, SparkContext, SparkEnv, SparkException, Task
 import org.apache.spark.storage.RDDBlockId
 
 /**
- * An RDD that reads from checkpoint files previously written to the local file system.
+ * An RDD that reads from checkpoint files previously written into Spark's caching layer.
  *
  * Since local checkpointing is not intended for recovery across applications, it is possible
  * to always know a priori the exact partitions to compute. There are no guarantees, however,
@@ -33,11 +33,12 @@ import org.apache.spark.storage.RDDBlockId
  */
 private[spark] class LocalCheckpointRDD[T: ClassTag](
     @transient sc: SparkContext,
+    rddId: Int,
     originalPartitionIndices: Array[Int])
   extends CheckpointRDD[T](sc) {
 
   def this(rdd: RDD[T]) {
-    this(rdd.context, rdd.partitions.map(_.index))
+    this(rdd.context, rdd.id, rdd.partitions.map(_.index))
   }
 
   /**
@@ -51,7 +52,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
    * Return the location of the checkpoint block associated with the given partition.
    */
   protected override def getPreferredLocations(partition: Partition): Seq[String] = {
-    val blockId = RDDBlockId(id, partition.index)
+    val blockId = RDDBlockId(rddId, partition.index)
     SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
   }
 
@@ -63,7 +64,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
    * cases, however, this block should already be local in this executor's disk store.
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
-    val blockId = RDDBlockId(id, partition.index)
+    val blockId = RDDBlockId(rddId, partition.index)
     SparkEnv.get.blockManager.get(blockId) match {
       case Some(result) => result.data.asInstanceOf[Iterator[T]]
       case None => throw new SparkException(
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 88dfa84ac1ffa..fd24b1516c93c 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -19,11 +19,11 @@ package org.apache.spark.rdd
 
 import scala.reflect.ClassTag
 
-import org.apache.spark.{Logging, SparkEnv, TaskContext}
+import org.apache.spark.{Logging, SparkEnv, SparkException, TaskContext}
 import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 
 /**
- * An implementation of checkpointing that writes the RDD data to a local file system.
+ * An implementation of checkpointing implemented on top of Spark's caching layer.
  *
  * Local checkpointing trades off fault tolerance for performance by skipping the expensive
  * step of replicating the checkpointed data in a reliable storage. This is useful for use
@@ -33,36 +33,50 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
   extends RDDCheckpointData[T](rdd) with Logging {
 
   /**
-   * Write the content of each partition to a local disk store.
+   * Transform the specified storage level to one that uses disk.
+   *
+   * This guarantees that the RDD can be recomputed multiple times correctly as long as
+   * executors do not fail. Otherwise, if the RDD is cached in memory only, for instance,
+   * the checkpoint data will be lost if the relevant block is evicted from memory.
+   *
+   * This should be called immediately before the first job on the RDD is run.
    */
-  protected override def doCheckpoint(): CheckpointRDD[T] = {
-    val checkpointRdd = new LocalCheckpointRDD[T](rdd)
-    val checkpointRddId = checkpointRdd.id
-    val persistPartition = (taskContext: TaskContext, values: Iterator[T]) => {
-
-      // This uses the existing caching interface to write the checkpoint files.
-      // Each partition is cached on disk without replication using the checkpoint RDD's ID.
-      //
-      // The reason why the original RDD's ID is not used is because the original RDD may
-      // already be cached with a different storage level. The alternative of modifying the
-      // original storage level is significantly more complicated downstream especially if
-      // replication is involved.
-      // TODO: if a partition is already in disk store, do not write it again
-
-      val blockId = RDDBlockId(checkpointRddId, taskContext.partitionId())
-      SparkEnv.get.blockManager.putIterator(blockId, values, StorageLevel.DISK_ONLY)
+  def transformStorageLevel(): Unit = {
+    rdd.getStorageLevel match {
+      case StorageLevel.NONE =>
+        // If this RDD is not already marked for caching, persist it on disk
+        rdd.persist(StorageLevel.DISK_ONLY)
+      case level if level.useOffHeap =>
+        // If this RDD is to be cached off-heap, fail fast since we cannot provide any
+        // correctness guarantees about subsequent computations after the first one
+        throw new SparkException("Local checkpointing is not compatible with off heap caching.")
+      case level =>
+        // Otherwise, adjust the existing storage level to use disk
+        // This guards against potential data losses caused by memory evictions
+        rdd.setStorageLevel(StorageLevel(
+          useDisk = true, level.useMemory, level.deserialized, level.replication))
     }
-    rdd.context.runJob(rdd, persistPartition)
+    assert(rdd.getStorageLevel.isValid, s"Resulting level is invalid: ${rdd.getStorageLevel}")
+  }
+
+  /**
+   * Ensure the RDD is fully cached and return a CheckpointRDD that reads from these blocks.
+   */
+  protected override def doCheckpoint(): CheckpointRDD[T] = {
+    val cm = SparkEnv.get.cacheManager
+    val bmm = SparkEnv.get.blockManager.master
 
-    // Since checkpoint files are just cached blocks that belong to the checkpoint RDD,
-    // we can just register it for clean up like any other RDD.
-    if (rdd.conf.getBoolean("spark.cleaner.referenceTracking.cleanCheckpoints", false)) {
-      rdd.context.cleaner.foreach { cleaner =>
-        cleaner.registerRDDForCleanup(checkpointRdd)
+    // Local checkpointing relies on the fact that all partitions of this RDD are cached.
+    // This may not be the case, however, if the job does not fully drain the RDD iterator.
+    // For this reason, we must force cache any partitions that are not already cached.
+    rdd.partitions.foreach { p =>
+      val blockId = RDDBlockId(rdd.id, p.index)
+      if (!bmm.contains(blockId)) {
+        cm.getOrCompute(rdd, p, TaskContext.empty(), rdd.getStorageLevel)
       }
     }
 
-    checkpointRdd
+    new LocalCheckpointRDD[T](rdd)
   }
 
 }
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index d17cd01185580..9729fa34eafc0 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -188,6 +188,14 @@ abstract class RDD[T: ClassTag](
   /** Get the RDD's current storage level, or StorageLevel.NONE if none is set. */
   def getStorageLevel: StorageLevel = storageLevel
 
+  /**
+   * Override any existing storage level with the one specified.
+   * This is for internal use only and is currently used only for local checkpointing.
+   */
+  private[spark] def setStorageLevel(newLevel: StorageLevel): Unit = {
+    storageLevel = newLevel
+  }
+
   // Our dependencies and partitions will be gotten by calling subclass's methods below, and will
   // be overwritten when we're checkpointed
   private var dependencies_ : Seq[Dependency[_]] = null
@@ -1597,6 +1605,18 @@ abstract class RDD[T: ClassTag](
     deps = null    // Forget the constructor argument for dependencies too
   }
 
+  /**
+   * Callback invoked immediately before a job is run on this RDD.
+   * This recursively calls itself on this RDD's parents.
+   */
+  private[spark] def beforeRunJob(): Unit = {
+    checkpointData match {
+      case Some(local: LocalRDDCheckpointData[T]) => local.transformStorageLevel()
+      case _ =>
+    }
+    dependencies.foreach(_.rdd.beforeRunJob())
+  }
+
   /**
    * Clears the dependencies of this RDD. This method must ensure that all references
    * to the original parent RDDs is removed to enable the parent RDDs to be garbage
diff --git a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
index 4439283fe774f..2425f11c21222 100644
--- a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
@@ -32,6 +32,42 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     sc = spy(new SparkContext("local[2]", "test"))
   }
 
+  test("transform storage level") {
+    val rdd = sc.parallelize(1 to 100).localCheckpoint()
+    assert(rdd.getStorageLevel === StorageLevel.NONE)
+    assert(rdd.checkpointData.isDefined)
+    val data = rdd.checkpointData match {
+      case Some(local: LocalRDDCheckpointData[Int]) => local
+      case _ => fail("Checkpoint data was not of expected type!")
+    }
+
+    // No storage level -> disk only
+    data.transformStorageLevel()
+    assert(rdd.getStorageLevel === StorageLevel.DISK_ONLY)
+
+    // Disk only -> disk only
+    data.transformStorageLevel()
+    assert(rdd.getStorageLevel === StorageLevel.DISK_ONLY)
+
+    // Memory only -> memory and disk
+    rdd.setStorageLevel(StorageLevel.MEMORY_ONLY)
+    data.transformStorageLevel()
+    assert(rdd.getStorageLevel === StorageLevel.MEMORY_AND_DISK)
+
+    // Memory and disk -> memory and disk
+    data.transformStorageLevel()
+    assert(rdd.getStorageLevel === StorageLevel.MEMORY_AND_DISK)
+
+    // Other properties (i.e. whether to serialize, replication factor) should stay
+    rdd.setStorageLevel(StorageLevel.MEMORY_ONLY_SER_2)
+    data.transformStorageLevel()
+    assert(rdd.getStorageLevel === StorageLevel.MEMORY_AND_DISK_SER_2)
+
+    // Off-heap should fail fast
+    rdd.setStorageLevel(StorageLevel.OFF_HEAP)
+    intercept[SparkException] { data.transformStorageLevel() }
+  }
+
   test("basic lineage truncation") {
     val numPartitions = 4
     val parallelRdd = sc.parallelize(1 to 100, numPartitions)
@@ -66,10 +102,10 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
   }
 
   test("indirect lineage truncation") {
-    val numPartitions = 4
-    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
-    val mappedRdd = parallelRdd.map { i => i + 1 }
-    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }.localCheckpoint()
+    val filteredRdd = sc.parallelize(1 to 100, 4)
+      .map { i => i + 1 }
+      .filter { i => i % 2 == 0 }
+      .localCheckpoint()
     val coalescedRdd = filteredRdd.repartition(10)
 
     // After an action, only the dependencies of the checkpointed RDD changes
@@ -84,11 +120,35 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     assert(coalescedRdd.collect() === result)
   }
 
+  test("checkpoint without draining iterator") {
+    val filteredRdd = sc.parallelize(1 to 100, 4)
+      .map { i => i + 1 }
+      .filter { i => i % 2 == 0 }
+      .sortBy(identity) // needed for determinism
+      .localCheckpoint()
+
+    // This does not drain the iterator, but checkpointing should still work
+    val first = filteredRdd.first()
+    assert(filteredRdd.count() === 50)
+    assert(filteredRdd.count() === 50)
+    assert(filteredRdd.first() === first)
+    assert(filteredRdd.first() === first)
+
+    // Test the same thing by calling actions on the child instead
+    val coalescedRdd = filteredRdd.repartition(10)
+    val first2 = coalescedRdd.first()
+    assert(coalescedRdd.count() === 50)
+    assert(coalescedRdd.count() === 50)
+    assert(coalescedRdd.first() === first2)
+    assert(coalescedRdd.first() === first2)
+  }
+
   test("checkpoint files are in disk store") {
     val numPartitions = 4
-    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
-    val mappedRdd = parallelRdd.map { i => i + 1 }
-    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }.localCheckpoint()
+    val filteredRdd = sc.parallelize(1 to 100, numPartitions)
+      .map { i => i + 1 }
+      .filter { i => i % 2 == 0 }
+      .localCheckpoint()
     val bmm = sc.env.blockManager.master
 
     // After an action, the blocks should be found somewhere on disk
@@ -98,71 +158,51 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     assert(filteredRdd.checkpointData.isDefined)
     assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
 
-    // These blocks should be cached using the checkpoint RDD's ID
-    val checkpointRddId = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get.id
     (0 until numPartitions).foreach { i =>
-      val blockId = RDDBlockId(checkpointRddId, i)
+      val blockId = RDDBlockId(filteredRdd.id, i)
       val status = bmm.getBlockStatus(blockId).values.head
       assert(status.storageLevel === StorageLevel.DISK_ONLY)
       assert(status.diskSize > 0)
     }
   }
 
-  test("checkpoint files do not interfere with caching") {
+  test("checkpoint files are in disk store with caching") {
     val numPartitions = 4
-    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
-    val mappedRdd = parallelRdd.map { i => i + 1 }
-    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }
-      .persist(StorageLevel.MEMORY_ONLY_2) // also cache it in memory
+    val filteredRdd = sc.parallelize(1 to 100, numPartitions)
+      .map { i => i + 1 }
+      .filter { i => i % 2 == 0 }
+      .persist(StorageLevel.MEMORY_ONLY)
       .localCheckpoint()
     val bmm = sc.env.blockManager.master
 
-    // After an action, the blocks should be found somewhere on disk
-    assert(bmm.getStorageStatus.forall(_.memUsed == 0))
-    assert(bmm.getStorageStatus.forall(_.diskUsed == 0))
+    // After an action, the blocks should be found in the
+    // block manager with a new storage level that uses disk
     filteredRdd.collect()
-    assert(bmm.getStorageStatus.forall(_.memUsed > 0))
-    assert(bmm.getStorageStatus.forall(_.diskUsed > 0))
-    assert(filteredRdd.checkpointData.isDefined)
-    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
-
-    /** Return whether blocks of the specified RDD are cached with a particular storage level. */
-    def areBlocksCached(rddId: Int, level: StorageLevel): Boolean = {
-      (0 until numPartitions).forall { i =>
-        val blockId = RDDBlockId(rddId, i)
-        val status = bmm.getBlockStatus(blockId).values
-        status.nonEmpty && status.head.storageLevel == level
-      }
+    (0 until numPartitions).foreach { i =>
+      val blockId = RDDBlockId(filteredRdd.id, i)
+      val status = bmm.getBlockStatus(blockId).values.head
+      assert(status.storageLevel === StorageLevel.MEMORY_AND_DISK)
+      assert(status.memSize > 0)
     }
-
-    // Checkpoint files should remain even if we unpersist the RDD
-    val checkpointRddId = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get.id
-    assert(areBlocksCached(filteredRdd.id, StorageLevel.MEMORY_ONLY_2))
-    assert(areBlocksCached(checkpointRddId, StorageLevel.DISK_ONLY))
-    filteredRdd.unpersist(blocking = true)
-    assert(!areBlocksCached(filteredRdd.id, StorageLevel.MEMORY_ONLY_2))
-    assert(areBlocksCached(checkpointRddId, StorageLevel.DISK_ONLY))
   }
 
   test("missing checkpoint file fails with informative message") {
     val numPartitions = 4
-    val parallelRdd = sc.parallelize(1 to 100, numPartitions)
-    val mappedRdd = parallelRdd.map { i => i + 1 }
-    val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }.localCheckpoint()
+    val filteredRdd = sc.parallelize(1 to 100, numPartitions)
+      .map { i => i + 1 }
+      .filter { i => i % 2 == 0 }
+      .localCheckpoint()
     val bmm = sc.env.blockManager.master
 
     // After an action, the blocks should be found somewhere on disk
     filteredRdd.collect()
-    assert(filteredRdd.checkpointData.isDefined)
-    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
-    val checkpointRddId = filteredRdd.checkpointData.flatMap(_.checkpointRDD).get.id
     (0 until numPartitions).foreach { i =>
-      assert(bmm.contains(RDDBlockId(checkpointRddId, i)))
+      assert(bmm.contains(RDDBlockId(filteredRdd.id, i)))
     }
 
     // Remove one of the blocks to simulate executor failure
     // Collecting the RDD should now fail with an informative exception
-    val blockId = RDDBlockId(checkpointRddId, numPartitions - 1)
+    val blockId = RDDBlockId(filteredRdd.id, numPartitions - 1)
     bmm.removeBlock(blockId)
     try {
       filteredRdd.collect()

From 48a999627eb3f3de02de771281553dbc71db0adf Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 14 Jul 2015 20:17:36 -0700
Subject: [PATCH 18/27] Avoid traversing dependency tree + rewrite tests

This commit fixes previous test failures in GraphX and Bagel,
where tests involving long dependency chains would time out.
This is because the previous implementation unconditionally
traversed the entire dependency tree every time before a job
is run, which is not safe if there are circular dependencies.

This commit changes the storage level mutation behavior such
that calling localCheckpoint() with persist(...) in any order
will yield the right storage level. LocalCheckpointSuite is
completely rewritten in a way that ensures this does not
break crucial functionality.
---
 .../scala/org/apache/spark/SparkContext.scala |   1 -
 .../apache/spark/rdd/LocalCheckpointRDD.scala |   8 +-
 .../spark/rdd/LocalRDDCheckpointData.scala    |  49 ++-
 .../main/scala/org/apache/spark/rdd/RDD.scala |  71 ++--
 .../apache/spark/ContextCleanerSuite.scala    |  15 +-
 .../spark/rdd/LocalCheckpointSuite.scala      | 343 ++++++++++++------
 6 files changed, 297 insertions(+), 190 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/SparkContext.scala b/core/src/main/scala/org/apache/spark/SparkContext.scala
index 62e2baea170d2..53c6ebb73ed3b 100644
--- a/core/src/main/scala/org/apache/spark/SparkContext.scala
+++ b/core/src/main/scala/org/apache/spark/SparkContext.scala
@@ -1741,7 +1741,6 @@ class SparkContext(config: SparkConf) extends Logging with ExecutorAllocationCli
     if (conf.getBoolean("spark.logLineage", false)) {
       logInfo("RDD's recursive dependencies:\n" + rdd.toDebugString)
     }
-    rdd.beforeRunJob()
     dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, allowLocal,
       resultHandler, localProperties.get)
     progressBar.foreach(_.finishAll())
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 16c8fc102d246..1e13a0ee33776 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -68,10 +68,10 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
     SparkEnv.get.blockManager.get(blockId) match {
       case Some(result) => result.data.asInstanceOf[Iterator[T]]
       case None => throw new SparkException(
-        s"Checkpoint block $blockId not found! It is likely that the executor that " +
-        "originally checkpointed this block is no longer alive. If this problem persists, " +
-        "you may consider using `rdd.checkpoint()` instead, which is slower than local " +
-        "checkpointing but more fault-tolerant.")
+        s"Checkpoint block $blockId not found! Either the executor that originally " +
+        "checkpointed this block is no longer alive, or the original RDD is unpersisted. " +
+        "If this problem persists, you may consider using `rdd.checkpoint()` instead, " +
+        "which is slower than local checkpointing but more fault-tolerant.")
     }
   }
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index fd24b1516c93c..8644e7e3c5347 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -32,33 +32,6 @@ import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
   extends RDDCheckpointData[T](rdd) with Logging {
 
-  /**
-   * Transform the specified storage level to one that uses disk.
-   *
-   * This guarantees that the RDD can be recomputed multiple times correctly as long as
-   * executors do not fail. Otherwise, if the RDD is cached in memory only, for instance,
-   * the checkpoint data will be lost if the relevant block is evicted from memory.
-   *
-   * This should be called immediately before the first job on the RDD is run.
-   */
-  def transformStorageLevel(): Unit = {
-    rdd.getStorageLevel match {
-      case StorageLevel.NONE =>
-        // If this RDD is not already marked for caching, persist it on disk
-        rdd.persist(StorageLevel.DISK_ONLY)
-      case level if level.useOffHeap =>
-        // If this RDD is to be cached off-heap, fail fast since we cannot provide any
-        // correctness guarantees about subsequent computations after the first one
-        throw new SparkException("Local checkpointing is not compatible with off heap caching.")
-      case level =>
-        // Otherwise, adjust the existing storage level to use disk
-        // This guards against potential data losses caused by memory evictions
-        rdd.setStorageLevel(StorageLevel(
-          useDisk = true, level.useMemory, level.deserialized, level.replication))
-    }
-    assert(rdd.getStorageLevel.isValid, s"Resulting level is invalid: ${rdd.getStorageLevel}")
-  }
-
   /**
    * Ensure the RDD is fully cached and return a CheckpointRDD that reads from these blocks.
    */
@@ -80,3 +53,25 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
   }
 
 }
+
+private[spark] object LocalRDDCheckpointData {
+
+  /**
+   * Transform the specified storage level to one that uses disk.
+   *
+   * This guarantees that the RDD can be recomputed multiple times correctly as long as
+   * executors do not fail. Otherwise, if the RDD is cached in memory only, for instance,
+   * the checkpoint data will be lost if the relevant block is evicted from memory.
+   *
+   * This method is idempotent.
+   */
+  def transformStorageLevel(level: StorageLevel): StorageLevel = {
+    // If this RDD is to be cached off-heap, fail fast since we cannot provide any
+    // correctness guarantees about subsequent computations after the first one
+    if (level.useOffHeap) {
+      throw new SparkException("Local checkpointing is not compatible with off heap caching.")
+    }
+
+    StorageLevel(useDisk = true, level.useMemory, level.deserialized, level.replication)
+  }
+}
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 9729fa34eafc0..3ea737a47f713 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -149,23 +149,42 @@ abstract class RDD[T: ClassTag](
   }
 
   /**
-   * Set this RDD's storage level to persist its values across operations after the first time
-   * it is computed. This can only be used to assign a new storage level if the RDD does not
-   * have a storage level set yet..
+   * Mark this RDD for persisting with the specified level.
+   *
+   * @param newLevel the target storage level
+   * @param allowOverride whether to override any existing level with the new one
    */
-  def persist(newLevel: StorageLevel): this.type = {
+  private def persist(newLevel: StorageLevel, allowOverride: Boolean): this.type = {
     // TODO: Handle changes of StorageLevel
-    if (storageLevel != StorageLevel.NONE && newLevel != storageLevel) {
+    if (storageLevel != StorageLevel.NONE && storageLevel != newLevel && !allowOverride) {
       throw new UnsupportedOperationException(
         "Cannot change storage level of an RDD after it was already assigned a level")
     }
-    sc.persistRDD(this)
-    // Register the RDD with the ContextCleaner for automatic GC-based cleanup
-    sc.cleaner.foreach(_.registerRDDForCleanup(this))
+    // If this is the first time this RDD is marked for persisting, register it
+    // with the SparkContext for cleanups and accounting. Do this only once.
+    if (storageLevel == StorageLevel.NONE) {
+      sc.cleaner.foreach(_.registerRDDForCleanup(this))
+      sc.persistRDD(this)
+    }
     storageLevel = newLevel
     this
   }
 
+  /**
+   * Set this RDD's storage level to persist its values across operations after the first time
+   * it is computed. This can only be used to assign a new storage level if the RDD does not
+   * have a storage level set yet.
+   */
+  def persist(newLevel: StorageLevel): this.type = {
+    if (isLocallyCheckpointed) {
+      // This could happen if the user calls localCheckpoint() before persist(), in which case
+      // we should respect the storage level set by the user explicitly and override the old one
+      persist(LocalRDDCheckpointData.transformStorageLevel(newLevel), allowOverride = true)
+    } else {
+      persist(newLevel, allowOverride = false)
+    }
+  }
+
   /** Persist this RDD with the default storage level (`MEMORY_ONLY`). */
   def persist(): this.type = persist(StorageLevel.MEMORY_ONLY)
 
@@ -188,14 +207,6 @@ abstract class RDD[T: ClassTag](
   /** Get the RDD's current storage level, or StorageLevel.NONE if none is set. */
   def getStorageLevel: StorageLevel = storageLevel
 
-  /**
-   * Override any existing storage level with the one specified.
-   * This is for internal use only and is currently used only for local checkpointing.
-   */
-  private[spark] def setStorageLevel(newLevel: StorageLevel): Unit = {
-    storageLevel = newLevel
-  }
-
   // Our dependencies and partitions will be gotten by calling subclass's methods below, and will
   // be overwritten when we're checkpointed
   private var dependencies_ : Seq[Dependency[_]] = null
@@ -1499,14 +1510,28 @@ abstract class RDD[T: ClassTag](
         "to a high value.")
     }
     checkpointData = Some(new LocalRDDCheckpointData(this))
+    // In case this RDD is already marked for caching, we need to override the
+    // existing storage level with one that is suitable for local checkpointing.
+    persist(LocalRDDCheckpointData.transformStorageLevel(storageLevel), allowOverride = true)
     this
   }
 
   /**
-   * Return whether this RDD has been checkpointed or not, either reliably or locally.
+   * Return whether this RDD is marked for checkpointing, either reliably or locally.
    */
   def isCheckpointed: Boolean = checkpointData.exists(_.isCheckpointed)
 
+  /**
+   * Return whether this RDD is marked for local checkpointing.
+   * Exposed for testing.
+   */
+  private[rdd] def isLocallyCheckpointed: Boolean = {
+    checkpointData match {
+      case Some(_: LocalRDDCheckpointData[T]) => true
+      case _ => false
+    }
+  }
+
   /**
    * Gets the name of the directory to which this RDD was checkpointed.
    * This is not defined if the RDD is checkpointed locally.
@@ -1605,18 +1630,6 @@ abstract class RDD[T: ClassTag](
     deps = null    // Forget the constructor argument for dependencies too
   }
 
-  /**
-   * Callback invoked immediately before a job is run on this RDD.
-   * This recursively calls itself on this RDD's parents.
-   */
-  private[spark] def beforeRunJob(): Unit = {
-    checkpointData match {
-      case Some(local: LocalRDDCheckpointData[T]) => local.transformStorageLevel()
-      case _ =>
-    }
-    dependencies.foreach(_.rdd.beforeRunJob())
-  }
-
   /**
    * Clears the dependencies of this RDD. This method must ensure that all references
    * to the original parent RDDs is removed to enable the parent RDDs to be garbage
diff --git a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
index 0f2eb847818eb..ac29fcdf03684 100644
--- a/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
+++ b/core/src/test/scala/org/apache/spark/ContextCleanerSuite.scala
@@ -262,29 +262,24 @@ class ContextCleanerSuite extends ContextCleanerSuiteBase {
   }
 
   test("automatically clean up local checkpoint") {
+    // Note that this test is similar to the RDD cleanup
+    // test because the same underlying mechanism is used!
     var rdd = newPairRDD().localCheckpoint()
     assert(rdd.checkpointData.isDefined)
     assert(rdd.checkpointData.get.checkpointRDD.isEmpty)
-
-    // The checkpoint RDD should be set after an action
     rdd.count()
     assert(rdd.checkpointData.get.checkpointRDD.isDefined)
-    var checkpointRdd = rdd.checkpointData.flatMap(_.checkpointRDD).get
-
-    // Do this to make CleanerTester happy (should not affect anything)
-    checkpointRdd.persist()
 
     // Test that GC does not cause checkpoint cleanup due to a strong reference
-    val preGCTester = new CleanerTester(sc, rddIds = Seq(checkpointRdd.id))
+    val preGCTester = new CleanerTester(sc, rddIds = Seq(rdd.id))
     runGC()
     intercept[Exception] {
       preGCTester.assertCleanup()(timeout(1000 millis))
     }
 
-    // Test that RDD going out of scope does cause the checkpoint files to be cleaned up
-    val postGCTester = new CleanerTester(sc, rddIds = Seq(checkpointRdd.id))
+    // Test that RDD going out of scope does cause the checkpoint blocks to be cleaned up
+    val postGCTester = new CleanerTester(sc, rddIds = Seq(rdd.id))
     rdd = null
-    checkpointRdd = null
     runGC()
     postGCTester.assertCleanup()
   }
diff --git a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
index 2425f11c21222..2114060abeccd 100644
--- a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
@@ -29,43 +29,26 @@ import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
 
   override def beforeEach(): Unit = {
-    sc = spy(new SparkContext("local[2]", "test"))
+    sc = new SparkContext("local[2]", "test")
   }
 
   test("transform storage level") {
-    val rdd = sc.parallelize(1 to 100).localCheckpoint()
-    assert(rdd.getStorageLevel === StorageLevel.NONE)
-    assert(rdd.checkpointData.isDefined)
-    val data = rdd.checkpointData match {
-      case Some(local: LocalRDDCheckpointData[Int]) => local
-      case _ => fail("Checkpoint data was not of expected type!")
+    val transform = LocalRDDCheckpointData.transformStorageLevel _
+    assert(transform(StorageLevel.NONE) === StorageLevel.DISK_ONLY)
+    assert(transform(StorageLevel.MEMORY_ONLY) === StorageLevel.MEMORY_AND_DISK)
+    assert(transform(StorageLevel.MEMORY_ONLY_SER) === StorageLevel.MEMORY_AND_DISK_SER)
+    assert(transform(StorageLevel.MEMORY_ONLY_2) === StorageLevel.MEMORY_AND_DISK_2)
+    assert(transform(StorageLevel.MEMORY_ONLY_SER_2) === StorageLevel.MEMORY_AND_DISK_SER_2)
+    assert(transform(StorageLevel.DISK_ONLY) === StorageLevel.DISK_ONLY)
+    assert(transform(StorageLevel.DISK_ONLY_2) === StorageLevel.DISK_ONLY_2)
+    assert(transform(StorageLevel.MEMORY_AND_DISK) === StorageLevel.MEMORY_AND_DISK)
+    assert(transform(StorageLevel.MEMORY_AND_DISK_SER) === StorageLevel.MEMORY_AND_DISK_SER)
+    assert(transform(StorageLevel.MEMORY_AND_DISK_2) === StorageLevel.MEMORY_AND_DISK_2)
+    assert(transform(StorageLevel.MEMORY_AND_DISK_SER_2) === StorageLevel.MEMORY_AND_DISK_SER_2)
+    // Off-heap is not supported and Spark should fail fast
+    intercept[SparkException] {
+      transform(StorageLevel.OFF_HEAP)
     }
-
-    // No storage level -> disk only
-    data.transformStorageLevel()
-    assert(rdd.getStorageLevel === StorageLevel.DISK_ONLY)
-
-    // Disk only -> disk only
-    data.transformStorageLevel()
-    assert(rdd.getStorageLevel === StorageLevel.DISK_ONLY)
-
-    // Memory only -> memory and disk
-    rdd.setStorageLevel(StorageLevel.MEMORY_ONLY)
-    data.transformStorageLevel()
-    assert(rdd.getStorageLevel === StorageLevel.MEMORY_AND_DISK)
-
-    // Memory and disk -> memory and disk
-    data.transformStorageLevel()
-    assert(rdd.getStorageLevel === StorageLevel.MEMORY_AND_DISK)
-
-    // Other properties (i.e. whether to serialize, replication factor) should stay
-    rdd.setStorageLevel(StorageLevel.MEMORY_ONLY_SER_2)
-    data.transformStorageLevel()
-    assert(rdd.getStorageLevel === StorageLevel.MEMORY_AND_DISK_SER_2)
-
-    // Off-heap should fail fast
-    rdd.setStorageLevel(StorageLevel.OFF_HEAP)
-    intercept[SparkException] { data.transformStorageLevel() }
   }
 
   test("basic lineage truncation") {
@@ -74,17 +57,21 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     val mappedRdd = parallelRdd.map { i => i + 1 }
     val filteredRdd = mappedRdd.filter { i => i % 2 == 0 }
     val expectedPartitionIndices = (0 until numPartitions).toArray
+    assert(filteredRdd.checkpointData.isEmpty)
+    assert(filteredRdd.getStorageLevel === StorageLevel.NONE)
+    assert(filteredRdd.partitions.map(_.index) === expectedPartitionIndices)
     assert(filteredRdd.dependencies.size === 1)
     assert(filteredRdd.dependencies.head.rdd === mappedRdd)
-    assert(filteredRdd.partitions.map(_.index) === expectedPartitionIndices)
-    assert(filteredRdd.checkpointData.isEmpty)
     assert(mappedRdd.dependencies.size === 1)
-    assert(parallelRdd.dependencies.size === 0)
     assert(mappedRdd.dependencies.head.rdd === parallelRdd)
+    assert(parallelRdd.dependencies.size === 0)
+
+    // Mark the RDD for local checkpointing
     filteredRdd.localCheckpoint()
     assert(filteredRdd.checkpointData.isDefined)
     assert(!filteredRdd.checkpointData.get.isCheckpointed)
     assert(!filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+    assert(filteredRdd.getStorageLevel === StorageLevel.DISK_ONLY)
 
     // After an action, the lineage is truncated
     val result = filteredRdd.collect()
@@ -101,111 +88,86 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     assert(filteredRdd.collect() === result)
   }
 
+  test("basic lineage truncation - caching before checkpointing") {
+    testBasicLineageTruncationWithCaching(
+      newRdd.persist(StorageLevel.MEMORY_ONLY).localCheckpoint(),
+      StorageLevel.MEMORY_AND_DISK)
+  }
+
+  test("basic lineage truncation - caching after checkpointing") {
+    testBasicLineageTruncationWithCaching(
+      newRdd.localCheckpoint().persist(StorageLevel.MEMORY_ONLY),
+      StorageLevel.MEMORY_AND_DISK)
+  }
+
   test("indirect lineage truncation") {
-    val filteredRdd = sc.parallelize(1 to 100, 4)
-      .map { i => i + 1 }
-      .filter { i => i % 2 == 0 }
-      .localCheckpoint()
-    val coalescedRdd = filteredRdd.repartition(10)
+    testIndirectLineageTruncation(newRdd.localCheckpoint(), StorageLevel.DISK_ONLY)
+  }
 
-    // After an action, only the dependencies of the checkpointed RDD changes
-    val filteredDependencies = filteredRdd.dependencies
-    val coalescedDependencies = coalescedRdd.dependencies
-    val result = coalescedRdd.collect()
-    assert(filteredRdd.dependencies !== filteredDependencies)
-    assert(coalescedRdd.dependencies === coalescedDependencies)
+  test("indirect lineage truncation - caching before checkpointing") {
+    testIndirectLineageTruncation(
+      newRdd.persist(StorageLevel.MEMORY_ONLY).localCheckpoint(),
+      StorageLevel.MEMORY_AND_DISK)
+  }
 
-    // Recomputation should yield the same result
-    assert(coalescedRdd.collect() === result)
-    assert(coalescedRdd.collect() === result)
+  test("indirect lineage truncation - caching after checkpointing") {
+    testIndirectLineageTruncation(
+      newRdd.localCheckpoint().persist(StorageLevel.MEMORY_ONLY),
+      StorageLevel.MEMORY_AND_DISK)
   }
 
   test("checkpoint without draining iterator") {
-    val filteredRdd = sc.parallelize(1 to 100, 4)
-      .map { i => i + 1 }
-      .filter { i => i % 2 == 0 }
-      .sortBy(identity) // needed for determinism
-      .localCheckpoint()
+    testWithoutDrainingIterator(newSortedRdd.localCheckpoint(), StorageLevel.DISK_ONLY, 50)
+  }
 
-    // This does not drain the iterator, but checkpointing should still work
-    val first = filteredRdd.first()
-    assert(filteredRdd.count() === 50)
-    assert(filteredRdd.count() === 50)
-    assert(filteredRdd.first() === first)
-    assert(filteredRdd.first() === first)
-
-    // Test the same thing by calling actions on the child instead
-    val coalescedRdd = filteredRdd.repartition(10)
-    val first2 = coalescedRdd.first()
-    assert(coalescedRdd.count() === 50)
-    assert(coalescedRdd.count() === 50)
-    assert(coalescedRdd.first() === first2)
-    assert(coalescedRdd.first() === first2)
-  }
-
-  test("checkpoint files are in disk store") {
-    val numPartitions = 4
-    val filteredRdd = sc.parallelize(1 to 100, numPartitions)
-      .map { i => i + 1 }
-      .filter { i => i % 2 == 0 }
-      .localCheckpoint()
-    val bmm = sc.env.blockManager.master
+  test("checkpoint without draining iterator - caching before checkpointing") {
+    testWithoutDrainingIterator(
+      newSortedRdd.persist(StorageLevel.MEMORY_ONLY).localCheckpoint(),
+      StorageLevel.MEMORY_AND_DISK,
+      50)
+  }
 
-    // After an action, the blocks should be found somewhere on disk
-    assert(bmm.getStorageStatus.forall(_.diskUsed == 0))
-    filteredRdd.collect()
-    assert(bmm.getStorageStatus.forall(_.diskUsed > 0))
-    assert(filteredRdd.checkpointData.isDefined)
-    assert(filteredRdd.checkpointData.get.checkpointRDD.isDefined)
+  test("checkpoint without draining iterator - caching after checkpointing") {
+    testWithoutDrainingIterator(
+      newSortedRdd.localCheckpoint().persist(StorageLevel.MEMORY_ONLY),
+      StorageLevel.MEMORY_AND_DISK,
+      50)
+  }
 
-    (0 until numPartitions).foreach { i =>
-      val blockId = RDDBlockId(filteredRdd.id, i)
-      val status = bmm.getBlockStatus(blockId).values.head
-      assert(status.storageLevel === StorageLevel.DISK_ONLY)
-      assert(status.diskSize > 0)
-    }
+  test("checkpoint blocks exist") {
+    testCheckpointBlocksExist(newRdd.localCheckpoint(), StorageLevel.DISK_ONLY)
   }
 
-  test("checkpoint files are in disk store with caching") {
-    val numPartitions = 4
-    val filteredRdd = sc.parallelize(1 to 100, numPartitions)
-      .map { i => i + 1 }
-      .filter { i => i % 2 == 0 }
-      .persist(StorageLevel.MEMORY_ONLY)
-      .localCheckpoint()
-    val bmm = sc.env.blockManager.master
+  test("checkpoint blocks exist - caching before checkpointing") {
+    testCheckpointBlocksExist(
+      newRdd.persist(StorageLevel.MEMORY_ONLY).localCheckpoint(),
+      StorageLevel.MEMORY_AND_DISK)
+  }
 
-    // After an action, the blocks should be found in the
-    // block manager with a new storage level that uses disk
-    filteredRdd.collect()
-    (0 until numPartitions).foreach { i =>
-      val blockId = RDDBlockId(filteredRdd.id, i)
-      val status = bmm.getBlockStatus(blockId).values.head
-      assert(status.storageLevel === StorageLevel.MEMORY_AND_DISK)
-      assert(status.memSize > 0)
-    }
+  test("checkpoint blocks exist - caching after checkpointing") {
+    testCheckpointBlocksExist(
+      newRdd.localCheckpoint().persist(StorageLevel.MEMORY_ONLY),
+      StorageLevel.MEMORY_AND_DISK)
   }
 
-  test("missing checkpoint file fails with informative message") {
-    val numPartitions = 4
-    val filteredRdd = sc.parallelize(1 to 100, numPartitions)
-      .map { i => i + 1 }
-      .filter { i => i % 2 == 0 }
-      .localCheckpoint()
+  test("missing checkpoint block fails with informative message") {
+    val rdd = newRdd.localCheckpoint()
+    val numPartitions = rdd.partitions.size
+    val partitionIndices = rdd.partitions.map(_.index)
     val bmm = sc.env.blockManager.master
 
-    // After an action, the blocks should be found somewhere on disk
-    filteredRdd.collect()
-    (0 until numPartitions).foreach { i =>
-      assert(bmm.contains(RDDBlockId(filteredRdd.id, i)))
+    // After an action, the blocks should be found somewhere in the cache
+    rdd.collect()
+    partitionIndices.foreach { i =>
+      assert(bmm.contains(RDDBlockId(rdd.id, i)))
     }
 
     // Remove one of the blocks to simulate executor failure
     // Collecting the RDD should now fail with an informative exception
-    val blockId = RDDBlockId(filteredRdd.id, numPartitions - 1)
+    val blockId = RDDBlockId(rdd.id, numPartitions - 1)
     bmm.removeBlock(blockId)
     try {
-      filteredRdd.collect()
+      rdd.collect()
       fail("Collect should have failed if local checkpoint block is removed...")
     } catch {
       case se: SparkException =>
@@ -215,4 +177,147 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     }
   }
 
+  /**
+   * Helper method to create a simple RDD.
+   */
+  private def newRdd: RDD[Int] = {
+    sc.parallelize(1 to 100, 4)
+      .map { i => i + 1 }
+      .filter { i => i % 2 == 0 }
+  }
+
+  /**
+   * Helper method to create a simple sorted RDD.
+   */
+  private def newSortedRdd: RDD[Int] = newRdd.sortBy(identity)
+
+  /**
+   * Helper method to test basic lineage truncation with caching.
+   *
+   * @param rdd an RDD that is both marked for caching and local checkpointing
+   */
+  private def testBasicLineageTruncationWithCaching[T](
+      rdd: RDD[T],
+      targetStorageLevel: StorageLevel): Unit = {
+    require(targetStorageLevel !== StorageLevel.NONE)
+    require(rdd.getStorageLevel !== StorageLevel.NONE)
+    require(rdd.isLocallyCheckpointed)
+    val result = rdd.collect()
+    assert(rdd.getStorageLevel === targetStorageLevel)
+    assert(rdd.checkpointData.isDefined)
+    assert(rdd.checkpointData.get.isCheckpointed)
+    assert(rdd.checkpointData.get.checkpointRDD.isDefined)
+    assert(rdd.dependencies.head.rdd === rdd.checkpointData.get.checkpointRDD.get)
+    assert(rdd.collect() === result)
+    assert(rdd.collect() === result)
+  }
+
+  /**
+   * Helper method to test indirect lineage truncation.
+   *
+   * Indirect lineage truncation here means the action is called on one of the
+   * checkpointed RDD's descendants, but not on the checkpointed RDD itself.
+   *
+   * @param rdd a locally checkpointed RDD
+   */
+  private def testIndirectLineageTruncation[T](
+      rdd: RDD[T],
+      targetStorageLevel: StorageLevel): Unit = {
+    require(targetStorageLevel !== StorageLevel.NONE)
+    require(rdd.isLocallyCheckpointed)
+    val rdd1 = rdd.map { i => i + "1" }
+    val rdd2 = rdd1.map { i => i + "2" }
+    val rdd3 = rdd2.map { i => i + "3" }
+    val rddDependencies = rdd.dependencies
+    val rdd1Dependencies = rdd1.dependencies
+    val rdd2Dependencies = rdd2.dependencies
+    val rdd3Dependencies = rdd3.dependencies
+    assert(rdd1Dependencies.size === 1)
+    assert(rdd1Dependencies.head.rdd === rdd)
+    assert(rdd2Dependencies.size === 1)
+    assert(rdd2Dependencies.head.rdd === rdd1)
+    assert(rdd3Dependencies.size === 1)
+    assert(rdd3Dependencies.head.rdd === rdd2)
+
+    // Only the locally checkpointed RDD should have special storage level
+    assert(rdd.getStorageLevel === targetStorageLevel)
+    assert(rdd1.getStorageLevel === StorageLevel.NONE)
+    assert(rdd2.getStorageLevel === StorageLevel.NONE)
+    assert(rdd3.getStorageLevel === StorageLevel.NONE)
+
+    // After an action, only the dependencies of the checkpointed RDD changes
+    val result = rdd3.collect()
+    assert(rdd.dependencies !== rddDependencies)
+    assert(rdd1.dependencies === rdd1Dependencies)
+    assert(rdd2.dependencies === rdd2Dependencies)
+    assert(rdd3.dependencies === rdd3Dependencies)
+    assert(rdd3.collect() === result)
+    assert(rdd3.collect() === result)
+  }
+
+  /**
+   * Helper method to test checkpointing without fully draining the iterator.
+   *
+   * Not all RDD actions fully consume the iterator. As a result, a subset of the partitions
+   * may not be cached. However, since we want to truncate the lineage safely, we explicitly
+   * ensure that *all* partitions are fully cached. This method asserts this behavior.
+   *
+   * @param rdd a locally checkpointed RDD
+   */
+  private def testWithoutDrainingIterator[T](
+      rdd: RDD[T],
+      targetStorageLevel: StorageLevel,
+      targetCount: Int): Unit = {
+    require(targetCount > 0)
+    require(targetStorageLevel !== StorageLevel.NONE)
+    require(rdd.isLocallyCheckpointed)
+
+    // This does not drain the iterator, but checkpointing should still work
+    val first = rdd.first()
+    assert(rdd.count() === targetCount)
+    assert(rdd.count() === targetCount)
+    assert(rdd.first() === first)
+    assert(rdd.first() === first)
+
+    // Test the same thing by calling actions on a descendant instead
+    val rdd1 = rdd.repartition(10)
+    val rdd2 = rdd1.repartition(100)
+    val rdd3 = rdd2.repartition(1000)
+    val first2 = rdd3.first()
+    assert(rdd3.count() === targetCount)
+    assert(rdd3.count() === targetCount)
+    assert(rdd3.first() === first2)
+    assert(rdd3.first() === first2)
+    assert(rdd.getStorageLevel === targetStorageLevel)
+    assert(rdd1.getStorageLevel === StorageLevel.NONE)
+    assert(rdd2.getStorageLevel === StorageLevel.NONE)
+    assert(rdd3.getStorageLevel === StorageLevel.NONE)
+  }
+
+  /**
+   * Helper method to test whether the checkpoint blocks are found in the cache.
+   *
+   * @param rdd a locally checkpointed RDD
+   */
+  private def testCheckpointBlocksExist[T](
+      rdd: RDD[T],
+      targetStorageLevel: StorageLevel): Unit = {
+    val bmm = sc.env.blockManager.master
+    val partitionIndices = rdd.partitions.map(_.index)
+
+    // The blocks should not exist before the action
+    partitionIndices.foreach { i =>
+      assert(!bmm.contains(RDDBlockId(rdd.id, i)))
+    }
+
+    // After an action, the blocks should be found in the cache with the expected level
+    rdd.collect()
+    partitionIndices.foreach { i =>
+      val blockId = RDDBlockId(rdd.id, i)
+      val status = bmm.getBlockStatus(blockId)
+      assert(status.nonEmpty)
+      assert(status.values.head.storageLevel === targetStorageLevel)
+    }
+  }
+
 }

From 1bbe154e43143fa6f2a2d6035e6253f244bb3785 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Tue, 14 Jul 2015 20:34:36 -0700
Subject: [PATCH 19/27] Simplify LocalCheckpointRDD

It doesn't even need to fetch blocks because the child should
already do this through getOrCompute. This commit makes it simply
a placeholder for providing error messages when things fail.
---
 .../apache/spark/rdd/LocalCheckpointRDD.scala | 46 +++++++------------
 1 file changed, 17 insertions(+), 29 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 1e13a0ee33776..3067d0fc14919 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -23,18 +23,21 @@ import org.apache.spark.{Partition, SparkContext, SparkEnv, SparkException, Task
 import org.apache.spark.storage.RDDBlockId
 
 /**
- * An RDD that reads from checkpoint files previously written into Spark's caching layer.
+ * A dummy CheckpointRDD that exists to provide informative error messages during failures.
  *
- * Since local checkpointing is not intended for recovery across applications, it is possible
- * to always know a priori the exact partitions to compute. There are no guarantees, however,
- * that the checkpoint files backing these partitions still exist when we try to read them
- * later. This is because the lifecycle of local checkpoint files is tied to that of executors,
- * whose failures are conducive to irrecoverable calamity.
+ * This is simply a placeholder because the original checkpointed RDD is expected to be
+ * fully cached. Only if an executor fails or if the user explicitly unpersists the original
+ * RDD will Spark ever attempt to compute this CheckpointRDD. When this happens, however,
+ * we must provide an informative error message.
+ *
+ * @param sc the active SparkContext
+ * @param rddId the ID of the checkpointed RDD
+ * @param partitionIndices the partitionIndices of the checkpointed RDD
  */
 private[spark] class LocalCheckpointRDD[T: ClassTag](
     @transient sc: SparkContext,
     rddId: Int,
-    originalPartitionIndices: Array[Int])
+    partitionIndices: Array[Int])
   extends CheckpointRDD[T](sc) {
 
   def this(rdd: RDD[T]) {
@@ -45,34 +48,19 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
    * Return partitions that describe how to recover the checkpointed data.
    */
   protected override def getPartitions: Array[Partition] = {
-    originalPartitionIndices.map { i => new CheckpointRDDPartition(i) }
-  }
-
-  /**
-   * Return the location of the checkpoint block associated with the given partition.
-   */
-  protected override def getPreferredLocations(partition: Partition): Seq[String] = {
-    val blockId = RDDBlockId(rddId, partition.index)
-    SparkEnv.get.blockManager.master.getLocations(blockId).map(_.host)
+    partitionIndices.map { i => new CheckpointRDDPartition(i) }
   }
 
   /**
-   * Read the content of the checkpoint block associated with this partition.
-   *
-   * Note that the block may not exist if the executor that wrote it is no longer alive.
-   * This is an irrecoverable failure and we should convey this to the user. In normal
-   * cases, however, this block should already be local in this executor's disk store.
+   * Throw an exception indicating that the relevant block is not found.
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
     val blockId = RDDBlockId(rddId, partition.index)
-    SparkEnv.get.blockManager.get(blockId) match {
-      case Some(result) => result.data.asInstanceOf[Iterator[T]]
-      case None => throw new SparkException(
-        s"Checkpoint block $blockId not found! Either the executor that originally " +
-        "checkpointed this block is no longer alive, or the original RDD is unpersisted. " +
-        "If this problem persists, you may consider using `rdd.checkpoint()` instead, " +
-        "which is slower than local checkpointing but more fault-tolerant.")
-    }
+    throw new SparkException(
+      s"Checkpoint block $blockId not found! Either the executor that originally " +
+      "checkpointed this block is no longer alive, or the original RDD is unpersisted. " +
+      "If this problem persists, you may consider using `rdd.checkpoint()` instead, " +
+      "which is slower than local checkpointing but more fault-tolerant.")
   }
 
 }

From a92657d815e7837a64d69546acc954a792ae1d1a Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Wed, 15 Jul 2015 17:59:24 -0700
Subject: [PATCH 20/27] Update a few comments

---
 .../org/apache/spark/rdd/CheckpointRDD.scala  |  4 +-
 .../apache/spark/rdd/LocalCheckpointRDD.scala | 19 ++++---
 .../spark/rdd/LocalRDDCheckpointData.scala    | 20 ++++---
 .../main/scala/org/apache/spark/rdd/RDD.scala | 57 +++++++++++--------
 .../apache/spark/rdd/RDDCheckpointData.scala  |  4 +-
 5 files changed, 59 insertions(+), 45 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
index 9e24b58ebcd54..72fe215dae73e 100644
--- a/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/CheckpointRDD.scala
@@ -22,12 +22,12 @@ import scala.reflect.ClassTag
 import org.apache.spark.{Partition, SparkContext, TaskContext}
 
 /**
- * An RDD partition that maps to a checkpoint file.
+ * An RDD partition used to recover checkpointed data.
  */
 private[spark] class CheckpointRDDPartition(val index: Int) extends Partition
 
 /**
- * An RDD that recovers checkpointed data from persisted files.
+ * An RDD that recovers checkpointed data from storage.
  */
 private[spark] abstract class CheckpointRDD[T: ClassTag](@transient sc: SparkContext)
   extends RDD[T](sc, Nil) {
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index 3067d0fc14919..c80548824d27c 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -32,7 +32,7 @@ import org.apache.spark.storage.RDDBlockId
  *
  * @param sc the active SparkContext
  * @param rddId the ID of the checkpointed RDD
- * @param partitionIndices the partitionIndices of the checkpointed RDD
+ * @param partitionIndices the partition indices of the checkpointed RDD
  */
 private[spark] class LocalCheckpointRDD[T: ClassTag](
     @transient sc: SparkContext,
@@ -44,23 +44,24 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
     this(rdd.context, rdd.id, rdd.partitions.map(_.index))
   }
 
-  /**
-   * Return partitions that describe how to recover the checkpointed data.
-   */
   protected override def getPartitions: Array[Partition] = {
     partitionIndices.map { i => new CheckpointRDDPartition(i) }
   }
 
   /**
    * Throw an exception indicating that the relevant block is not found.
+   *
+   * This should only be called if the original RDD is explicitly unpersisted or if an
+   * executor is lost. Under normal circumstances, however, the original RDD (our child)
+   * is expected to be fully cached and so all partitions should already be computed and
+   * available in the block storage.
    */
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
-    val blockId = RDDBlockId(rddId, partition.index)
     throw new SparkException(
-      s"Checkpoint block $blockId not found! Either the executor that originally " +
-      "checkpointed this block is no longer alive, or the original RDD is unpersisted. " +
-      "If this problem persists, you may consider using `rdd.checkpoint()` instead, " +
-      "which is slower than local checkpointing but more fault-tolerant.")
+      s"Checkpoint block ${RDDBlockId(rddId, partition.index)} not found! Either the executor " +
+      s"that originally checkpointed this block is no longer alive, or the original RDD is " +
+      s"unpersisted. If this problem persists, you may consider using `rdd.checkpoint()` " +
+      s"instead, which is slower than local checkpointing but more fault-tolerant.")
   }
 
 }
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 8644e7e3c5347..2a492b2fdeee8 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -26,26 +26,30 @@ import org.apache.spark.storage.{RDDBlockId, StorageLevel}
  * An implementation of checkpointing implemented on top of Spark's caching layer.
  *
  * Local checkpointing trades off fault tolerance for performance by skipping the expensive
- * step of replicating the checkpointed data in a reliable storage. This is useful for use
- * cases where RDDs build up long lineages that need to be truncated often (e.g. GraphX).
+ * step of replicating the checkpointed data in a reliable storage. Instead, checkpoint data
+ * is written to the local, ephemeral block storage that lives in each executor. This is useful
+ * for use cases where RDDs build up long lineages that need to be truncated often (e.g. GraphX).
  */
 private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
   extends RDDCheckpointData[T](rdd) with Logging {
 
   /**
-   * Ensure the RDD is fully cached and return a CheckpointRDD that reads from these blocks.
+   * Ensure the RDD is fully cached so the partitions can be recovered later.
    */
   protected override def doCheckpoint(): CheckpointRDD[T] = {
     val cm = SparkEnv.get.cacheManager
     val bmm = SparkEnv.get.blockManager.master
+    val level = rdd.getStorageLevel
 
-    // Local checkpointing relies on the fact that all partitions of this RDD are cached.
-    // This may not be the case, however, if the job does not fully drain the RDD iterator.
-    // For this reason, we must force cache any partitions that are not already cached.
+    // Assume storage level uses disk; otherwise memory eviction may cause data loss
+    assume(level.useDisk, s"Storage level $level is not appropriate for local checkpointing")
+
+    // Not all RDD actions compute all partitions of the RDD (e.g. take)
+    // We must compute and cache any missing partitions for correctness reasons
     rdd.partitions.foreach { p =>
       val blockId = RDDBlockId(rdd.id, p.index)
       if (!bmm.contains(blockId)) {
-        cm.getOrCompute(rdd, p, TaskContext.empty(), rdd.getStorageLevel)
+        cm.getOrCompute(rdd, p, TaskContext.empty(), level)
       }
     }
 
@@ -69,7 +73,7 @@ private[spark] object LocalRDDCheckpointData {
     // If this RDD is to be cached off-heap, fail fast since we cannot provide any
     // correctness guarantees about subsequent computations after the first one
     if (level.useOffHeap) {
-      throw new SparkException("Local checkpointing is not compatible with off heap caching.")
+      throw new SparkException("Local checkpointing is not compatible with off-heap caching.")
     }
 
     StorageLevel(useDisk = true, level.useMemory, level.deserialized, level.replication)
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 3ea737a47f713..7c61e54ea28f3 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -149,14 +149,14 @@ abstract class RDD[T: ClassTag](
   }
 
   /**
-   * Mark this RDD for persisting with the specified level.
+   * Mark this RDD for persisting using the specified level.
    *
    * @param newLevel the target storage level
    * @param allowOverride whether to override any existing level with the new one
    */
   private def persist(newLevel: StorageLevel, allowOverride: Boolean): this.type = {
     // TODO: Handle changes of StorageLevel
-    if (storageLevel != StorageLevel.NONE && storageLevel != newLevel && !allowOverride) {
+    if (storageLevel != StorageLevel.NONE && newLevel != storageLevel && !allowOverride) {
       throw new UnsupportedOperationException(
         "Cannot change storage level of an RDD after it was already assigned a level")
     }
@@ -173,12 +173,13 @@ abstract class RDD[T: ClassTag](
   /**
    * Set this RDD's storage level to persist its values across operations after the first time
    * it is computed. This can only be used to assign a new storage level if the RDD does not
-   * have a storage level set yet.
+   * have a storage level set yet. Local checkpointing is an exception.
    */
   def persist(newLevel: StorageLevel): this.type = {
     if (isLocallyCheckpointed) {
-      // This could happen if the user calls localCheckpoint() before persist(), in which case
-      // we should respect the storage level set by the user explicitly and override the old one
+      // This means the user previously called localCheckpoint(), which should have already
+      // marked this RDD for persisting. Here we should override the old storage level with
+      // one that is explicitly requested by the user (after adapting it to use disk).
       persist(LocalRDDCheckpointData.transformStorageLevel(newLevel), allowOverride = true)
     } else {
       persist(newLevel, allowOverride = false)
@@ -1484,35 +1485,43 @@ abstract class RDD[T: ClassTag](
   }
 
   /**
-   * Mark this RDD for unsafe checkpointing using a local file system.
+   * Mark this RDD for local checkpointing using Spark's existing caching layer.
    *
-   * This method is for users who wish to truncate the lineage of the RDD while skipping the
-   * expensive step of replicating the materialized data in a distributed file system. This is
+   * This method is for users who wish to truncate RDD lineages while skipping the expensive
+   * step of replicating the materialized data in a reliable distributed file system. This is
    * useful for RDDs with long lineages that need to be truncated periodically (e.g. GraphX).
    *
-   * It is unsafe because we sacrifice fault-tolerance for performance. In particular, we
-   * persist the materialized values only to an ephemeral local storage. The effect is that
-   * if an executor fails during the computation, the checkpointed data will no longer be
-   * accessible.
+   * Local checkpointing sacrifices fault-tolerance for performance. In particular,
+   * checkpointed data is written to ephemeral local storage (memory or disk) instead of
+   * to a reliable distributed storage. The effect is that if an executor fails during the
+   * computation, the checkpointed data may no longer be accessible, causing an irrecoverable
+   * job failure.
    *
-   * The actual persisting occurs after the first job involving this RDD has completed.
-   * The checkpoint directory set through `SparkContext#setCheckpointDir` is not read.
-   *
-   * This is NOT safe to use with dynamic allocation, which removes cached blocks with
-   * executors that it removes. If you must use both features, you are advised to set
+   * This is NOT safe to use with dynamic allocation, which removes executors along
+   * with their cached blocks. If you must use both features, you are advised to set
    * `spark.dynamicAllocation.cachedExecutorIdleTimeout` to a high value.
+   *
+   * The checkpoint directory set through `SparkContext#setCheckpointDir` is not read.
    */
   def localCheckpoint(): this.type = RDDCheckpointData.synchronized {
-    if (conf.getBoolean("spark.dynamicAllocation.enabled", false)) {
+    if (conf.getBoolean("spark.dynamicAllocation.enabled", false) &&
+        conf.contains("spark.dynamicAllocation.cachedExecutorIdleTimeout")) {
       logWarning("Local checkpointing is NOT safe to use with dynamic allocation, " +
-        "removes cached blocks with executors that it removes. If you must use both " +
-        "features, you are advised to set `spark.dynamicAllocation.cachedExecutorIdleTimout`" +
+        "which removes executors along with their cached blocks. If you must use both " +
+        "features, you are advised to set `spark.dynamicAllocation.cachedExecutorIdleTimeout` " +
         "to a high value.")
     }
-    checkpointData = Some(new LocalRDDCheckpointData(this))
-    // In case this RDD is already marked for caching, we need to override the
-    // existing storage level with one that is suitable for local checkpointing.
+
+    // Note: At this point we do not actually know whether the user will call persist() on
+    // this RDD later, so we must explicitly call it here ourselves to ensure the cached
+    // blocks are registered for cleanup later in the SparkContext.
+    //
+    // If, however, the user has already called persist() on this RDD, then we must adapt
+    // the storage level he/she specified to one that is appropriate for local checkpointing
+    // (i.e. uses disk) to guarantee correctness.
+
     persist(LocalRDDCheckpointData.transformStorageLevel(storageLevel), allowOverride = true)
+    checkpointData = Some(new LocalRDDCheckpointData(this))
     this
   }
 
@@ -1624,7 +1633,7 @@ abstract class RDD[T: ClassTag](
    * Changes the dependencies of this RDD from its original parents to a new RDD (`newRDD`)
    * created from the checkpoint file, and forget its old dependencies and partitions.
    */
-  private[spark] def markCheckpointed(checkpointRDD: RDD[_]) {
+  private[spark] def markCheckpointed(): Unit = {
     clearDependencies()
     partitions_ = null
     deps = null    // Forget the constructor argument for dependencies too
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
index 30bcdd95d8a47..0e43520870c0a 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDDCheckpointData.scala
@@ -73,11 +73,11 @@ private[spark] abstract class RDDCheckpointData[T: ClassTag](@transient rdd: RDD
 
     val newRDD = doCheckpoint()
 
-    // Update our state and mark the RDD as checkpointed
+    // Update our state and truncate the RDD lineage
     RDDCheckpointData.synchronized {
       cpRDD = Some(newRDD)
       cpState = Checkpointed
-      rdd.markCheckpointed(newRDD)
+      rdd.markCheckpointed()
     }
   }
 

From f5be0f37c6e9941ad03df053072f9c2b0d4fbffa Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Thu, 16 Jul 2015 16:44:08 -0700
Subject: [PATCH 21/27] Use MEMORY_AND_DISK as the default local checkpoint
 level

For performance!
---
 .../scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala | 2 ++
 core/src/main/scala/org/apache/spark/rdd/RDD.scala          | 6 +++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 2a492b2fdeee8..b7f86c4145dd2 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -60,6 +60,8 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
 
 private[spark] object LocalRDDCheckpointData {
 
+  val DEFAULT_STORAGE_LEVEL = StorageLevel.MEMORY_AND_DISK
+
   /**
    * Transform the specified storage level to one that uses disk.
    *
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 7c61e54ea28f3..f08ead284add7 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1520,7 +1520,11 @@ abstract class RDD[T: ClassTag](
     // the storage level he/she specified to one that is appropriate for local checkpointing
     // (i.e. uses disk) to guarantee correctness.
 
-    persist(LocalRDDCheckpointData.transformStorageLevel(storageLevel), allowOverride = true)
+    if (storageLevel == StorageLevel.NONE) {
+      persist(LocalRDDCheckpointData.DEFAULT_STORAGE_LEVEL)
+    } else {
+      persist(LocalRDDCheckpointData.transformStorageLevel(storageLevel), allowOverride = true)
+    }
     checkpointData = Some(new LocalRDDCheckpointData(this))
     this
   }

From e908a42c698d295e78ffd65f42b1c3ca97e9f46d Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Thu, 16 Jul 2015 19:19:55 -0700
Subject: [PATCH 22/27] Fix tests

---
 .../apache/spark/rdd/LocalCheckpointSuite.scala   | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
index 2114060abeccd..5103eb74b2457 100644
--- a/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/LocalCheckpointSuite.scala
@@ -71,7 +71,7 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
     assert(filteredRdd.checkpointData.isDefined)
     assert(!filteredRdd.checkpointData.get.isCheckpointed)
     assert(!filteredRdd.checkpointData.get.checkpointRDD.isDefined)
-    assert(filteredRdd.getStorageLevel === StorageLevel.DISK_ONLY)
+    assert(filteredRdd.getStorageLevel === LocalRDDCheckpointData.DEFAULT_STORAGE_LEVEL)
 
     // After an action, the lineage is truncated
     val result = filteredRdd.collect()
@@ -101,7 +101,9 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
   }
 
   test("indirect lineage truncation") {
-    testIndirectLineageTruncation(newRdd.localCheckpoint(), StorageLevel.DISK_ONLY)
+    testIndirectLineageTruncation(
+      newRdd.localCheckpoint(),
+      LocalRDDCheckpointData.DEFAULT_STORAGE_LEVEL)
   }
 
   test("indirect lineage truncation - caching before checkpointing") {
@@ -117,7 +119,10 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
   }
 
   test("checkpoint without draining iterator") {
-    testWithoutDrainingIterator(newSortedRdd.localCheckpoint(), StorageLevel.DISK_ONLY, 50)
+    testWithoutDrainingIterator(
+      newSortedRdd.localCheckpoint(),
+      LocalRDDCheckpointData.DEFAULT_STORAGE_LEVEL,
+      50)
   }
 
   test("checkpoint without draining iterator - caching before checkpointing") {
@@ -135,7 +140,9 @@ class LocalCheckpointSuite extends SparkFunSuite with LocalSparkContext {
   }
 
   test("checkpoint blocks exist") {
-    testCheckpointBlocksExist(newRdd.localCheckpoint(), StorageLevel.DISK_ONLY)
+    testCheckpointBlocksExist(
+      newRdd.localCheckpoint(),
+      LocalRDDCheckpointData.DEFAULT_STORAGE_LEVEL)
   }
 
   test("checkpoint blocks exist - caching before checkpointing") {

From c2e111b3ee7be3efa7dc320a09f9fa80507dd772 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Mon, 27 Jul 2015 18:27:08 -0700
Subject: [PATCH 23/27] Address comments

---
 .../apache/spark/rdd/LocalCheckpointRDD.scala | 10 ++++----
 .../spark/rdd/LocalRDDCheckpointData.scala    |  2 +-
 .../main/scala/org/apache/spark/rdd/RDD.scala | 23 ++++++++----------
 .../spark/rdd/ReliableCheckpointRDD.scala     | 24 ++++---------------
 4 files changed, 20 insertions(+), 39 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index c80548824d27c..f76ac50b6fe93 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -32,20 +32,20 @@ import org.apache.spark.storage.RDDBlockId
  *
  * @param sc the active SparkContext
  * @param rddId the ID of the checkpointed RDD
- * @param partitionIndices the partition indices of the checkpointed RDD
+ * @param numPartitions the number of partitions in the checkpointed RDD
  */
 private[spark] class LocalCheckpointRDD[T: ClassTag](
     @transient sc: SparkContext,
     rddId: Int,
-    partitionIndices: Array[Int])
+    numPartitions: Int)
   extends CheckpointRDD[T](sc) {
 
   def this(rdd: RDD[T]) {
-    this(rdd.context, rdd.id, rdd.partitions.map(_.index))
+    this(rdd.context, rdd.id, rdd.partitions.size)
   }
 
   protected override def getPartitions: Array[Partition] = {
-    partitionIndices.map { i => new CheckpointRDDPartition(i) }
+    (0 until partitionIndices).map { i => new CheckpointRDDPartition(i) }
   }
 
   /**
@@ -59,7 +59,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
   override def compute(partition: Partition, context: TaskContext): Iterator[T] = {
     throw new SparkException(
       s"Checkpoint block ${RDDBlockId(rddId, partition.index)} not found! Either the executor " +
-      s"that originally checkpointed this block is no longer alive, or the original RDD is " +
+      s"that originally checkpointed this partition is no longer alive, or the original RDD is " +
       s"unpersisted. If this problem persists, you may consider using `rdd.checkpoint()` " +
       s"instead, which is slower than local checkpointing but more fault-tolerant.")
   }
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index b7f86c4145dd2..53d711bc2b149 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -26,7 +26,7 @@ import org.apache.spark.storage.{RDDBlockId, StorageLevel}
  * An implementation of checkpointing implemented on top of Spark's caching layer.
  *
  * Local checkpointing trades off fault tolerance for performance by skipping the expensive
- * step of replicating the checkpointed data in a reliable storage. Instead, checkpoint data
+ * step of saving the RDD data to a reliable and fault-tolerant storage. Instead, the data
  * is written to the local, ephemeral block storage that lives in each executor. This is useful
  * for use cases where RDDs build up long lineages that need to be truncated often (e.g. GraphX).
  */
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index 99883c4381264..aa1da48c3ab5d 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1473,16 +1473,14 @@ abstract class RDD[T: ClassTag](
    * executed on this RDD. It is strongly recommended that this RDD is persisted in
    * memory, otherwise saving it on a file will require recomputation.
    */
-  def checkpoint(): Unit = {
+  def checkpoint(): Unit = RDDCheckpointData.synchronized {
+    // NOTE: we use a global lock here due to complexities downstream with ensuring
+    // children RDD partitions point to the correct parent partitions. In the future
+    // we should revisit this consideration.
     if (context.checkpointDir.isEmpty) {
       throw new SparkException("Checkpoint directory has not been set in the SparkContext")
     } else if (checkpointData.isEmpty) {
-      // NOTE: we use a global lock here due to complexities downstream with ensuring
-      // children RDD partitions point to the correct parent partitions. In the future
-      // we should revisit this consideration.
-      RDDCheckpointData.synchronized {
-        checkpointData = Some(new ReliableRDDCheckpointData(this))
-      }
+      checkpointData = Some(new ReliableRDDCheckpointData(this))
     }
   }
 
@@ -1493,17 +1491,16 @@ abstract class RDD[T: ClassTag](
    * step of replicating the materialized data in a reliable distributed file system. This is
    * useful for RDDs with long lineages that need to be truncated periodically (e.g. GraphX).
    *
-   * Local checkpointing sacrifices fault-tolerance for performance. In particular,
-   * checkpointed data is written to ephemeral local storage (memory or disk) instead of
-   * to a reliable distributed storage. The effect is that if an executor fails during the
-   * computation, the checkpointed data may no longer be accessible, causing an irrecoverable
-   * job failure.
+   * Local checkpointing sacrifices fault-tolerance for performance. In particular, checkpointed
+   * data is written to ephemeral local storage in the executors instead of to a reliable,
+   * fault-tolerant storage. The effect is that if an executor fails during the computation,
+   * the checkpointed data may no longer be accessible, causing an irrecoverable job failure.
    *
    * This is NOT safe to use with dynamic allocation, which removes executors along
    * with their cached blocks. If you must use both features, you are advised to set
    * `spark.dynamicAllocation.cachedExecutorIdleTimeout` to a high value.
    *
-   * The checkpoint directory set through `SparkContext#setCheckpointDir` is not read.
+   * The checkpoint directory set through `SparkContext#setCheckpointDir` is not used.
    */
   def localCheckpoint(): this.type = RDDCheckpointData.synchronized {
     if (conf.getBoolean("spark.dynamicAllocation.enabled", false) &&
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index 2dc5cd85a1111..9838447cb0670 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -40,6 +40,9 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
   @transient private val fs = new Path(checkpointPath).getFileSystem(hadoopConf)
   private val broadcastedConf = sc.broadcast(new SerializableConfiguration(hadoopConf))
 
+  // Fail fast if checkpoint directory does not exist
+  require(fs.exists(checkpointPath), s"Checkpoint directory does not exist: $checkpointPath")
+
   /**
    * Return the path of the checkpoint directory this RDD reads data from.
    */
@@ -143,7 +146,7 @@ private[spark] object ReliableCheckpointRDD extends Logging {
         logInfo(s"Deleting tempOutputPath $tempOutputPath")
         fs.delete(tempOutputPath, false)
         throw new IOException("Checkpoint failed: failed to save output of task: " +
-          s"${ctx.attemptNumber()} and final output path does not exist")
+          s"${ctx.attemptNumber()} and final output path does not exist: $finalOutputPath")
       } else {
         // Some other copy of this task must've finished before us and renamed it
         logInfo(s"Final output path $finalOutputPath already exists; not overwriting it")
@@ -172,23 +175,4 @@ private[spark] object ReliableCheckpointRDD extends Logging {
     deserializeStream.asIterator.asInstanceOf[Iterator[T]]
   }
 
-  // Test whether CheckpointRDD generate expected number of partitions despite
-  // each split file having multiple blocks. This needs to be run on a
-  // cluster (mesos or standalone) using HDFS.
-  def main(args: Array[String]) {
-    val Array(cluster, hdfsPath) = args
-    val sc = new SparkContext(cluster, "CheckpointRDD Test")
-    val rdd = sc.makeRDD(1 to 10, 10).flatMap(x => 1 to 10000)
-    val path = new Path(hdfsPath, "temp")
-    val conf = SparkHadoopUtil.get.newConfiguration(new SparkConf())
-    val fs = path.getFileSystem(conf)
-    val broadcastedConf = sc.broadcast(new SerializableConfiguration(conf))
-    sc.runJob(
-      rdd, ReliableCheckpointRDD.writeCheckpointFile[Int](path.toString, broadcastedConf, 1024) _)
-    val cpRDD = new ReliableCheckpointRDD[Int](sc, path.toString)
-    assert(cpRDD.partitions.length == rdd.partitions.length, "Number of partitions is not the same")
-    assert(cpRDD.collect().toList == rdd.collect().toList, "Data of partitions not the same")
-    fs.delete(path, true)
-  }
-
 }

From ab003a38ca764d059af390796e3ff19287c7d350 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Mon, 27 Jul 2015 20:41:23 -0700
Subject: [PATCH 24/27] Fix compile

---
 .../scala/org/apache/spark/TaskContext.scala  |  2 +-
 .../apache/spark/rdd/LocalCheckpointRDD.scala |  2 +-
 .../spark/rdd/ReliableCheckpointRDD.scala     | 34 ++++++++-----------
 3 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/TaskContext.scala b/core/src/main/scala/org/apache/spark/TaskContext.scala
index bf65a787da048..5d2c551d58514 100644
--- a/core/src/main/scala/org/apache/spark/TaskContext.scala
+++ b/core/src/main/scala/org/apache/spark/TaskContext.scala
@@ -65,7 +65,7 @@ object TaskContext {
    * Internal use only.
    */
   private[spark] def empty(): TaskContext = {
-    new TaskContextImpl(0, 0, 0, 0, null, false)
+    new TaskContextImpl(0, 0, 0, 0, null, null)
   }
 }
 
diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
index f76ac50b6fe93..daa5779d688cc 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalCheckpointRDD.scala
@@ -45,7 +45,7 @@ private[spark] class LocalCheckpointRDD[T: ClassTag](
   }
 
   protected override def getPartitions: Array[Partition] = {
-    (0 until partitionIndices).map { i => new CheckpointRDDPartition(i) }
+    (0 until numPartitions).toArray.map { i => new CheckpointRDDPartition(i) }
   }
 
   /**
diff --git a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
index 9838447cb0670..35d8b0bfd18c5 100644
--- a/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/ReliableCheckpointRDD.scala
@@ -37,11 +37,12 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
   extends CheckpointRDD[T](sc) {
 
   @transient private val hadoopConf = sc.hadoopConfiguration
-  @transient private val fs = new Path(checkpointPath).getFileSystem(hadoopConf)
+  @transient private val cpath = new Path(checkpointPath)
+  @transient private val fs = cpath.getFileSystem(hadoopConf)
   private val broadcastedConf = sc.broadcast(new SerializableConfiguration(hadoopConf))
 
   // Fail fast if checkpoint directory does not exist
-  require(fs.exists(checkpointPath), s"Checkpoint directory does not exist: $checkpointPath")
+  require(fs.exists(cpath), s"Checkpoint directory does not exist: $checkpointPath")
 
   /**
    * Return the path of the checkpoint directory this RDD reads data from.
@@ -56,25 +57,18 @@ private[spark] class ReliableCheckpointRDD[T: ClassTag](
    * checkpoint files are fully preserved in a reliable storage across application lifespans.
    */
   protected override def getPartitions: Array[Partition] = {
-    val cpath = new Path(checkpointPath)
-    val numPartitions =
-      // listStatus can throw exception if path does not exist.
-      if (fs.exists(cpath)) {
-        val inputFiles = fs.listStatus(cpath)
-          .map(_.getPath)
-          .filter(_.getName.startsWith("part-"))
-          .sortBy(_.toString)
-        // Fail fast if input files are invalid
-        inputFiles.zipWithIndex.foreach { case (path, i) =>
-          if (!path.toString.endsWith(ReliableCheckpointRDD.checkpointFileName(i))) {
-            throw new SparkException(s"Invalid checkpoint file: $path")
-          }
-        }
-        inputFiles.length
-      } else {
-        0
+    // listStatus can throw exception if path does not exist.
+    val inputFiles = fs.listStatus(cpath)
+      .map(_.getPath)
+      .filter(_.getName.startsWith("part-"))
+      .sortBy(_.toString)
+    // Fail fast if input files are invalid
+    inputFiles.zipWithIndex.foreach { case (path, i) =>
+      if (!path.toString.endsWith(ReliableCheckpointRDD.checkpointFileName(i))) {
+        throw new SparkException(s"Invalid checkpoint file: $path")
       }
-    Array.tabulate(numPartitions)(i => new CheckpointRDDPartition(i))
+    }
+    Array.tabulate(inputFiles.length)(i => new CheckpointRDDPartition(i))
   }
 
   /**

From 3be5aea1d93937251d0b42d7200a2421f08b9f4a Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Sat, 1 Aug 2015 22:26:20 -0700
Subject: [PATCH 25/27] Address comments

---
 .../spark/rdd/LocalRDDCheckpointData.scala    |  10 +-
 .../main/scala/org/apache/spark/rdd/RDD.scala |   9 +-
 .../org/apache/spark/CheckpointSuite.scala    | 118 +++++++++---------
 3 files changed, 70 insertions(+), 67 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 53d711bc2b149..5e7b0b55c9607 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -37,8 +37,6 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
    * Ensure the RDD is fully cached so the partitions can be recovered later.
    */
   protected override def doCheckpoint(): CheckpointRDD[T] = {
-    val cm = SparkEnv.get.cacheManager
-    val bmm = SparkEnv.get.blockManager.master
     val level = rdd.getStorageLevel
 
     // Assume storage level uses disk; otherwise memory eviction may cause data loss
@@ -46,12 +44,8 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
 
     // Not all RDD actions compute all partitions of the RDD (e.g. take)
     // We must compute and cache any missing partitions for correctness reasons
-    rdd.partitions.foreach { p =>
-      val blockId = RDDBlockId(rdd.id, p.index)
-      if (!bmm.contains(blockId)) {
-        cm.getOrCompute(rdd, p, TaskContext.empty(), level)
-      }
-    }
+    // TODO: avoid running another job here (SPARK-8582)
+    rdd.count()
 
     new LocalCheckpointRDD[T](rdd)
   }
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
index aa1da48c3ab5d..081c721f23687 100644
--- a/core/src/main/scala/org/apache/spark/rdd/RDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -1508,7 +1508,8 @@ abstract class RDD[T: ClassTag](
       logWarning("Local checkpointing is NOT safe to use with dynamic allocation, " +
         "which removes executors along with their cached blocks. If you must use both " +
         "features, you are advised to set `spark.dynamicAllocation.cachedExecutorIdleTimeout` " +
-        "to a high value.")
+        "to a high value. E.g. If you plan to use the RDD for 1 hour, set the timeout to " +
+        "at least 1 hour.")
     }
 
     // Note: At this point we do not actually know whether the user will call persist() on
@@ -1524,6 +1525,12 @@ abstract class RDD[T: ClassTag](
     } else {
       persist(LocalRDDCheckpointData.transformStorageLevel(storageLevel), allowOverride = true)
     }
+
+    checkpointData match {
+      case Some(reliable: ReliableRDDCheckpointData[_]) => logWarning(
+        "RDD was already marked for reliable checkpointing: overriding with local checkpoint.")
+      case _ =>
+    }
     checkpointData = Some(new LocalRDDCheckpointData(this))
     this
   }
diff --git a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
index c6bbfe8725bcd..d343bb95cb68c 100644
--- a/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
+++ b/core/src/test/scala/org/apache/spark/CheckpointSuite.scala
@@ -27,7 +27,7 @@ import org.apache.spark.util.Utils
 
 /**
  * Test suite for end-to-end checkpointing functionality.
- * This tests both normal checkpoints and local checkpoints.
+ * This tests both reliable checkpoints and local checkpoints.
  */
 class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging {
   private var checkpointDir: File = _
@@ -46,35 +46,36 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     Utils.deleteRecursively(checkpointDir)
   }
 
-  runTest("basic checkpointing") { normal: Boolean =>
+  runTest("basic checkpointing") { reliableCheckpoint: Boolean =>
     val parCollection = sc.makeRDD(1 to 4)
     val flatMappedRDD = parCollection.flatMap(x => 1 to x)
-    checkpoint(flatMappedRDD, normal)
+    checkpoint(flatMappedRDD, reliableCheckpoint)
     assert(flatMappedRDD.dependencies.head.rdd === parCollection)
     val result = flatMappedRDD.collect()
     assert(flatMappedRDD.dependencies.head.rdd != parCollection)
     assert(flatMappedRDD.collect() === result)
   }
 
-  runTest("RDDs with one-to-one dependencies") { normal: Boolean =>
-    testRDD(_.map(x => x.toString), normal)
-    testRDD(_.flatMap(x => 1 to x), normal)
-    testRDD(_.filter(_ % 2 == 0), normal)
-    testRDD(_.sample(false, 0.5, 0), normal)
-    testRDD(_.glom(), normal)
-    testRDD(_.mapPartitions(_.map(_.toString)), normal)
-    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).mapValues(_.toString), normal)
-    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).flatMapValues(x => 1 to x), normal)
-    testRDD(_.pipe(Seq("cat")), normal)
+  runTest("RDDs with one-to-one dependencies") { reliableCheckpoint: Boolean =>
+    testRDD(_.map(x => x.toString), reliableCheckpoint)
+    testRDD(_.flatMap(x => 1 to x), reliableCheckpoint)
+    testRDD(_.filter(_ % 2 == 0), reliableCheckpoint)
+    testRDD(_.sample(false, 0.5, 0), reliableCheckpoint)
+    testRDD(_.glom(), reliableCheckpoint)
+    testRDD(_.mapPartitions(_.map(_.toString)), reliableCheckpoint)
+    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).mapValues(_.toString), reliableCheckpoint)
+    testRDD(_.map(x => (x % 2, 1)).reduceByKey(_ + _).flatMapValues(x => 1 to x),
+      reliableCheckpoint)
+    testRDD(_.pipe(Seq("cat")), reliableCheckpoint)
   }
 
-  runTest("ParallelCollectionRDD") { normal: Boolean =>
+  runTest("ParallelCollectionRDD") { reliableCheckpoint: Boolean =>
     val parCollection = sc.makeRDD(1 to 4, 2)
     val numPartitions = parCollection.partitions.size
-    checkpoint(parCollection, normal)
+    checkpoint(parCollection, reliableCheckpoint)
     assert(parCollection.dependencies === Nil)
     val result = parCollection.collect()
-    if (normal) {
+    if (reliableCheckpoint) {
       assert(sc.checkpointFile[Int](parCollection.getCheckpointFile.get).collect() === result)
     }
     assert(parCollection.dependencies != Nil)
@@ -84,15 +85,15 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     assert(parCollection.collect() === result)
   }
 
-  runTest("BlockRDD") { normal: Boolean =>
+  runTest("BlockRDD") { reliableCheckpoint: Boolean =>
     val blockId = TestBlockId("id")
     val blockManager = SparkEnv.get.blockManager
     blockManager.putSingle(blockId, "test", StorageLevel.MEMORY_ONLY)
     val blockRDD = new BlockRDD[String](sc, Array(blockId))
     val numPartitions = blockRDD.partitions.size
-    checkpoint(blockRDD, normal)
+    checkpoint(blockRDD, reliableCheckpoint)
     val result = blockRDD.collect()
-    if (normal) {
+    if (reliableCheckpoint) {
       assert(sc.checkpointFile[String](blockRDD.getCheckpointFile.get).collect() === result)
     }
     assert(blockRDD.dependencies != Nil)
@@ -101,29 +102,29 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     assert(blockRDD.collect() === result)
   }
 
-  runTest("ShuffleRDD") { normal: Boolean =>
+  runTest("ShuffleRDD") { reliableCheckpoint: Boolean =>
     testRDD(rdd => {
       // Creating ShuffledRDD directly as PairRDDFunctions.combineByKey produces a MapPartitionedRDD
       new ShuffledRDD[Int, Int, Int](rdd.map(x => (x % 2, 1)), partitioner)
-    }, normal)
+    }, reliableCheckpoint)
   }
 
-  runTest("UnionRDD") { normal: Boolean =>
+  runTest("UnionRDD") { reliableCheckpoint: Boolean =>
     def otherRDD: RDD[Int] = sc.makeRDD(1 to 10, 1)
-    testRDD(_.union(otherRDD), normal)
-    testRDDPartitions(_.union(otherRDD), normal)
+    testRDD(_.union(otherRDD), reliableCheckpoint)
+    testRDDPartitions(_.union(otherRDD), reliableCheckpoint)
   }
 
-  runTest("CartesianRDD") { normal: Boolean =>
+  runTest("CartesianRDD") { reliableCheckpoint: Boolean =>
     def otherRDD: RDD[Int] = sc.makeRDD(1 to 10, 1)
-    testRDD(new CartesianRDD(sc, _, otherRDD), normal)
-    testRDDPartitions(new CartesianRDD(sc, _, otherRDD), normal)
+    testRDD(new CartesianRDD(sc, _, otherRDD), reliableCheckpoint)
+    testRDDPartitions(new CartesianRDD(sc, _, otherRDD), reliableCheckpoint)
 
     // Test that the CartesianRDD updates parent partitions (CartesianRDD.s1/s2) after
     // the parent RDD has been checkpointed and parent partitions have been changed.
     // Note that this test is very specific to the current implementation of CartesianRDD.
     val ones = sc.makeRDD(1 to 100, 10).map(x => x)
-    checkpoint(ones, normal) // checkpoint that MappedRDD
+    checkpoint(ones, reliableCheckpoint) // checkpoint that MappedRDD
     val cartesian = new CartesianRDD(sc, ones, ones)
     val splitBeforeCheckpoint =
       serializeDeserialize(cartesian.partitions.head.asInstanceOf[CartesianPartition])
@@ -137,16 +138,16 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  runTest("CoalescedRDD") { normal: Boolean =>
-    testRDD(_.coalesce(2), normal)
-    testRDDPartitions(_.coalesce(2), normal)
+  runTest("CoalescedRDD") { reliableCheckpoint: Boolean =>
+    testRDD(_.coalesce(2), reliableCheckpoint)
+    testRDDPartitions(_.coalesce(2), reliableCheckpoint)
 
     // Test that the CoalescedRDDPartition updates parent partitions (CoalescedRDDPartition.parents)
     // after the parent RDD has been checkpointed and parent partitions have been changed.
     // Note that this test is very specific to the current implementation of
     // CoalescedRDDPartitions.
     val ones = sc.makeRDD(1 to 100, 10).map(x => x)
-    checkpoint(ones, normal) // checkpoint that MappedRDD
+    checkpoint(ones, reliableCheckpoint) // checkpoint that MappedRDD
     val coalesced = new CoalescedRDD(ones, 2)
     val splitBeforeCheckpoint =
       serializeDeserialize(coalesced.partitions.head.asInstanceOf[CoalescedRDDPartition])
@@ -159,7 +160,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  runTest("CoGroupedRDD") { normal: Boolean =>
+  runTest("CoGroupedRDD") { reliableCheckpoint: Boolean =>
     val longLineageRDD1 = generateFatPairRDD()
 
     // Collect the RDD as sequences instead of arrays to enable equality tests in testRDD
@@ -168,26 +169,26 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
 
     testRDD(rdd => {
       CheckpointSuite.cogroup(longLineageRDD1, rdd.map(x => (x % 2, 1)), partitioner)
-    }, normal, seqCollectFunc)
+    }, reliableCheckpoint, seqCollectFunc)
 
     val longLineageRDD2 = generateFatPairRDD()
     testRDDPartitions(rdd => {
       CheckpointSuite.cogroup(
         longLineageRDD2, sc.makeRDD(1 to 2, 2).map(x => (x % 2, 1)), partitioner)
-    }, normal, seqCollectFunc)
+    }, reliableCheckpoint, seqCollectFunc)
   }
 
-  runTest("ZippedPartitionsRDD") { normal: Boolean =>
-    testRDD(rdd => rdd.zip(rdd.map(x => x)), normal)
-    testRDDPartitions(rdd => rdd.zip(rdd.map(x => x)), normal)
+  runTest("ZippedPartitionsRDD") { reliableCheckpoint: Boolean =>
+    testRDD(rdd => rdd.zip(rdd.map(x => x)), reliableCheckpoint)
+    testRDDPartitions(rdd => rdd.zip(rdd.map(x => x)), reliableCheckpoint)
 
     // Test that ZippedPartitionsRDD updates parent partitions after parent RDDs have
     // been checkpointed and parent partitions have been changed.
     // Note that this test is very specific to the implementation of ZippedPartitionsRDD.
     val rdd = generateFatRDD()
     val zippedRDD = rdd.zip(rdd.map(x => x)).asInstanceOf[ZippedPartitionsRDD2[_, _, _]]
-    checkpoint(zippedRDD.rdd1, normal)
-    checkpoint(zippedRDD.rdd2, normal)
+    checkpoint(zippedRDD.rdd1, reliableCheckpoint)
+    checkpoint(zippedRDD.rdd2, reliableCheckpoint)
     val partitionBeforeCheckpoint =
       serializeDeserialize(zippedRDD.partitions.head.asInstanceOf[ZippedPartitionsPartition])
     zippedRDD.count()
@@ -202,27 +203,27 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  runTest("PartitionerAwareUnionRDD") { normal: Boolean =>
+  runTest("PartitionerAwareUnionRDD") { reliableCheckpoint: Boolean =>
     testRDD(rdd => {
       new PartitionerAwareUnionRDD[(Int, Int)](sc, Array(
         generateFatPairRDD(),
         rdd.map(x => (x % 2, 1)).reduceByKey(partitioner, _ + _)
       ))
-    }, normal)
+    }, reliableCheckpoint)
 
     testRDDPartitions(rdd => {
       new PartitionerAwareUnionRDD[(Int, Int)](sc, Array(
         generateFatPairRDD(),
         rdd.map(x => (x % 2, 1)).reduceByKey(partitioner, _ + _)
       ))
-    }, normal)
+    }, reliableCheckpoint)
 
     // Test that the PartitionerAwareUnionRDD updates parent partitions
     // (PartitionerAwareUnionRDD.parents) after the parent RDD has been checkpointed and parent
     // partitions have been changed. Note that this test is very specific to the current
     // implementation of PartitionerAwareUnionRDD.
     val pairRDD = generateFatPairRDD()
-    checkpoint(pairRDD, normal)
+    checkpoint(pairRDD, reliableCheckpoint)
     val unionRDD = new PartitionerAwareUnionRDD(sc, Array(pairRDD))
     val partitionBeforeCheckpoint = serializeDeserialize(
       unionRDD.partitions.head.asInstanceOf[PartitionerAwareUnionRDDPartition])
@@ -236,11 +237,11 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     )
   }
 
-  runTest("CheckpointRDD with zero partitions") { normal: Boolean =>
+  runTest("CheckpointRDD with zero partitions") { reliableCheckpoint: Boolean =>
     val rdd = new BlockRDD[Int](sc, Array[BlockId]())
     assert(rdd.partitions.size === 0)
     assert(rdd.isCheckpointed === false)
-    checkpoint(rdd, normal)
+    checkpoint(rdd, reliableCheckpoint)
     assert(rdd.count() === 0)
     assert(rdd.isCheckpointed === true)
     assert(rdd.partitions.size === 0)
@@ -248,19 +249,19 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
 
   // Utility test methods
 
-  /** Checkpoint the RDD either locally or normally. */
-  private def checkpoint(rdd: RDD[_], normal: Boolean): Unit = {
-    if (normal) {
+  /** Checkpoint the RDD either locally or reliably. */
+  private def checkpoint(rdd: RDD[_], reliableCheckpoint: Boolean): Unit = {
+    if (reliableCheckpoint) {
       rdd.checkpoint()
     } else {
       rdd.localCheckpoint()
     }
   }
 
-  /** Run a test twice, once for local checkpointing and once for normal checkpointing. */
+  /** Run a test twice, once for local checkpointing and once for reliable checkpointing. */
   private def runTest(name: String)(body: Boolean => Unit): Unit = {
-    test(name + " [normal checkpointing]")(body(true))
-    test(name + " [local checkpointing]")(body(false))
+    test(name + " [reliable checkpoint]")(body(true))
+    test(name + " [local checkpoint]")(body(false))
   }
 
   private def defaultCollectFunc[T](rdd: RDD[T]): Any = rdd.collect()
@@ -271,13 +272,13 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * on all RDDs that have a parent RDD (i.e., do not call on ParallelCollection, BlockRDD, etc.).
    *
    * @param op an operation to run on the RDD
+   * @param reliableCheckpoint if true, use reliable checkpoints, otherwise use local checkpoints
    * @param collectFunc a function for collecting the values in the RDD, in case there are
    *   non-comparable types like arrays that we want to convert to something that supports ==
-   * @param normal if true, use normal checkpoints, otherwise use local checkpoints
    */
   private def testRDD[U: ClassTag](
       op: (RDD[Int]) => RDD[U],
-      normal: Boolean,
+      reliableCheckpoint: Boolean,
       collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
     // Generate the final RDD using given RDD operation
     val baseRDD = generateFatRDD()
@@ -295,14 +296,14 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     // Find serialized sizes before and after the checkpoint
     logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
     val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
-    checkpoint(operatedRDD, normal)
+    checkpoint(operatedRDD, reliableCheckpoint)
     val result = collectFunc(operatedRDD)
     operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables
     val (rddSizeAfterCheckpoint, partitionSizeAfterCheckpoint) = getSerializedSizes(operatedRDD)
     logInfo("RDD after checkpoint: " + operatedRDD + "\n" + operatedRDD.toDebugString)
 
     // Test whether the checkpoint file has been created
-    if (normal) {
+    if (reliableCheckpoint) {
       assert(collectFunc(sc.checkpointFile[U](operatedRDD.getCheckpointFile.get)) === result)
     }
 
@@ -340,12 +341,13 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
    * partitions (i.e., do not call it on simple RDD like MappedRDD).
    *
    * @param op an operation to run on the RDD
+   * @param reliableCheckpoint if true, use reliable checkpoints, otherwise use local checkpoints
    * @param collectFunc a function for collecting the values in the RDD, in case there are
    *   non-comparable types like arrays that we want to convert to something that supports ==
    */
   private def testRDDPartitions[U: ClassTag](
       op: (RDD[Int]) => RDD[U],
-      normal: Boolean,
+      reliableCheckpoint: Boolean,
       collectFunc: RDD[U] => Any = defaultCollectFunc[U] _): Unit = {
     // Generate the final RDD using given RDD operation
     val baseRDD = generateFatRDD()
@@ -362,7 +364,7 @@ class CheckpointSuite extends SparkFunSuite with LocalSparkContext with Logging
     val (rddSizeBeforeCheckpoint, partitionSizeBeforeCheckpoint) = getSerializedSizes(operatedRDD)
     // checkpoint the parent RDD, not the generated one
     parentRDDs.foreach { rdd =>
-      checkpoint(rdd, normal)
+      checkpoint(rdd, reliableCheckpoint)
     }
     val result = collectFunc(operatedRDD)  // force checkpointing
     operatedRDD.collect() // force re-initialization of post-checkpoint lazy variables

From 34bc05901b49ef46cd537242ccae4417cb015127 Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Sun, 2 Aug 2015 16:21:23 -0700
Subject: [PATCH 26/27] Avoid computing all partitions in local checkpoint

---
 .../apache/spark/rdd/LocalRDDCheckpointData.scala  | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 5e7b0b55c9607..8e83d4fdd9104 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -21,6 +21,7 @@ import scala.reflect.ClassTag
 
 import org.apache.spark.{Logging, SparkEnv, SparkException, TaskContext}
 import org.apache.spark.storage.{RDDBlockId, StorageLevel}
+import org.apache.spark.util.Utils
 
 /**
  * An implementation of checkpointing implemented on top of Spark's caching layer.
@@ -42,10 +43,15 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
     // Assume storage level uses disk; otherwise memory eviction may cause data loss
     assume(level.useDisk, s"Storage level $level is not appropriate for local checkpointing")
 
-    // Not all RDD actions compute all partitions of the RDD (e.g. take)
-    // We must compute and cache any missing partitions for correctness reasons
-    // TODO: avoid running another job here (SPARK-8582)
-    rdd.count()
+    // Not all actions compute all partitions of the RDD (e.g. take). For correctness, we
+    // must cache any missing partitions. TODO: avoid running another job here (SPARK-8582).
+    val action = (tc: TaskContext, iterator: Iterator[T]) => Utils.getIteratorSize(iterator)
+    val missingPartitionIndices = rdd.partitions.map(_.index).filter { i =>
+      SparkEnv.get.blockManager.master.contains(RDDBlockId(rdd.id, i))
+    }
+    if (missingPartitionIndices.nonEmpty) {
+      rdd.sparkContext.runJob(rdd, action, missingPartitionIndices)
+    }
 
     new LocalCheckpointRDD[T](rdd)
   }

From 729600f83072422824689fb22e5148ccccd211cf Mon Sep 17 00:00:00 2001
From: Andrew Or <andrew@databricks.com>
Date: Sun, 2 Aug 2015 19:23:51 -0700
Subject: [PATCH 27/27] Oops, fix tests

This proves that the test is valid!
---
 .../scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
index 8e83d4fdd9104..d6fad896845f6 100644
--- a/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/LocalRDDCheckpointData.scala
@@ -47,7 +47,7 @@ private[spark] class LocalRDDCheckpointData[T: ClassTag](@transient rdd: RDD[T])
     // must cache any missing partitions. TODO: avoid running another job here (SPARK-8582).
     val action = (tc: TaskContext, iterator: Iterator[T]) => Utils.getIteratorSize(iterator)
     val missingPartitionIndices = rdd.partitions.map(_.index).filter { i =>
-      SparkEnv.get.blockManager.master.contains(RDDBlockId(rdd.id, i))
+      !SparkEnv.get.blockManager.master.contains(RDDBlockId(rdd.id, i))
     }
     if (missingPartitionIndices.nonEmpty) {
       rdd.sparkContext.runJob(rdd, action, missingPartitionIndices)