update

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

@@ -139,12 +139,11 @@ private[this] class ParamsFactory(rawParams: Map[String, Any], sc: SparkContext)
     // back-compatible with "gpu_hist"
     val runOnGpu = treeMethod.exists(_ == "gpu_hist") || deviceIsGpu
 
-    val trackerConf = overridedParams.get("tracker_conf") match {
-      case None => TrackerConf()
-      case Some(conf: TrackerConf) => conf
-      case _ => throw new IllegalArgumentException("parameter \"tracker_conf\" must be an " +
-        "instance of TrackerConf.")
-    }
+    val trackerConf = TrackerConf(
+      overridedParams.get("rabit_tracker_timeout").asInstanceOf[Int],
+      overridedParams.get("rabit_tracker_host_ip").asInstanceOf[String],
+      overridedParams.get("rabit_tracker_port").asInstanceOf[Int],
+    )
 
     val earlyStoppingRounds = overridedParams.getOrElse(
       "num_early_stopping_rounds", 0).asInstanceOf[Int]
@@ -299,7 +298,7 @@ private[spark] trait StageLevelScheduling extends Serializable {
   }
 }
 
-private[spark] object NewXGBoost extends StageLevelScheduling {
+private[spark] object XGBoost extends StageLevelScheduling {
   private val logger = LogFactory.getLog("XGBoostSpark")
 
   def getGPUAddrFromResources: Int = {

jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostEstimator.scala

@@ -227,7 +227,7 @@ private[spark] abstract class XGBoostEstimator[
       "num_class" -> 3,
     )
 
-    val (booster, metrics) = NewXGBoost.train(
+    val (booster, metrics) = XGBoost.train(
       dataset.sparkSession.sparkContext, rdd, paramMap)
 
     val summary = XGBoostTrainingSummary(metrics)

jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/RabitParams.scala

@@ -0,0 +1,44 @@
+/*
+ Copyright (c) 2014-2022 by Contributors
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ */
+
+package ml.dmlc.xgboost4j.scala.spark.params
+
+import org.apache.spark.ml.param.{IntParam, Param, ParamValidators, Params}
+
+private[spark] trait RabitParams extends Params {
+
+  final val rabitTrackerTimeout = new IntParam(this, "rabitTrackerTimeout", "The number of " +
+    "seconds before timeout waiting for workers to connect. and for the tracker to shutdown.",
+    ParamValidators.gtEq(0))
+
+  final def getRabitTrackerTimeout: Int = $(rabitTrackerTimeout)
+
+  final val rabitTrackerHostIp = new Param[String](this, "rabitTrackerHostIp", "The Rabit " +
+    "Tracker host IP address. This is only needed if the host IP cannot be automatically " +
+    "guessed.")
+
+  final def getRabitTrackerHostIp: String = $(rabitTrackerHostIp)
+
+  final val rabitTrackerPort = new IntParam(this, "rabitTrackerPort", "The port number for the " +
+    "tracker to listen to. Use a system allocated one by default.",
+    ParamValidators.gtEq(0))
+
+  final def getRabitTrackerPort: Int = $(rabitTrackerPort)
+
+  setDefault(rabitTrackerTimeout -> 0, rabitTrackerHostIp -> "", rabitTrackerPort -> 0)
+}
+
+