Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix bug: doing rabit call after finalize in spark prediction phase (#1420)

Browse Source
This commit is contained in:
Fangzhou 2016-07-29 12:11:20 +08:00 committed by Yuan (Terry) Tang
parent 328e8e4c69
commit a8adf16228
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -20,8 +20,9 @@ import org.apache.hadoop.fs.{Path, FileSystem}
import org.apache.spark.{TaskContext, SparkContext}
import org.apache.spark.mllib.linalg.{DenseVector, Vector}
import org.apache.spark.rdd.RDD
import ml.dmlc.xgboost4j.java.{DMatrix => JDMatrix}
import ml.dmlc.xgboost4j.scala.{DMatrix, Booster}
import ml.dmlc.xgboost4j.java.{Rabit, DMatrix => JDMatrix}
import ml.dmlc.xgboost4j.scala.{Booster, DMatrix}
import scala.collection.JavaConverters._
class XGBoostModel(_booster: Booster) extends Serializable {
@ -37,6 +38,8 @@ class XGBoostModel(_booster: Booster) extends Serializable {
val appName = testSet.context.appName
testSet.mapPartitions { testSamples =>
if (testSamples.hasNext) {
val rabitEnv = Array("DMLC_TASK_ID" -> TaskContext.getPartitionId().toString).toMap
Rabit.init(rabitEnv.asJava)
val cacheFileName = {
if (useExternalCache) {
s"$appName-dtest_cache-${TaskContext.getPartitionId()}"
@ -45,7 +48,9 @@ class XGBoostModel(_booster: Booster) extends Serializable {
}
}
val dMatrix = new DMatrix(new JDMatrix(testSamples, cacheFileName))
Iterator(broadcastBooster.value.predict(dMatrix))
val res = broadcastBooster.value.predict(dMatrix)
Rabit.shutdown()
Iterator(res)
} else {
Iterator()
}