[pyspark] Handle the device parameter in pyspark. (dmlc#9390)

- Handle the new `device` parameter in PySpark.
- Deprecate the old `use_gpu` parameter.

Author: trivialfis

doc/tutorials/spark_estimator.rst

@@ -35,13 +35,13 @@ We can create a ``SparkXGBRegressor`` estimator like:
   )
 
 
-The above snippet creates a spark estimator which can fit on a spark dataset,
-and return a spark model that can transform a spark dataset and generate dataset
-with prediction column. We can set almost all of xgboost sklearn estimator parameters
-as ``SparkXGBRegressor`` parameters, but some parameter such as ``nthread`` is forbidden
-in spark estimator, and some parameters are replaced with pyspark specific parameters
-such as ``weight_col``, ``validation_indicator_col``, ``use_gpu``, for details please see
-``SparkXGBRegressor`` doc.
+The above snippet creates a spark estimator which can fit on a spark dataset, and return a
+spark model that can transform a spark dataset and generate dataset with prediction
+column. We can set almost all of xgboost sklearn estimator parameters as
+``SparkXGBRegressor`` parameters, but some parameter such as ``nthread`` is forbidden in
+spark estimator, and some parameters are replaced with pyspark specific parameters such as
+``weight_col``, ``validation_indicator_col``, for details please see ``SparkXGBRegressor``
+doc.
 
 The following code snippet shows how to train a spark xgboost regressor model,
 first we need to prepare a training dataset as a spark dataframe contains
@@ -88,7 +88,7 @@ XGBoost PySpark fully supports GPU acceleration. Users are not only able to enab
 efficient training but also utilize their GPUs for the whole PySpark pipeline including
 ETL and inference. In below sections, we will walk through an example of training on a
 PySpark standalone GPU cluster. To get started, first we need to install some additional
-packages, then we can set the ``use_gpu`` parameter to ``True``.
+packages, then we can set the ``device`` parameter to ``cuda`` or ``gpu``.
 
 Prepare the necessary packages
 ==============================
@@ -128,7 +128,7 @@ Write your PySpark application
 ==============================
 
 Below snippet is a small example for training xgboost model with PySpark. Notice that we are
-using a list of feature names and the additional parameter ``use_gpu``:
+using a list of feature names and the additional parameter ``device``:
 
 .. code-block:: python
 
@@ -148,12 +148,12 @@ using a list of feature names and the additional parameter ``use_gpu``:
   # get a list with feature column names
   feature_names = [x.name for x in train_df.schema if x.name != label_name]
 
-  # create a xgboost pyspark regressor estimator and set use_gpu=True
+  # create a xgboost pyspark regressor estimator and set device="cuda"
   regressor = SparkXGBRegressor(
     features_col=feature_names,
     label_col=label_name,
     num_workers=2,
-    use_gpu=True,
+    device="cuda",
   )
 
   # train and return the model
@@ -163,6 +163,7 @@ using a list of feature names and the additional parameter ``use_gpu``:
   predict_df = model.transform(test_df)
   predict_df.show()
 
