diff --git a/dmlc-core b/dmlc-core index 4e6459b0b..1db0792e1 160000 --- a/dmlc-core +++ b/dmlc-core @@ -1 +1 @@ -Subproject commit 4e6459b0bc15e6cf9b315cc75e2e5495c03cd417 +Subproject commit 1db0792e1a55355b1f07699bba18c88ded996953 diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala index 5d1a05c17..0122445c6 100644 --- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala +++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala @@ -43,7 +43,16 @@ object XGBoost extends Serializable { rabitEnv: mutable.Map[String, String], numWorkers: Int, round: Int, obj: ObjectiveTrait, eval: EvalTrait): RDD[Booster] = { import DataUtils._ - trainingData.repartition(numWorkers).mapPartitions { + val partitionedData = { + if (numWorkers > trainingData.partitions.length) { + trainingData.repartition(numWorkers) + } else if (numWorkers < trainingData.partitions.length) { + trainingData.coalesce(numWorkers) + } else { + trainingData + } + } + partitionedData.mapPartitions { trainingSamples => rabitEnv.put("DMLC_TASK_ID", TaskContext.getPartitionId().toString) Rabit.init(rabitEnv.asJava) @@ -60,6 +69,8 @@ object XGBoost extends Serializable { * @param trainingData the trainingset represented as RDD * @param configMap Map containing the configuration entries * @param round the number of iterations + * @param nWorkers the number of xgboost workers, 0 by default which means that the number of + * workers equals to the partition number of trainingData RDD * @param obj the user-defined objective function, null by default * @param eval the user-defined evaluation function, null by default * @throws ml.dmlc.xgboost4j.java.XGBoostError when the model training is failed @@ -67,9 +78,22 @@ object XGBoost extends Serializable { */ @throws(classOf[XGBoostError]) def train(trainingData: RDD[LabeledPoint], configMap: Map[String, Any], round: Int, - obj: ObjectiveTrait = null, eval: EvalTrait = null): XGBoostModel = { - val numWorkers = trainingData.partitions.length + nWorkers: Int = 0, obj: ObjectiveTrait = null, eval: EvalTrait = null): XGBoostModel = { implicit val sc = trainingData.sparkContext + if (configMap.contains("nthread")) { + val nThread = configMap("nthread") + val coresPerTask = sc.getConf.get("spark.task.cpus", "1") + require(nThread.toString <= coresPerTask, + s"the nthread configuration ($nThread) must be no larger than " + + s"spark.task.cpus ($coresPerTask)") + } + val numWorkers = { + if (nWorkers > 0) { + nWorkers + } else { + trainingData.partitions.length + } + } val tracker = new RabitTracker(numWorkers) require(tracker.start(), "FAULT: Failed to start tracker") val boosters = buildDistributedBoosters(trainingData, configMap, diff --git a/jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostSuite.scala b/jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostSuite.scala index 98964f23e..91a12530a 100644 --- a/jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostSuite.scala +++ b/jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostSuite.scala @@ -27,15 +27,15 @@ import org.apache.spark.mllib.linalg.DenseVector import org.apache.spark.mllib.regression.LabeledPoint import org.apache.spark.rdd.RDD import org.apache.spark.{SparkConf, SparkContext} -import org.scalatest.{BeforeAndAfterAll, FunSuite} +import org.scalatest.{BeforeAndAfter, FunSuite} import ml.dmlc.xgboost4j.java.{DMatrix => JDMatrix, XGBoostError} import ml.dmlc.xgboost4j.scala.{DMatrix, EvalTrait} -class XGBoostSuite extends FunSuite with BeforeAndAfterAll { +class XGBoostSuite extends FunSuite with BeforeAndAfter { private implicit var sc: SparkContext = null - private val numWorker = 2 + private val numWorkers = 4 private class EvalError extends EvalTrait { @@ -79,13 +79,13 @@ class XGBoostSuite extends FunSuite with BeforeAndAfterAll { } } - override def beforeAll(): Unit = { + before { // build SparkContext val sparkConf = new SparkConf().setMaster("local[*]").setAppName("XGBoostSuite") sc = new SparkContext(sparkConf) } - override def afterAll(): Unit = { + after { if (sc != null) { sc.stop() } @@ -112,12 +112,12 @@ class XGBoostSuite extends FunSuite with BeforeAndAfterAll { sampleList.toList } - private def buildTrainingRDD(): RDD[LabeledPoint] = { + private def buildTrainingRDD(sparkContext: Option[SparkContext] = None): RDD[LabeledPoint] = { val sampleList = readFile(getClass.getResource("/agaricus.txt.train").getFile) - sc.parallelize(sampleList, numWorker) + sparkContext.getOrElse(sc).parallelize(sampleList, numWorkers) } - test("build RDD containing boosters") { + test("build RDD containing boosters with the specified worker number") { val trainingRDD = buildTrainingRDD() val testSet = readFile(getClass.getResource("/agaricus.txt.test").getFile).iterator import DataUtils._ @@ -127,13 +127,13 @@ class XGBoostSuite extends FunSuite with BeforeAndAfterAll { List("eta" -> "1", "max_depth" -> "2", "silent" -> "0", "objective" -> "binary:logistic").toMap, new scala.collection.mutable.HashMap[String, String], - numWorker, 2, null, null) + numWorkers = 2, round = 5, null, null) val boosterCount = boosterRDD.count() - assert(boosterCount === numWorker) + assert(boosterCount === 2) val boosters = boosterRDD.collect() for (booster <- boosters) { val predicts = booster.predict(testSetDMatrix, true) - assert(new EvalError().eval(predicts, testSetDMatrix) < 0.1) + assert(new EvalError().eval(predicts, testSetDMatrix) < 0.17) } } @@ -155,4 +155,20 @@ class XGBoostSuite extends FunSuite with BeforeAndAfterAll { val predicts = loadedXGBooostModel.predict(testSetDMatrix) assert(eval.eval(predicts, testSetDMatrix) < 0.1) } + + test("nthread configuration must be equal to spark.task.cpus") { + sc.stop() + sc = null + val sparkConf = new SparkConf().setMaster("local[*]").setAppName("XGBoostSuite"). + set("spark.task.cpus", "4") + val customSparkContext = new SparkContext(sparkConf) + // start another app + val trainingRDD = buildTrainingRDD(Some(customSparkContext)) + val paramMap = List("eta" -> "1", "max_depth" -> "2", "silent" -> "0", + "objective" -> "binary:logistic", "nthread" -> 6).toMap + intercept[IllegalArgumentException] { + XGBoost.train(trainingRDD, paramMap, 5) + } + customSparkContext.stop() + } }