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

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

View File

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

View File

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