Support Cancellation of Spark Jobs

core/src/main/scala/spark/CacheManager.scala

@@ -19,7 +19,8 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
       case Some(cachedValues) =>
         // Partition is in cache, so just return its values
         logInfo("Found partition in cache!")
-        return cachedValues.asInstanceOf[Iterator[T]]
+        val iter = cachedValues.asInstanceOf[Iterator[T]]
+        return new InterruptibleIteratorDecorator(iter)
 
       case None =>
         // Mark the split as loading (unless someone else marks it first)
@@ -37,7 +38,8 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
             // downside of the current code is that threads wait serially if this does happen.
             blockManager.get(key) match {
               case Some(values) =>
-                return values.asInstanceOf[Iterator[T]]
+                val iter =  values.asInstanceOf[Iterator[T]]
+                return new InterruptibleIteratorDecorator(iter)
               case None =>
                 logInfo("Whoever was loading " + key + " failed; we'll try it ourselves")
                 loading.add(key)
@@ -53,7 +55,8 @@ private[spark] class CacheManager(blockManager: BlockManager) extends Logging {
           elements ++= rdd.computeOrReadCheckpoint(split, context)
           // Try to put this block in the blockManager
           blockManager.put(key, elements, storageLevel, true)
-          return elements.iterator.asInstanceOf[Iterator[T]]
+            val iter = elements.iterator.asInstanceOf[Iterator[T]]
+          return new InterruptibleIteratorDecorator(iter)
         } finally {
           loading.synchronized {
             loading.remove(key)

core/src/main/scala/spark/InterruptibleIterator.scala

@@ -0,0 +1,26 @@
+package spark
+
+trait InterruptibleIterator[+T] extends  Iterator[T]{
+
+  override def hasNext(): Boolean = {
+    if (!Thread.currentThread().isInterrupted()) {
+      true
+    } else {
+      throw new InterruptedException ("Thread interrupted during RDD iteration")
+    }
+  }
+
+}
+
+class InterruptibleIteratorDecorator[T](delegate: Iterator[T])
+	extends AnyRef with InterruptibleIterator[T] {
+
+  override def hasNext(): Boolean = {
+    super.hasNext
+    delegate.hasNext
+  }
+
+  override def next(): T = {
+    delegate.next()
+  }
+}

core/src/main/scala/spark/RDD.scala

@@ -81,7 +81,12 @@ abstract class RDD[T: ClassManifest](
   // =======================================================================
 
   /** Implemented by subclasses to compute a given partition. */
-  def compute(split: Partition, context: TaskContext): Iterator[T]
+  protected def compute(split: Partition, context: TaskContext): Iterator[T]
+
+  def computeInterruptibly(split: Partition, context: TaskContext): Iterator[T]	= {
+    val iter = compute(split, context)
+    new InterruptibleIteratorDecorator(iter)
+  }
 
   /**
    * Implemented by subclasses to return the set of partitions in this RDD. This method will only
@@ -229,7 +234,7 @@ abstract class RDD[T: ClassManifest](
     if (isCheckpointed) {
       firstParent[T].iterator(split, context)
     } else {
-      compute(split, context)
+      computeInterruptibly(split, context)
     }
   }
 

core/src/main/scala/spark/SparkContext.scala

@@ -617,6 +617,10 @@ class SparkContext(
     }
   }
 
+  def killJob(jobId: Int, reason: String="") {
+    dagScheduler.killJob(jobId, reason)
+  }
+
   /**
    * Run a function on a given set of partitions in an RDD and pass the results to the given
    * handler function. This is the main entry point for all actions in Spark. The allowLocal

core/src/main/scala/spark/executor/Executor.scala

@@ -2,16 +2,16 @@ package spark.executor
 
 import java.io.{File, FileOutputStream}
 import java.net.{URI, URL, URLClassLoader}
+import java.nio.ByteBuffer
 import java.util.concurrent._
-
+import scala.collection.JavaConversions._
+import scala.collection.mutable.{ArrayBuffer,ConcurrentMap, HashMap, Map}
+import scala.concurrent.JavaConversions._
 import org.apache.hadoop.fs.FileUtil
-
-import scala.collection.mutable.{ArrayBuffer, Map, HashMap}
-
 import spark.broadcast._
 import spark.scheduler._
 import spark._
-import java.nio.ByteBuffer
+import spark.scheduler.cluster.TaskDescription
 
 /**
  * The Mesos executor for Spark.
@@ -79,14 +79,26 @@ private[spark] class Executor(executorId: String, slaveHostname: String, propert
   val threadPool = new ThreadPoolExecutor(
     1, 128, 600, TimeUnit.SECONDS, new SynchronousQueue[Runnable])
 
+  val tasks: ConcurrentMap[Long, FutureTask[_]] = new ConcurrentHashMap[Long, FutureTask[_]]()
   def launchTask(context: ExecutorBackend, taskId: Long, serializedTask: ByteBuffer) {
-    threadPool.execute(new TaskRunner(context, taskId, serializedTask))
+    val runner = new TaskRunner(context, taskId, serializedTask)
+    val task = threadPool.submit(runner).asInstanceOf[FutureTask[_]]
+    tasks.put(taskId, task)
+
+  }
+
+  def killTask(context: ExecutorBackend, taskId: Long, executorId: String) {
+    val task = tasks.get(taskId)
+    task match {
+      case Some(t) => t.cancel(true)
+      case None => 
+    }
   }
 
   class TaskRunner(context: ExecutorBackend, taskId: Long, serializedTask: ByteBuffer)
     extends Runnable {
 
-    override def run() {
+    override def run(): Unit = {
       val startTime = System.currentTimeMillis()
       SparkEnv.set(env)
       Thread.currentThread.setContextClassLoader(replClassLoader)
@@ -138,6 +150,8 @@ private[spark] class Executor(executorId: String, slaveHostname: String, propert
           logError("Exception in task ID " + taskId, t)
           //System.exit(1)
         }
+      } finally {
+        tasks.remove(taskId)
       }
     }
   }

core/src/main/scala/spark/executor/StandaloneExecutorBackend.scala

@@ -55,6 +55,12 @@ private[spark] class StandaloneExecutorBackend(
       } else {
         executor.launchTask(this, taskDesc.taskId, taskDesc.serializedTask)
       }
+
+    case KillTask(taskId, executorId) => 
+      logInfo("Killing Task %s %s".format(taskId, executorId))
+      if (executor != null) {
+        executor.killTask(this, taskId, executorId)
+      }
 
     case Terminated(_) | RemoteClientDisconnected(_, _) | RemoteClientShutdown(_, _) =>
       logError("Driver terminated or disconnected! Shutting down.")

core/src/main/scala/spark/scheduler/DAGScheduler.scala

@@ -278,6 +278,41 @@ class DAGScheduler(
     return listener.awaitResult()    // Will throw an exception if the job fails
   }
 
+  def killJob(jobId: Int, reason: String)
+  {
+    logInfo("Killing Job %s".format(jobId))
+    val j = activeJobs.find(j => j.runId.equals(jobId))
+    j match {
+      case Some(job) => killJob(job, reason)
+      case None => Unit
+    }
+  }
+
+  private def killJob(job: ActiveJob, reason: String) {
+    logInfo("Killing Job and cleaning up stages %s".format(job.runId))
+    activeJobs.remove(job)
+    idToActiveJob.remove(job.runId)
+    val stage = job.finalStage
+    resultStageToJob.remove(stage)
+    killStage(stage)
+    // recursively remove all parent stages
+    stage.parents.foreach(p => killStage(p))
+    job.listener.jobFailed(new SparkException("Job failed: " + reason))
+  }
+
+  private def killStage(stage: Stage) {
+    logInfo("Killing Stage %s".format(stage.id))
+    idToStage.remove(stage.id)
+    if (stage.isShuffleMap) {
+      shuffleToMapStage.remove(stage.id)
+    }
+    waiting.remove(stage)
+    pendingTasks.remove(stage)
+    running.remove(stage)
+    taskSched.killTasks(stage.id)
+    stage.parents.foreach(p => killStage(p))
+  }
+
   /**
    * Process one event retrieved from the event queue.
    * Returns true if we should stop the event loop.
@@ -495,7 +530,11 @@ class DAGScheduler(
    */
   private def handleTaskCompletion(event: CompletionEvent) {
     val task = event.task
-    val stage = idToStage(task.stageId)
+    val stageId = task.stageId
+    if (!idToStage.contains(stageId)) {
+      return;
+    }
+    val stage = idToStage(stageId)
 
     def markStageAsFinished(stage: Stage) = {
       val serviceTime = stage.submissionTime match {

core/src/main/scala/spark/scheduler/ResultTask.scala

@@ -77,7 +77,7 @@ private[spark] class ResultTask[T, U](
     preferredLocs.foreach (hostPort => Utils.checkHost(Utils.parseHostPort(hostPort)._1, "preferredLocs : " + preferredLocs))
   }
 
-  override def run(attemptId: Long): U = {
+  override def runInterruptibly(attemptId: Long): U = {
     val context = new TaskContext(stageId, partition, attemptId)
     metrics = Some(context.taskMetrics)
     try {

core/src/main/scala/spark/scheduler/ShuffleMapTask.scala

@@ -98,6 +98,11 @@ private[spark] class ShuffleMapTask(
     rdd.partitions(partition)
   }
 
+  override def kill() {
+    logDebug("Killing Task %s %s".format(rdd.id, partition))
+    super.kill()
+  }
+
   override def writeExternal(out: ObjectOutput) {
     RDDCheckpointData.synchronized {
       split = rdd.partitions(partition)
@@ -124,7 +129,7 @@ private[spark] class ShuffleMapTask(
     split = in.readObject().asInstanceOf[Partition]
   }
 
-  override def run(attemptId: Long): MapStatus = {
+  override def runInterruptibly(attemptId: Long): MapStatus = {
     val numOutputSplits = dep.partitioner.numPartitions
 
     val taskContext = new TaskContext(stageId, partition, attemptId)
@@ -133,7 +138,6 @@ private[spark] class ShuffleMapTask(
     val blockManager = SparkEnv.get.blockManager
     var shuffle: ShuffleBlocks = null
     var buckets: ShuffleWriterGroup = null
-
     try {
       // Obtain all the block writers for shuffle blocks.
       val ser = SparkEnv.get.serializerManager.get(dep.serializerClass)

core/src/main/scala/spark/scheduler/Task.scala

@@ -1,23 +1,48 @@
 package spark.scheduler
 
-import spark.serializer.SerializerInstance
 import java.io.{DataInputStream, DataOutputStream}
 import java.nio.ByteBuffer
-import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
-import spark.util.ByteBufferInputStream
+import java.util.concurrent.{Callable, ExecutionException, Future, FutureTask}
 import scala.collection.mutable.HashMap
+import it.unimi.dsi.fastutil.io.FastByteArrayOutputStream
 import spark.executor.TaskMetrics
+import spark.serializer.SerializerInstance
+import spark.util.ByteBufferInputStream
+
 
 /**
  * A task to execute on a worker node.
  */
 private[spark] abstract class Task[T](val stageId: Int) extends Serializable {
-  def run(attemptId: Long): T
+  @volatile @transient var f: FutureTask[T] = null
+
+  def run(attemptId: Long): T = {
+    f = new FutureTask(new Callable[T] {
+      def call(): T = {
+        runInterruptibly(attemptId)
+      }
+    })
+    try {
+      f.run()
+      f.get()
+    } catch {
+      case e: Exception => throw e.getCause()
+    }
+  }
+
+  def runInterruptibly(attemptId: Long): T
+
   def preferredLocations: Seq[String] = Nil
 
   var generation: Long = -1   // Map output tracker generation. Will be set by TaskScheduler.
 
   var metrics: Option[TaskMetrics] = None
+
+  def kill(): Unit = {
+    if (f != null) {
+      f.cancel(true)
+    }
+  }
 
 }
 

core/src/main/scala/spark/scheduler/TaskScheduler.scala

@@ -19,6 +19,8 @@ private[spark] trait TaskScheduler {
 
   // Submit a sequence of tasks to run.
   def submitTasks(taskSet: TaskSet): Unit
+
+  def killTasks(stageId: Int): Unit
 
   // Set a listener for upcalls. This is guaranteed to be set before submitTasks is called.
   def setListener(listener: TaskSchedulerListener): Unit

core/src/main/scala/spark/scheduler/TaskSet.scala

@@ -15,4 +15,10 @@ private[spark] class TaskSet(
     val id: String = stageId + "." + attempt
 
   override def toString: String = "TaskSet " + id
+
+  def kill() = {
+    tasks.foreach {
+      _.kill()
+    }
+  }
 }

core/src/main/scala/spark/scheduler/cluster/ClusterScheduler.scala

@@ -199,6 +199,24 @@ private[spark] class ClusterScheduler(val sc: SparkContext)
     backend.reviveOffers()
   }
 
+  override def killTasks(stageId: Int) {
+    synchronized {
+      schedulableBuilder.popTaskSetManagers(stageId).foreach { 
+        t =>
+         val ts = t.asInstanceOf[TaskSetManager].taskSet
+         ts.kill()
+         val taskIds = taskSetTaskIds(ts.id)
+         taskIds.foreach {
+           tid => 
+             val execId = taskIdToExecutorId(tid)
+             backend.killTask(tid, execId)
+         }
+      }
+
+    }
+
+  }
+
   def taskSetFinished(manager: TaskSetManager) {
     this.synchronized {
       activeTaskSets -= manager.taskSet.id

core/src/main/scala/spark/scheduler/cluster/SchedulableBuilder.scala

@@ -22,6 +22,7 @@ import java.util.Properties
 private[spark] trait SchedulableBuilder {
   def buildPools()
   def addTaskSetManager(manager: Schedulable, properties: Properties)
+  def popTaskSetManagers(stageId: Int): Iterable[Schedulable]
 }
 
 private[spark] class FIFOSchedulableBuilder(val rootPool: Pool) extends SchedulableBuilder with Logging {
@@ -33,6 +34,16 @@ private[spark] class FIFOSchedulableBuilder(val rootPool: Pool) extends Schedula
   override def addTaskSetManager(manager: Schedulable, properties: Properties) {
     rootPool.addSchedulable(manager)
   }
+
+  override def popTaskSetManagers(stageId: Int) = {
+    val s = rootPool.schedulableNameToSchedulable.values.filter {
+      _.stageId == stageId
+    }
+    s.foreach {
+      rootPool.removeSchedulable(_)
+    }
+    s
+  }
 }
 
 private[spark] class FairSchedulableBuilder(val rootPool: Pool) extends SchedulableBuilder with Logging {
@@ -112,4 +123,14 @@ private[spark] class FairSchedulableBuilder(val rootPool: Pool) extends Schedula
     parentPool.addSchedulable(manager)
     logInfo("Added task set " + manager.name + " tasks to pool "+poolName)
   }
+
+  override def popTaskSetManagers(stageId: Int) = {
+    val s = rootPool.schedulableNameToSchedulable.values.filter {
+      _.stageId == stageId
+    }
+    s.foreach {
+      rootPool.removeSchedulable(_)
+    }
+    s
+  }
 }