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

enable train multiple models by distinguishing stage IDs (#1493)

Browse Source
This commit is contained in:
Nan Zhu 2016-08-20 16:37:07 -04:00 committed by GitHub
parent bdfa8c0e09
commit 582ee63e34
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -65,7 +65,8 @@ object XGBoost extends Serializable {
if (trainingSamples.hasNext) {
val cacheFileName: String = {
if (useExternalMemory && trainingSamples.hasNext) {
s"$appName-dtrain_cache-${TaskContext.getPartitionId()}"
s"$appName-${TaskContext.get().stageId()}-" +
s"dtrain_cache-${TaskContext.getPartitionId()}"
} else {
null
}

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

@ -42,7 +42,7 @@ class XGBoostModel(_booster: Booster) extends Serializable {
Rabit.init(rabitEnv.asJava)
val cacheFileName = {
if (useExternalCache) {
s"$appName-dtest_cache-${TaskContext.getPartitionId()}"
s"$appName-${TaskContext.get().stageId()}-dtest_cache-${TaskContext.getPartitionId()}"
} else {
null
}