+Like other distributed interfaces, the ```device`` parameter doesn't support specifying ordinal as GPUs are managed by Spark instead of XGBoost (good: ``device=cuda``, bad: ``device=cuda:0``).
 
 Submit the PySpark application
 ==============================

python-package/xgboost/core.py

@@ -276,6 +276,27 @@ def _check_call(ret: int) -> None:
         raise XGBoostError(py_str(_LIB.XGBGetLastError()))
 
 
+def _check_distributed_params(kwargs: Dict[str, Any]) -> None:
+    """Validate parameters in distributed environments."""
+    device = kwargs.get("device", None)
+    if device and not isinstance(device, str):
+        msg = "Invalid type for the `device` parameter"
+        msg += _expect((str,), type(device))
+        raise TypeError(msg)
+
+    if device and device.find(":") != -1:
+        raise ValueError(
+            "Distributed training doesn't support selecting device ordinal as GPUs are"
+            " managed by the distributed framework. use `device=cuda` or `device=gpu`"
+            " instead."
+        )
+
+    if kwargs.get("booster", None) == "gblinear":
+        raise NotImplementedError(
+            f"booster `{kwargs['booster']}` is not supported for distributed training."
+        )
+
+
 def build_info() -> dict:
     """Build information of XGBoost.  The returned value format is not stable. Also,
     please note that build time dependency is not the same as runtime dependency. For

python-package/xgboost/dask.py

@@ -70,6 +70,7 @@
     Metric,
     Objective,
     QuantileDMatrix,
+    _check_distributed_params,
     _deprecate_positional_args,
     _expect,
 )
@@ -924,17 +925,7 @@ async def _train_async(
 ) -> Optional[TrainReturnT]:
     workers = _get_workers_from_data(dtrain, evals)
     _rabit_args = await _get_rabit_args(len(workers), dconfig, client)
-
-    if params.get("booster", None) == "gblinear":
-        raise NotImplementedError(
-            f"booster `{params['booster']}` is not yet supported for dask."
-        )
-    device = params.get("device", None)
-    if device and device.find(":") != -1:
-        raise ValueError(
-            "The dask interface for XGBoost doesn't support selecting specific device"
-            " ordinal. Use `device=cpu` or `device=cuda` instead."
-        )
+    _check_distributed_params(params)
 
     def dispatched_train(
         parameters: Dict,

python-package/xgboost/sklearn.py

@@ -1004,13 +1004,17 @@ def fit(
             Validation metrics will help us track the performance of the model.
 
         eval_metric : str, list of str, or callable, optional
+
             .. deprecated:: 1.6.0
-                Use `eval_metric` in :py:meth:`__init__` or :py:meth:`set_params` instead.
+
+            Use `eval_metric` in :py:meth:`__init__` or :py:meth:`set_params` instead.
 
         early_stopping_rounds : int
+
             .. deprecated:: 1.6.0
-                Use `early_stopping_rounds` in :py:meth:`__init__` or
-                :py:meth:`set_params` instead.
+
+            Use `early_stopping_rounds` in :py:meth:`__init__` or :py:meth:`set_params`
+            instead.
         verbose :
             If `verbose` is True and an evaluation set is used, the evaluation metric
             measured on the validation set is printed to stdout at each boosting stage.

python-package/xgboost/spark/core.py

@@ -60,7 +60,7 @@
 import xgboost
 from xgboost import XGBClassifier
 from xgboost.compat import is_cudf_available
-from xgboost.core import Booster
+from xgboost.core import Booster, _check_distributed_params
 from xgboost.sklearn import DEFAULT_N_ESTIMATORS, XGBModel, _can_use_qdm
 from xgboost.training import train as worker_train
 
@@ -92,6 +92,7 @@
     get_class_name,
     get_logger,
     serialize_booster,
+    use_cuda,
 )
 
 # Put pyspark specific params here, they won't be passed to XGBoost.
@@ -108,7 +109,6 @@
     "arbitrary_params_dict",
     "force_repartition",
     "num_workers",
-    "use_gpu",
     "feature_names",
     "features_cols",
     "enable_sparse_data_optim",
@@ -132,8 +132,7 @@
 _inverse_pyspark_param_alias_map = {v: k for k, v in _pyspark_param_alias_map.items()}
 
 _unsupported_xgb_params = [
-    "gpu_id",  # we have "use_gpu" pyspark param instead.
-    "device",  # we have "use_gpu" pyspark param instead.
+    "gpu_id",  # we have "device" pyspark param instead.
     "enable_categorical",  # Use feature_types param to specify categorical feature instead
     "use_label_encoder",
     "n_jobs",  # Do not allow user to set it, will use `spark.task.cpus` value instead.
@@ -198,11 +197,24 @@ class _SparkXGBParams(
         "The number of XGBoost workers. Each XGBoost worker corresponds to one spark task.",
         TypeConverters.toInt,
     )
+    device = Param(
+        Params._dummy(),
+        "device",
+        (
+            "The device type for XGBoost executors. Available options are `cpu`,`cuda`"
+            " and `gpu`. Set `device` to `cuda` or `gpu` if the executors are running "
+            "on GPU instances. Currently, only one GPU per task is supported."
+        ),
+        TypeConverters.toString,
+    )
     use_gpu = Param(
         Params._dummy(),
         "use_gpu",
-        "A boolean variable. Set use_gpu=true if the executors "
-        + "are running on GPU instances. Currently, only one GPU per task is supported.",
+        (
+            "Deprecated, use `device` instead. A boolean variable. Set use_gpu=true "
+            "if the executors are running on GPU instances. Currently, only one GPU per"
+            " task is supported."
+        ),
         TypeConverters.toBoolean,
     )
     force_repartition = Param(
@@ -336,10 +348,20 @@ def _validate_params(self) -> None:
                 f"It cannot be less than 1 [Default is 1]"
             )
 
+        tree_method = self.getOrDefault(self.getParam("tree_method"))
+        if (
+            self.getOrDefault(self.use_gpu) or use_cuda(self.getOrDefault(self.device))
+        ) and not _can_use_qdm(tree_method):
+            raise ValueError(
+                f"The `{tree_method}` tree method is not supported on GPU."
+            )
+
         if self.getOrDefault(self.features_cols):
-            if not self.getOrDefault(self.use_gpu):
+            if not use_cuda(self.getOrDefault(self.device)) and not self.getOrDefault(
+                self.use_gpu
+            ):
                 raise ValueError(
-                    "features_col param with list value requires enabling use_gpu."
+                    "features_col param with list value requires `device=cuda`."
                 )
 
         if self.getOrDefault("objective") is not None:
@@ -392,17 +414,7 @@ def _validate_params(self) -> None:
                     "`pyspark.ml.linalg.Vector` type."
                 )
 
-        if self.getOrDefault(self.use_gpu):
-            tree_method = self.getParam("tree_method")
-            if (
-                self.getOrDefault(tree_method) is not None
-                and self.getOrDefault(tree_method) != "gpu_hist"
-            ):
-                raise ValueError(
-                    f"tree_method should be 'gpu_hist' or None when use_gpu is True,"
-                    f"found {self.getOrDefault(tree_method)}."
-                )
-
+        if use_cuda(self.getOrDefault(self.device)) or self.getOrDefault(self.use_gpu):
             gpu_per_task = (
                 _get_spark_session()
                 .sparkContext.getConf()
@@ -424,8 +436,8 @@ def _validate_params(self) -> None:
                 # so it's okay for printing the below warning instead of checking the real
                 # gpu numbers and raising the exception.
                 get_logger(self.__class__.__name__).warning(
-                    "You enabled use_gpu in spark local mode. Please make sure your local node "
-                    "has at least %d GPUs",
+                    "You enabled GPU in spark local mode. Please make sure your local "
+                    "node has at least %d GPUs",
                     self.getOrDefault(self.num_workers),
                 )
             else:
@@ -558,6 +570,7 @@ def __init__(self) -> None:
         #  they are added in `setParams`.
         self._setDefault(
             num_workers=1,
+            device="cpu",
             use_gpu=False,
             force_repartition=False,
             repartition_random_shuffle=False,
@@ -566,9 +579,7 @@ def __init__(self) -> None:
             arbitrary_params_dict={},
         )
 
-    def setParams(
-        self, **kwargs: Dict[str, Any]
-    ) -> None:  # pylint: disable=invalid-name
+    def setParams(self, **kwargs: Any) -> None:  # pylint: disable=invalid-name
         """
         Set params for the estimator.
         """
@@ -613,6 +624,8 @@ def setParams(
                     )
                     raise ValueError(err_msg)
                 _extra_params[k] = v
+
+        _check_distributed_params(kwargs)
         _existing_extra_params = self.getOrDefault(self.arbitrary_params_dict)
         self._set(arbitrary_params_dict={**_existing_extra_params, **_extra_params})
 
@@ -709,9 +722,6 @@ def _get_distributed_train_params(self, dataset: DataFrame) -> Dict[str, Any]:
         # TODO: support "num_parallel_tree" for random forest
         params["num_boost_round"] = self.getOrDefault("n_estimators")
 
-        if self.getOrDefault(self.use_gpu):
-            params["tree_method"] = "gpu_hist"
-
         return params
 
     @classmethod
@@ -883,8 +893,9 @@ def _fit(self, dataset: DataFrame) -> "_SparkXGBModel":
             dmatrix_kwargs,
         ) = self._get_xgb_parameters(dataset)
 
-        use_gpu = self.getOrDefault(self.use_gpu)
-
+        run_on_gpu = use_cuda(self.getOrDefault(self.device)) or self.getOrDefault(
+            self.use_gpu
+        )
         is_local = _is_local(_get_spark_session().sparkContext)
 
         num_workers = self.getOrDefault(self.num_workers)
@@ -903,7 +914,7 @@ def _train_booster(
             dev_ordinal = None
             use_qdm = _can_use_qdm(booster_params.get("tree_method", None))
 
-            if use_gpu:
+            if run_on_gpu:
                 dev_ordinal = (
                     context.partitionId() if is_local else _get_gpu_id(context)
                 )