test case for XGBoostSpark

jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala

@@ -19,14 +19,14 @@ package ml.dmlc.xgboost4j.scala.spark
 import scala.collection.immutable.HashMap
 
 import com.typesafe.config.Config
-import org.apache.spark.SparkContext
+import org.apache.spark.{TaskContext, SparkContext}
 import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.rdd.RDD
 
-import ml.dmlc.xgboost4j.java.{DMatrix => JDMatrix}
+import ml.dmlc.xgboost4j.java.{DMatrix => JDMatrix, Rabit, RabitTracker}
 import ml.dmlc.xgboost4j.scala.{XGBoost => SXGBoost, _}
 
-object XGBoost {
+object XGBoost extends Serializable {
 
   implicit def convertBoosterToXGBoostModel(booster: Booster): XGBoostModel = {
     new XGBoostModel(booster)
@@ -38,28 +38,43 @@ object XGBoost {
       numWorkers: Int, round: Int, obj: ObjectiveTrait, eval: EvalTrait): RDD[Booster] = {
     import DataUtils._
     val sc = trainingData.sparkContext
-    val dataUtilsBroadcast = sc.broadcast(DataUtils)
-    trainingData.repartition(numWorkers).mapPartitions {
-      trainingSamples =>
-        val dMatrix = new DMatrix(new JDMatrix(trainingSamples, null))
-        Iterator(SXGBoost.train(xgBoostConfMap, dMatrix, round,
-          watches = new HashMap[String, DMatrix], obj, eval))
-    }.cache()
+    val tracker = new RabitTracker(numWorkers)
+    if (tracker.start()) {
+      trainingData.repartition(numWorkers).mapPartitions {
+        trainingSamples =>
+          Rabit.init(new java.util.HashMap[String, String]() {
+            put("DMLC_TASK_ID", TaskContext.getPartitionId().toString)
+          })
+          val dMatrix = new DMatrix(new JDMatrix(trainingSamples, null))
+          val booster = SXGBoost.train(xgBoostConfMap, dMatrix, round,
+            watches = new HashMap[String, DMatrix], obj, eval)
+          Rabit.shutdown()
+          Iterator(booster)
+      }.cache()
+    } else {
+      null
+    }
   }
 
   def train(config: Config, trainingData: RDD[LabeledPoint], obj: ObjectiveTrait = null,
-      eval: EvalTrait = null): XGBoostModel = {
+      eval: EvalTrait = null): Option[XGBoostModel] = {
     import DataUtils._
     val numWorkers = config.getInt("numWorkers")
     val round = config.getInt("round")
     val sc = trainingData.sparkContext
-    // TODO: build configuration map from config
-    val xgBoostConfigMap = new HashMap[String, AnyRef]()
-    val boosters = buildDistributedBoosters(trainingData, xgBoostConfigMap, numWorkers, round,
-      obj, eval)
-    // force the job
-    sc.runJob(boosters, (boosters: Iterator[Booster]) => boosters)
-    // TODO: how to choose best model
-    boosters.first()
+    val tracker = new RabitTracker(numWorkers)
+    if (tracker.start()) {
+      // TODO: build configuration map from config
+      val xgBoostConfigMap = new HashMap[String, AnyRef]()
+      val boosters = buildDistributedBoosters(trainingData, xgBoostConfigMap, numWorkers, round,
+        obj, eval)
+      // force the job
+      sc.runJob(boosters, (boosters: Iterator[Booster]) => boosters)
+      tracker.waitFor()
+      // TODO: how to choose best model
+      Some(boosters.first())
+    } else {
+      None
+    }
   }
 }

jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostModel.scala

@@ -33,4 +33,8 @@ class XGBoostModel(booster: Booster) extends Serializable {
       Iterator(broadcastBooster.value.predict(dMatrix))
     }
   }
+
+  def predict(testSet: DMatrix): Array[Array[Float]] = {
+    booster.predict(testSet)
+  }
 }