Rabit update. (#5978)

* Remove parameter on JVM Packages.
This commit is contained in:
Jiaming Yuan
2020-08-11 09:17:32 +08:00
committed by GitHub
parent 0b2a26fa74
commit f93f1c03fc
4 changed files with 3 additions and 7 deletions

View File

@@ -381,7 +381,6 @@ object XGBoost extends Serializable {
val attempt = TaskContext.get().attemptNumber.toString
rabitEnv.put("DMLC_TASK_ID", taskId)
rabitEnv.put("DMLC_NUM_ATTEMPT", attempt)
rabitEnv.put("DMLC_WORKER_STOP_PROCESS_ON_ERROR", "false")
val numRounds = xgbExecutionParam.numRounds
val makeCheckpoint = xgbExecutionParam.checkpointParam.isDefined && taskId.toInt == 0
try {
@@ -997,4 +996,3 @@ private[spark] class LabeledPointGroupIterator(base: Iterator[XGBLabeledPoint])
group
}
}

View File

@@ -308,8 +308,7 @@ class XGBoostClassificationModel private[ml](
private val batchIterImpl = rowIterator.grouped($(inferBatchSize)).flatMap { batchRow =>
if (batchCnt == 0) {
val rabitEnv = Array(
"DMLC_TASK_ID" -> TaskContext.getPartitionId().toString,
"DMLC_WORKER_STOP_PROCESS_ON_ERROR" -> "false").toMap
"DMLC_TASK_ID" -> TaskContext.getPartitionId().toString).toMap
Rabit.init(rabitEnv.asJava)
}

View File

@@ -286,8 +286,7 @@ class XGBoostRegressionModel private[ml] (
private val batchIterImpl = rowIterator.grouped($(inferBatchSize)).flatMap { batchRow =>
if (batchCnt == 0) {
val rabitEnv = Array(
"DMLC_TASK_ID" -> TaskContext.getPartitionId().toString,
"DMLC_WORKER_STOP_PROCESS_ON_ERROR" -> "false").toMap
"DMLC_TASK_ID" -> TaskContext.getPartitionId().toString).toMap
Rabit.init(rabitEnv.asJava)
}