Some refactoring to make cluster scheduler pluggable.

Author: mateiz

core/src/main/scala/spark/PairRDDFunctions.scala

@@ -307,9 +307,12 @@ class PairRDDFunctions[K: ClassManifest, V: ClassManifest](
     val jobtrackerID = formatter.format(new Date())
     val stageId = self.id
     def writeShard(context: spark.TaskContext, iter: Iterator[(K,V)]): Int = {
+      // Hadoop wants a 32-bit task attempt ID, so if ours is bigger than Int.MaxValue, roll it
+      // around by taking a mod. We expect that no task will be attempted 2 billion times.
+      val attemptNumber = (context.attemptId % Int.MaxValue).toInt
       /* "reduce task" <split #> <attempt # = spark task #> */
       val attemptId = new TaskAttemptID(jobtrackerID,
-        stageId, false, context.splitId, context.attemptId)
+        stageId, false, context.splitId, attemptNumber)
       val hadoopContext = new TaskAttemptContext(wrappedConf.value, attemptId)
       val format = outputFormatClass.newInstance
       val committer = format.getOutputCommitter(hadoopContext)
@@ -371,7 +374,11 @@ class PairRDDFunctions[K: ClassManifest, V: ClassManifest](
     writer.preSetup()
 
     def writeToFile(context: TaskContext, iter: Iterator[(K,V)]) {
-      writer.setup(context.stageId, context.splitId, context.attemptId)
+      // Hadoop wants a 32-bit task attempt ID, so if ours is bigger than Int.MaxValue, roll it
+      // around by taking a mod. We expect that no task will be attempted 2 billion times.
+      val attemptNumber = (context.attemptId % Int.MaxValue).toInt
+
+      writer.setup(context.stageId, context.splitId, attemptNumber)
       writer.open()
 
       var count = 0

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

@@ -41,8 +41,8 @@ import spark.scheduler.ShuffleMapTask
 import spark.scheduler.DAGScheduler
 import spark.scheduler.TaskScheduler
 import spark.scheduler.local.LocalScheduler
+import spark.scheduler.cluster.ClusterScheduler
 import spark.scheduler.mesos.MesosScheduler
-import spark.scheduler.mesos.CoarseMesosScheduler
 import spark.storage.BlockManagerMaster
 
 class SparkContext(
@@ -89,11 +89,17 @@ class SparkContext(
         new LocalScheduler(threads.toInt, maxFailures.toInt)
       case _ =>
         MesosNativeLibrary.load()
+        val sched = new ClusterScheduler(this)
+        val schedContext = new MesosScheduler(sched, this, master, frameworkName)
+        sched.initialize(schedContext)
+        sched
+        /*
         if (System.getProperty("spark.mesos.coarse", "false") == "true") {
           new CoarseMesosScheduler(this, master, frameworkName)
         } else {
           new MesosScheduler(this, master, frameworkName)
         }
+        */
     }
   }
   taskScheduler.start()
@@ -272,11 +278,6 @@ class SparkContext(
     logInfo("Successfully stopped SparkContext")
   }
 
-  // Wait for the scheduler to be registered with the cluster manager
-  def waitForRegister() {
-    taskScheduler.waitForRegister()
-  }
-
   // Get Spark's home location from either a value set through the constructor,
   // or the spark.home Java property, or the SPARK_HOME environment variable
   // (in that order of preference). If neither of these is set, return None.

core/src/main/scala/spark/TaskContext.scala

@@ -1,3 +1,3 @@
 package spark
 
-class TaskContext(val stageId: Int, val splitId: Int, val attemptId: Int) extends Serializable
+class TaskContext(val stageId: Int, val splitId: Int, val attemptId: Long) extends Serializable

core/src/main/scala/spark/TaskState.scala

@@ -0,0 +1,32 @@
+package spark
+
+import org.apache.mesos.Protos.{TaskState => MesosTaskState}
+
+object TaskState
+  extends Enumeration("LAUNCHING", "RUNNING", "FINISHED", "FAILED", "KILLED", "LOST") {
+
+  val LAUNCHING, RUNNING, FINISHED, FAILED, KILLED, LOST = Value
+
+  type TaskState = Value
+
+  def isFinished(state: TaskState) = Seq(FINISHED, FAILED, LOST).contains(state)
+
+  def toMesos(state: TaskState): MesosTaskState = state match {
+    case LAUNCHING => MesosTaskState.TASK_STARTING
+    case RUNNING => MesosTaskState.TASK_RUNNING
+    case FINISHED => MesosTaskState.TASK_FINISHED
+    case FAILED => MesosTaskState.TASK_FAILED
+    case KILLED => MesosTaskState.TASK_KILLED
+    case LOST => MesosTaskState.TASK_LOST
+  }
+
+  def fromMesos(mesosState: MesosTaskState): TaskState = mesosState match {
+    case MesosTaskState.TASK_STAGING => LAUNCHING
+    case MesosTaskState.TASK_STARTING => LAUNCHING
+    case MesosTaskState.TASK_RUNNING => RUNNING
+    case MesosTaskState.TASK_FINISHED => FINISHED
+    case MesosTaskState.TASK_FAILED => FAILED
+    case MesosTaskState.TASK_KILLED => KILLED
+    case MesosTaskState.TASK_LOST => LOST
+  }
+}

core/src/main/scala/spark/Utils.scala

@@ -13,6 +13,7 @@ import scala.io.Source
  * Various utility methods used by Spark.
  */
 object Utils {
+  /** Serialize an object using Java serialization */
   def serialize[T](o: T): Array[Byte] = {
     val bos = new ByteArrayOutputStream()
     val oos = new ObjectOutputStream(bos)
@@ -21,12 +22,14 @@ object Utils {
     return bos.toByteArray
   }
 
+  /** Deserialize an object using Java serialization */
   def deserialize[T](bytes: Array[Byte]): T = {
     val bis = new ByteArrayInputStream(bytes)
     val ois = new ObjectInputStream(bis)
     return ois.readObject.asInstanceOf[T]
   }
 
+  /** Deserialize an object using Java serialization and the given ClassLoader */
   def deserialize[T](bytes: Array[Byte], loader: ClassLoader): T = {
     val bis = new ByteArrayInputStream(bytes)
     val ois = new ObjectInputStream(bis) {
@@ -106,6 +109,13 @@ object Utils {
     }
   }
 
+  /** Copy a file on the local file system */
+  def copyFile(source: File, dest: File) {
+    val in = new FileInputStream(source)
+    val out = new FileOutputStream(dest)
+    copyStream(in, out, true)
+  }
+
   /**
    * Shuffle the elements of a collection into a random order, returning the
    * result in a new collection. Unlike scala.util.Random.shuffle, this method

core/src/main/scala/spark/deploy/DeployMessage.scala

@@ -1,5 +1,7 @@
 package spark.deploy
 
+import spark.deploy.ExecutorState.ExecutorState
+
 sealed trait DeployMessage extends Serializable
 
 // Worker to Master
@@ -10,8 +12,7 @@ case class RegisterWorker(id: String, host: String, port: Int, cores: Int, memor
 case class ExecutorStateChanged(
     jobId: String,
     execId: Int,
-    state:
-    ExecutorState.Value,
+    state: ExecutorState,
     message: Option[String])
   extends DeployMessage
 
@@ -38,7 +39,7 @@ case class RegisterJob(jobDescription: JobDescription) extends DeployMessage
 
 case class RegisteredJob(jobId: String) extends DeployMessage
 case class ExecutorAdded(id: Int, workerId: String, host: String, cores: Int, memory: Int)
-case class ExecutorUpdated(id: Int, state: ExecutorState.Value, message: Option[String])
+case class ExecutorUpdated(id: Int, state: ExecutorState, message: Option[String])
 case class JobKilled(message: String)
 
 // Internal message in Client

core/src/main/scala/spark/deploy/ExecutorState.scala

@@ -5,5 +5,7 @@ object ExecutorState
 
   val LAUNCHING, LOADING, RUNNING, KILLED, FAILED, LOST = Value
 
-  def isFinished(state: Value): Boolean = (state == KILLED || state == FAILED || state == LOST)
+  type ExecutorState = Value
+
+  def isFinished(state: ExecutorState): Boolean = Seq(KILLED, FAILED, LOST).contains(state)
 }

core/src/main/scala/spark/deploy/JobDescription.scala

@@ -4,7 +4,6 @@ class JobDescription(
     val name: String,
     val cores: Int,
     val memoryPerSlave: Int,
-    val fileUrls: Seq[String],
     val command: Command)
   extends Serializable {
 

core/src/main/scala/spark/deploy/client/TestClient.scala

@@ -24,8 +24,8 @@ object TestClient {
   def main(args: Array[String]) {
     val url = args(0)
     val (actorSystem, port) = AkkaUtils.createActorSystem("spark", Utils.localIpAddress(), 0)
-    val desc = new JobDescription("TestClient", 1, 512, Seq(),
-      Command("spark.deploy.client.TestExecutor", Seq(), Map()))
+    val desc = new JobDescription(
+      "TestClient", 1, 512, Command("spark.deploy.client.TestExecutor", Seq(), Map()))
     val listener = new TestListener
     val client = new Client(actorSystem, url, desc, listener)
     client.start()

core/src/main/scala/spark/deploy/client/TestExecutor.scala

@@ -3,5 +3,8 @@ package spark.deploy.client
 object TestExecutor {
   def main(args: Array[String]) {
     println("Hello world!")
+    while (true) {
+      Thread.sleep(1000)
+    }
   }
 }