[jvm-packages] Add the new device parameter. (#9385)

jvm-packages/xgboost4j-spark-gpu/src/main/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuPreXGBoost.scala

@@ -137,8 +137,12 @@ object GpuPreXGBoost extends PreXGBoostProvider {
     val (Seq(labelName, weightName, marginName), feturesCols, groupName, evalSets) =
       estimator match {
         case est: XGBoostEstimatorCommon =>
-          require(est.isDefined(est.treeMethod) && est.getTreeMethod.equals("gpu_hist"),
-            s"GPU train requires tree_method set to gpu_hist")
+          require(
+            est.isDefined(est.device) &&
+              (est.getDevice.equals("cuda") || est.getDevice.equals("gpu")) ||
+              est.isDefined(est.treeMethod) && est.getTreeMethod.equals("gpu_hist"),
+            s"GPU train requires `device` set to `cuda` or `gpu`."
+          )
           val groupName = estimator match {
             case regressor: XGBoostRegressor => if (regressor.isDefined(regressor.groupCol)) {
               regressor.getGroupCol } else ""

jvm-packages/xgboost4j-spark-gpu/src/test/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuXGBoostGeneralSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2021-2022 by Contributors
+ Copyright (c) 2021-2023 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -50,9 +50,12 @@ class GpuXGBoostGeneralSuite extends GpuTestSuite {
     withGpuSparkSession() { spark =>
       import spark.implicits._
       val trainingDf = trainingData.toDF(allColumnNames: _*)
-      val xgbParam = Map("eta" -> 0.1f, "max_depth" -> 2, "objective" -> "multi:softprob",
-        "num_class" -> 3, "num_round" -> 5, "num_workers" -> 1, "tree_method" -> "gpu_hist",
-        "features_cols" -> featureNames, "label_col" -> labelName)
+      val xgbParam = Map(
+        "eta" -> 0.1f, "max_depth" -> 2, "objective" -> "multi:softprob",
+        "num_class" -> 3, "num_round" -> 5, "num_workers" -> 1,
+        "tree_method" -> "hist", "device" -> "cuda",
+        "features_cols" -> featureNames, "label_col" -> labelName
+      )
       new XGBoostClassifier(xgbParam)
         .fit(trainingDf)
     }
@@ -65,8 +68,11 @@ class GpuXGBoostGeneralSuite extends GpuTestSuite {
 
       trainingDf = trainingDf.select(labelName, "f2", weightName, "f3", baseMarginName, "f1")
 
-      val xgbParam = Map("eta" -> 0.1f, "max_depth" -> 2, "objective" -> "multi:softprob",
-        "num_class" -> 3, "num_round" -> 5, "num_workers" -> 1, "tree_method" -> "gpu_hist")
+      val xgbParam = Map(
+        "eta" -> 0.1f, "max_depth" -> 2, "objective" -> "multi:softprob",
+        "num_class" -> 3, "num_round" -> 5, "num_workers" -> 1,
+        "tree_method" -> "hist", "device" -> "cuda"
+      )
       new XGBoostClassifier(xgbParam)
         .setFeaturesCol(featureNames)
         .setLabelCol(labelName)
@@ -127,7 +133,7 @@ class GpuXGBoostGeneralSuite extends GpuTestSuite {
     }
   }
 
-  test("Throw exception when tree method is not set to gpu_hist") {
+  test("Throw exception when device is not set to cuda") {
     withGpuSparkSession() { spark =>
       import spark.implicits._
       val trainingDf = trainingData.toDF(allColumnNames: _*)
@@ -139,12 +145,11 @@ class GpuXGBoostGeneralSuite extends GpuTestSuite {
           .setLabelCol(labelName)
           .fit(trainingDf)
       }
-      assert(thrown.getMessage.contains("GPU train requires tree_method set to gpu_hist"))
+      assert(thrown.getMessage.contains("GPU train requires `device` set to `cuda`"))
     }
   }
 
   test("Train with eval") {
-
     withGpuSparkSession() { spark =>
       import spark.implicits._
       val Array(trainingDf, eval1, eval2) = trainingData.toDF(allColumnNames: _*)
@@ -184,4 +189,24 @@ class GpuXGBoostGeneralSuite extends GpuTestSuite {
     }
   }
 
+  test("device ordinal should not be specified") {
+    withGpuSparkSession() { spark =>
+      import spark.implicits._
+      val trainingDf = trainingData.toDF(allColumnNames: _*)
+      val params = Map(
+        "objective" -> "multi:softprob",
+        "num_class" -> 3,
+        "num_round" -> 5,
+        "num_workers" -> 1
+      )
+      val thrown = intercept[IllegalArgumentException] {
+        new XGBoostClassifier(params)
+        .setFeaturesCol(featureNames)
+        .setLabelCol(labelName)
+        .setDevice("cuda:1")
+        .fit(trainingDf)
+      }
+      assert(thrown.getMessage.contains("`cuda` or `gpu`"))
+    }
+  }
 }

jvm-packages/xgboost4j-spark-gpu/src/test/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuXGBoostRegressorSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2021-2022 by Contributors
+ Copyright (c) 2021-2023 by Contributors
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -40,7 +40,7 @@ class GpuXGBoostRegressorSuite extends GpuTestSuite {
   test("The transform result should be same for several runs on same model") {
     withGpuSparkSession(enableCsvConf()) { spark =>
       val xgbParam = Map("eta" -> 0.1f, "max_depth" -> 2, "objective" -> "reg:squarederror",
-        "num_round" -> 10, "num_workers" -> 1, "tree_method" -> "gpu_hist",
+        "num_round" -> 10, "num_workers" -> 1, "tree_method" -> "hist", "device" -> "cuda",
         "features_cols" -> featureNames, "label_col" -> labelName)
       val Array(originalDf, testDf) = spark.read.option("header", "true").schema(schema)
         .csv(getResourcePath("/rank.train.csv")).randomSplit(Array(0.7, 0.3), seed = 1)
@@ -54,10 +54,30 @@ class GpuXGBoostRegressorSuite extends GpuTestSuite {
     }
   }
 
+  test("Tree method gpu_hist still works") {
+    withGpuSparkSession(enableCsvConf()) { spark =>
+      val params = Map(
+        "tree_method" -> "gpu_hist",
+        "features_cols" -> featureNames,
+        "label_col" -> labelName,
+        "num_round" -> 10,
+        "num_workers" -> 1
+      )
+      val Array(originalDf, testDf) = spark.read.option("header", "true").schema(schema)
+        .csv(getResourcePath("/rank.train.csv")).randomSplit(Array(0.7, 0.3), seed = 1)
+      // Get a model
+      val model = new XGBoostRegressor(params).fit(originalDf)
+      val left = model.transform(testDf).collect()
+      val right = model.transform(testDf).collect()
+      // The left should be same with right
+      assert(compareResults(true, 0.000001, left, right))
+    }
+  }
+
   test("use weight") {
     withGpuSparkSession(enableCsvConf()) { spark =>
       val xgbParam = Map("eta" -> 0.1f, "max_depth" -> 2, "objective" -> "reg:squarederror",
-        "num_round" -> 10, "num_workers" -> 1, "tree_method" -> "gpu_hist",
+        "num_round" -> 10, "num_workers" -> 1, "tree_method" -> "hist", "device" -> "cuda",
         "features_cols" -> featureNames, "label_col" -> labelName)
       val Array(originalDf, testDf) = spark.read.option("header", "true").schema(schema)
         .csv(getResourcePath("/rank.train.csv")).randomSplit(Array(0.7, 0.3), seed = 1)
@@ -88,7 +108,8 @@ class GpuXGBoostRegressorSuite extends GpuTestSuite {
       val classifier = new XGBoostRegressor(xgbParam)
         .setFeaturesCol(featureNames)
         .setLabelCol(labelName)
-        .setTreeMethod("gpu_hist")
+        .setTreeMethod("hist")
+        .setDevice("cuda")
       (classifier.fit(rawInput), testDf)
     }
 
@@ -175,7 +196,7 @@ class GpuXGBoostRegressorSuite extends GpuTestSuite {
       val classifier = new XGBoostRegressor(xgbParam)
         .setFeaturesCol(featureNames)
         .setLabelCol(labelName)
-        .setTreeMethod("gpu_hist")
+        .setDevice("cuda")
       classifier.fit(rawInput)
     }
 
@@ -234,5 +255,4 @@ class GpuXGBoostRegressorSuite extends GpuTestSuite {
       assert(testDf.count() === ret.length)
     }
   }
-
 }