Fix typo. (#8192)

* Fix loading DMatrix binary in distributed env. (#8149)

- Try to load DMatrix binary before trying to parse text input.
- Remove some unmaintained code.

* Fix.

CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.14 FATAL_ERROR)
-project(xgboost LANGUAGES CXX C VERSION 1.6.0)
+project(xgboost LANGUAGES CXX C VERSION 1.6.2)
 include(cmake/Utils.cmake)
 list(APPEND CMAKE_MODULE_PATH "${xgboost_SOURCE_DIR}/cmake/modules")
 cmake_policy(SET CMP0022 NEW)
@@ -200,6 +200,10 @@ endif (JVM_BINDINGS)
 # Plugin
 add_subdirectory(${xgboost_SOURCE_DIR}/plugin)
 
+if (PLUGIN_RMM)
+  find_package(rmm REQUIRED)
+endif (PLUGIN_RMM)
+
 #-- library
 if (BUILD_STATIC_LIB)
   add_library(xgboost STATIC)

Jenkinsfile

@@ -7,7 +7,7 @@
 dockerRun = 'tests/ci_build/ci_build.sh'
 
 // Which CUDA version to use when building reference distribution wheel
-ref_cuda_ver = '11.0'
+ref_cuda_ver = '11.0.3'
 
 import groovy.transform.Field
 
@@ -60,9 +60,9 @@ pipeline {
             'build-cpu-rabit-mock': { BuildCPUMock() },
             // Build reference, distribution-ready Python wheel with CUDA 11.0
             // using CentOS 7 image
-            'build-gpu-cuda11.0': { BuildCUDA(cuda_version: '11.0', build_rmm: true) },
+            'build-gpu-cuda11.0': { BuildCUDA(cuda_version: '11.0.3', build_rmm: true) },
-            'build-gpu-rpkg': { BuildRPackageWithCUDA(cuda_version: '11.0') },
+            'build-gpu-rpkg': { BuildRPackageWithCUDA(cuda_version: '11.0.3') },
-            'build-jvm-packages-gpu-cuda11.0': { BuildJVMPackagesWithCUDA(spark_version: '3.0.1', cuda_version: '11.0') },
+            'build-jvm-packages-gpu-cuda11.0': { BuildJVMPackagesWithCUDA(spark_version: '3.0.1', cuda_version: '11.0.3') },
             'build-jvm-packages': { BuildJVMPackages(spark_version: '3.0.1') },
             'build-jvm-doc': { BuildJVMDoc() }
           ])
@@ -77,9 +77,9 @@ pipeline {
             'test-python-cpu': { TestPythonCPU() },
             'test-python-cpu-arm64': { TestPythonCPUARM64() },
             // artifact_cuda_version doesn't apply to RMM tests; RMM tests will always match CUDA version between artifact and host env
-            'test-python-gpu-cuda11.0': { TestPythonGPU(artifact_cuda_version: '11.0', host_cuda_version: '11.0', test_rmm: true) },
+            'test-python-gpu-cuda11.0': { TestPythonGPU(artifact_cuda_version: '11.0.3', host_cuda_version: '11.0.3', test_rmm: true) },
-            'test-python-mgpu-cuda11.0': { TestPythonGPU(artifact_cuda_version: '11.0', host_cuda_version: '11.0', multi_gpu: true, test_rmm: true) },
+            'test-python-mgpu-cuda11.0': { TestPythonGPU(artifact_cuda_version: '11.0.3', host_cuda_version: '11.0.3', multi_gpu: true, test_rmm: true) },
-            'test-cpp-gpu-cuda11.0': { TestCppGPU(artifact_cuda_version: '11.0', host_cuda_version: '11.0', test_rmm: true) },
+            'test-cpp-gpu-cuda11.0': { TestCppGPU(artifact_cuda_version: '11.0.3', host_cuda_version: '11.0.3', test_rmm: true) },
             'test-jvm-jdk8': { CrossTestJVMwithJDK(jdk_version: '8', spark_version: '3.0.0') }
           ])
         }
@@ -123,7 +123,7 @@ def ClangTidy() {
     echo "Running clang-tidy job..."
     def container_type = "clang_tidy"
     def docker_binary = "docker"
-    def dockerArgs = "--build-arg CUDA_VERSION_ARG=11.0"
+    def dockerArgs = "--build-arg CUDA_VERSION_ARG=11.0.3"
     sh """
     ${dockerRun} ${container_type} ${docker_binary} ${dockerArgs} python3 tests/ci_build/tidy.py --cuda-archs 75
     """
@@ -397,7 +397,7 @@ def TestCppGPU(args) {
   node(nodeReq) {
     unstash name: "xgboost_cpp_tests_cuda${artifact_cuda_version}"
     unstash name: 'srcs'
-    echo "Test C++, CUDA ${args.host_cuda_version}"
+    echo "Test C++, CUDA ${args.host_cuda_version}, rmm: ${args.test_rmm}"
     def container_type = "gpu"
     def docker_binary = "nvidia-docker"
     def docker_args = "--build-arg CUDA_VERSION_ARG=${args.host_cuda_version}"
@@ -410,7 +410,7 @@ def TestCppGPU(args) {
       docker_binary = "nvidia-docker"
       docker_args = "--build-arg CUDA_VERSION_ARG=${args.host_cuda_version}"
       sh """
-      ${dockerRun} ${container_type} ${docker_binary} ${docker_args} bash -c "source activate gpu_test && build/testxgboost --use-rmm-pool --gtest_filter=-*DeathTest.*"
+      ${dockerRun} ${container_type} ${docker_binary} ${docker_args} bash -c "source activate gpu_test && build/testxgboost --use-rmm-pool"
       """
     }
     deleteDir()
@@ -445,7 +445,7 @@ def DeployJVMPackages(args) {
     if (env.BRANCH_NAME == 'master' || env.BRANCH_NAME.startsWith('release')) {
       echo 'Deploying to xgboost-maven-repo S3 repo...'
       sh """
-      ${dockerRun} jvm_gpu_build docker --build-arg CUDA_VERSION_ARG=11.0 tests/ci_build/deploy_jvm_packages.sh ${args.spark_version}
+      ${dockerRun} jvm_gpu_build docker --build-arg CUDA_VERSION_ARG=11.0.3 tests/ci_build/deploy_jvm_packages.sh ${args.spark_version}
       """
     }
     deleteDir()

R-package/DESCRIPTION

@@ -1,7 +1,7 @@
 Package: xgboost
 Type: Package
 Title: Extreme Gradient Boosting
-Version: 1.6.0.1
+Version: 1.6.2.1
 Date: 2022-03-29
 Authors@R: c(
   person("Tianqi", "Chen", role = c("aut"),

cmake/Utils.cmake

@@ -169,10 +169,17 @@ function(xgboost_set_cuda_flags target)
       $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/utf-8>)
   endif (MSVC)
 
-  set_target_properties(${target} PROPERTIES
+  if (PLUGIN_RMM)
-    CUDA_STANDARD 14
+    set_target_properties(${target} PROPERTIES
-    CUDA_STANDARD_REQUIRED ON
+      CUDA_STANDARD 17
-    CUDA_SEPARABLE_COMPILATION OFF)
+      CUDA_STANDARD_REQUIRED ON
+      CUDA_SEPARABLE_COMPILATION OFF)
+  else ()
+    set_target_properties(${target} PROPERTIES
+      CUDA_STANDARD 14
+      CUDA_STANDARD_REQUIRED ON
+      CUDA_SEPARABLE_COMPILATION OFF)
+  endif (PLUGIN_RMM)
 endfunction(xgboost_set_cuda_flags)
 
 macro(xgboost_link_nccl target)
@@ -189,10 +196,18 @@ endmacro(xgboost_link_nccl)
 
 # compile options
 macro(xgboost_target_properties target)
-  set_target_properties(${target} PROPERTIES
+  if (PLUGIN_RMM)
-    CXX_STANDARD 14
+    set_target_properties(${target} PROPERTIES
-    CXX_STANDARD_REQUIRED ON
+      CXX_STANDARD 17
-    POSITION_INDEPENDENT_CODE ON)
+      CXX_STANDARD_REQUIRED ON
+      POSITION_INDEPENDENT_CODE ON)
+  else ()
+    set_target_properties(${target} PROPERTIES
+      CXX_STANDARD 14
+      CXX_STANDARD_REQUIRED ON
+      POSITION_INDEPENDENT_CODE ON)
+  endif (PLUGIN_RMM)
+
   if (HIDE_CXX_SYMBOLS)
     #-- Hide all C++ symbols
     set_target_properties(${target} PROPERTIES
@@ -247,6 +262,10 @@ macro(xgboost_target_defs target)
       PRIVATE
       -DXGBOOST_BUILTIN_PREFETCH_PRESENT=1)
   endif (XGBOOST_BUILTIN_PREFETCH_PRESENT)
+
+  if (PLUGIN_RMM)
+    target_compile_definitions(objxgboost PUBLIC -DXGBOOST_USE_RMM=1)
+  endif (PLUGIN_RMM)
 endmacro(xgboost_target_defs)
 
 # handles dependencies
@@ -269,6 +288,10 @@ macro(xgboost_target_link_libraries target)
     xgboost_set_cuda_flags(${target})
   endif (USE_CUDA)
 
+  if (PLUGIN_RMM)
+    target_link_libraries(${target} PRIVATE rmm::rmm)
+  endif (PLUGIN_RMM)
+
   if (USE_NCCL)
     xgboost_link_nccl(${target})
   endif (USE_NCCL)

dev/release-py-r.py

@@ -75,7 +75,7 @@ def download_wheels(
     return filenames
 
 
-def download_py_packages(major: int, minor: int, commit_hash: str):
+def download_py_packages(branch: str, major: int, minor: int, commit_hash: str) -> None:
     platforms = [
         "win_amd64",
         "manylinux2014_x86_64",
@@ -84,7 +84,8 @@ def download_py_packages(major: int, minor: int, commit_hash: str):
         "macosx_12_0_arm64"
     ]
 
-    dir_URL = PREFIX + str(major) + "." + str(minor) + ".0" + "/"
+    branch = branch.split("_")[1]  # release_x.y.z
+    dir_URL = PREFIX + branch + "/"
     src_filename_prefix = "xgboost-" + args.release + "%2B" + commit_hash + "-py3-none-"
     target_filename_prefix = "xgboost-" + args.release + "-py3-none-"
 
@@ -105,16 +106,17 @@ Following steps should be done manually:
     )
 
 
-def download_r_packages(release: str, rc: str, commit: str) -> None:
+def download_r_packages(release: str, branch: str, rc: str, commit: str) -> None:
     platforms = ["win64", "linux"]
     dirname = "./r-packages"
     if not os.path.exists(dirname):
         os.mkdir(dirname)
 
     filenames = []
+    branch = branch.split("_")[1]  # release_x.y.z
 
     for plat in platforms:
-        url = f"{PREFIX}{release}/xgboost_r_gpu_{plat}_{commit}.tar.gz"
+        url = f"{PREFIX}{branch}/xgboost_r_gpu_{plat}_{commit}.tar.gz"
 
         if not rc:
             filename = f"xgboost_r_gpu_{plat}_{release}.tar.gz"
@@ -152,7 +154,11 @@ def main(args: argparse.Namespace) -> None:
         assert rc == "rc"
 
     release = str(major) + "." + str(minor) + "." + str(patch)
-    branch = "release_" + release
+    if args.branch is not None:
+        branch = args.branch
+    else:
+        branch = "release_" + str(major) + "." + str(minor) + ".0"
+
     git.clean("-xdf")
     git.checkout(branch)
     git.pull("origin", branch)
@@ -160,10 +166,10 @@ def main(args: argparse.Namespace) -> None:
     commit_hash = latest_hash()
 
     download_r_packages(
-        release, "" if rc is None else rc + str(rc_ver), commit_hash
+        release, branch, "" if rc is None else rc + str(rc_ver), commit_hash
     )
 
-    download_py_packages(major, minor, commit_hash)
+    download_py_packages(branch, major, minor, commit_hash)
 
 
 if __name__ == "__main__":
@@ -174,5 +180,14 @@ if __name__ == "__main__":
         required=True,
         help="Version tag, e.g. '1.3.2', or '1.5.0rc1'"
     )
+    parser.add_argument(
+        "--branch",
+        type=str,
+        default=None,
+        help=(
+            "Optional branch. Usually patch releases reuse the same branch of the"
+            " major release, but there can be exception."
+        )
+    )
     args = parser.parse_args()
     main(args)

doc/jvm/xgboost4j_spark_gpu_tutorial.rst

@@ -1,5 +1,5 @@
 #############################################
-XGBoost4J-Spark-GPU Tutorial (version 1.6.0+)
+XGBoost4J-Spark-GPU Tutorial (version 1.6.1+)
 #############################################
 
 **XGBoost4J-Spark-GPU** is an open source library aiming to accelerate distributed XGBoost training on Apache Spark cluster from
@@ -220,7 +220,7 @@ application jar is iris-1.0.0.jar
 
   cudf_version=22.02.0
   rapids_version=22.02.0
-  xgboost_version=1.6.0
+  xgboost_version=1.6.1
   main_class=Iris
   app_jar=iris-1.0.0.jar
 

doc/jvm/xgboost4j_spark_tutorial.rst

@@ -16,12 +16,6 @@ This tutorial is to cover the end-to-end process to build a machine learning pip
 * Building a Machine Learning Pipeline with XGBoost4J-Spark
 * Running XGBoost4J-Spark in Production
 
-.. note::
-
-  **SparkContext will be stopped by default when XGBoost training task fails**.
-
-  XGBoost4J-Spark 1.2.0+ exposes a parameter **kill_spark_context_on_worker_failure**. Set **kill_spark_context_on_worker_failure** to **false** so that the SparkContext will not be stopping on training failure. Instead of stopping the SparkContext, XGBoost4J-Spark will throw an exception instead. Users who want to re-use the SparkContext should wrap the training code in a try-catch block.
-
 .. contents::
   :backlinks: none
   :local:
@@ -129,7 +123,7 @@ labels. A DataFrame like this (containing vector-represented features and numeri
 
 .. note::
 
-  There is no need to assemble feature columns from version 1.6.0+. Instead, users can specify an array of
+  There is no need to assemble feature columns from version 1.6.1+. Instead, users can specify an array of
   feture column names by ``setFeaturesCol(value: Array[String])`` and XGBoost4j-Spark will do it.
 
 Dealing with missing values

doc/tutorials/categorical.rst

@@ -74,23 +74,20 @@ Optimal Partitioning
 .. versionadded:: 1.6
 
 Optimal partitioning is a technique for partitioning the categorical predictors for each
-node split, the proof of optimality for numerical objectives like ``RMSE`` was first
+node split, the proof of optimality for numerical output was first introduced by `[1]
-introduced by `[1] <#references>`__. The algorithm is used in decision trees for handling
+<#references>`__. The algorithm is used in decision trees `[2] <#references>`__, later
-regression and binary classification tasks `[2] <#references>`__, later LightGBM `[3]
+LightGBM `[3] <#references>`__ brought it to the context of gradient boosting trees and
-<#references>`__ brought it to the context of gradient boosting trees and now is also
+now is also adopted in XGBoost as an optional feature for handling categorical
-adopted in XGBoost as an optional feature for handling categorical splits. More
+splits. More specifically, the proof by Fisher `[1] <#references>`__ states that, when
-specifically, the proof by Fisher `[1] <#references>`__ states that, when trying to
+trying to partition a set of discrete values into groups based on the distances between a
-partition a set of discrete values into groups based on the distances between a measure of
+measure of these values, one only needs to look at sorted partitions instead of
-these values, one only needs to look at sorted partitions instead of enumerating all
+enumerating all possible permutations. In the context of decision trees, the discrete
-possible permutations. In the context of decision trees, the discrete values are
+values are categories, and the measure is the output leaf value.  Intuitively, we want to
-categories, and the measure is the output leaf value.  Intuitively, we want to group the
+group the categories that output similar leaf values. During split finding, we first sort
-categories that output similar leaf values. During split finding, we first sort the
+the gradient histogram to prepare the contiguous partitions then enumerate the splits
-gradient histogram to prepare the contiguous partitions then enumerate the splits
 according to these sorted values. One of the related parameters for XGBoost is
 ``max_cat_to_one_hot``, which controls whether one-hot encoding or partitioning should be
-used for each feature, see :doc:`/parameter` for details.  When objective is not
+used for each feature, see :doc:`/parameter` for details.
-regression or binary classification, XGBoost will fallback to using onehot encoding
-instead.
 
 
 **********************

include/xgboost/task.h

@@ -36,10 +36,6 @@ struct ObjInfo {
 
   explicit ObjInfo(Task t) : task{t} {}
   ObjInfo(Task t, bool khess) : task{t}, const_hess{khess} {}
-
-  XGBOOST_DEVICE bool UseOneHot() const {
-    return (task != ObjInfo::kRegression && task != ObjInfo::kBinary);
-  }
 };
 }  // namespace xgboost
 #endif  // XGBOOST_TASK_H_

include/xgboost/version_config.h

@@ -6,6 +6,6 @@
 
 #define XGBOOST_VER_MAJOR 1
 #define XGBOOST_VER_MINOR 6
-#define XGBOOST_VER_PATCH 0
+#define XGBOOST_VER_PATCH 2
 
 #endif  // XGBOOST_VERSION_CONFIG_H_

jvm-packages/pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>ml.dmlc</groupId>
     <artifactId>xgboost-jvm_2.12</artifactId>
-    <version>1.6.0</version>
+    <version>1.6.2</version>
     <packaging>pom</packaging>
     <name>XGBoost JVM Package</name>
     <description>JVM Package for XGBoost</description>

jvm-packages/xgboost4j-example/pom.xml

@@ -6,10 +6,10 @@
     <parent>
         <groupId>ml.dmlc</groupId>
         <artifactId>xgboost-jvm_2.12</artifactId>
-        <version>1.6.0</version>
+        <version>1.6.2</version>
     </parent>
     <artifactId>xgboost4j-example_2.12</artifactId>
-    <version>1.6.0</version>
+    <version>1.6.2</version>
     <packaging>jar</packaging>
     <build>
         <plugins>
@@ -26,7 +26,7 @@
         <dependency>
             <groupId>ml.dmlc</groupId>
             <artifactId>xgboost4j-spark_${scala.binary.version}</artifactId>
-            <version>1.6.0</version>
+            <version>1.6.2</version>
         </dependency>
         <dependency>
             <groupId>org.apache.spark</groupId>
@@ -37,7 +37,7 @@
         <dependency>
             <groupId>ml.dmlc</groupId>
             <artifactId>xgboost4j-flink_${scala.binary.version}</artifactId>
-            <version>1.6.0</version>
+            <version>1.6.2</version>
         </dependency>
         <dependency>
             <groupId>org.apache.commons</groupId>

jvm-packages/xgboost4j-flink/pom.xml

@@ -6,10 +6,10 @@
     <parent>
         <groupId>ml.dmlc</groupId>
         <artifactId>xgboost-jvm_2.12</artifactId>
-        <version>1.6.0</version>
+        <version>1.6.2</version>
     </parent>
     <artifactId>xgboost4j-flink_2.12</artifactId>
-    <version>1.6.0</version>
+    <version>1.6.2</version>
     <build>
         <plugins>
             <plugin>
@@ -26,7 +26,7 @@
         <dependency>
             <groupId>ml.dmlc</groupId>
             <artifactId>xgboost4j_${scala.binary.version}</artifactId>
-            <version>1.6.0</version>
+            <version>1.6.2</version>
         </dependency>
         <dependency>
             <groupId>org.apache.commons</groupId>

jvm-packages/xgboost4j-gpu/pom.xml

@@ -6,10 +6,10 @@
     <parent>
         <groupId>ml.dmlc</groupId>
         <artifactId>xgboost-jvm_2.12</artifactId>
-        <version>1.6.0</version>
+        <version>1.6.2</version>
     </parent>
     <artifactId>xgboost4j-gpu_2.12</artifactId>
-    <version>1.6.0</version>
+    <version>1.6.2</version>
     <packaging>jar</packaging>
 
     <dependencies>

jvm-packages/xgboost4j-gpu/src/test/java/ml/dmlc/xgboost4j/gpu/java/BoosterTest.java

@@ -69,7 +69,7 @@ public class BoosterTest {
       .hasHeader().build();
 
     int maxBin = 16;
-    int round = 100;
+    int round = 10;
     //set params
     Map<String, Object> paramMap = new HashMap<String, Object>() {
       {

jvm-packages/xgboost4j-spark-gpu/pom.xml

@@ -6,7 +6,7 @@
     <parent>
         <groupId>ml.dmlc</groupId>
         <artifactId>xgboost-jvm_2.12</artifactId>
-        <version>1.6.0</version>
+        <version>1.6.2</version>
     </parent>
     <artifactId>xgboost4j-spark-gpu_2.12</artifactId>
     <build>
@@ -24,7 +24,7 @@
         <dependency>
             <groupId>ml.dmlc</groupId>
             <artifactId>xgboost4j-gpu_${scala.binary.version}</artifactId>
-            <version>1.6.0</version>
+            <version>1.6.2</version>
         </dependency>
         <dependency>
             <groupId>org.apache.spark</groupId>

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

@@ -56,18 +56,20 @@ class GpuPreXGBoost extends PreXGBoostProvider {
   }
 
   /**
-   * Convert the Dataset[_] to RDD[Watches] which will be fed to XGBoost
+   * Convert the Dataset[_] to RDD[() => Watches] which will be fed to XGBoost
    *
    * @param estimator [[XGBoostClassifier]] or [[XGBoostRegressor]]
    * @param dataset   the training data
    * @param params    all user defined and defaulted params
-   * @return [[XGBoostExecutionParams]] => (RDD[[Watches]], Option[ RDD[_] ])
+   * @return [[XGBoostExecutionParams]] => (Boolean, RDD[[() => Watches]], Option[ RDD[_] ])
-   *         RDD[Watches] will be used as the training input
+   *         Boolean if building DMatrix in rabit context
+   *         RDD[() => Watches] will be used as the training input
    *         Option[ RDD[_] ] is the optional cached RDD
    */
   override def buildDatasetToRDD(estimator: Estimator[_],
       dataset: Dataset[_],
-      params: Map[String, Any]): XGBoostExecutionParams => (RDD[Watches], Option[RDD[_]]) = {
+      params: Map[String, Any]):
+    XGBoostExecutionParams => (Boolean, RDD[() => Watches], Option[RDD[_]]) = {
     GpuPreXGBoost.buildDatasetToRDD(estimator, dataset, params)
   }
 
@@ -116,19 +118,21 @@ object GpuPreXGBoost extends PreXGBoostProvider {
   }
 
   /**
-   * Convert the Dataset[_] to RDD[Watches] which will be fed to XGBoost
+   * Convert the Dataset[_] to RDD[() => Watches] which will be fed to XGBoost
    *
    * @param estimator supports XGBoostClassifier and XGBoostRegressor
    * @param dataset   the training data
    * @param params    all user defined and defaulted params
-   * @return [[XGBoostExecutionParams]] => (RDD[[Watches]], Option[ RDD[_] ])
+   * @return [[XGBoostExecutionParams]] => (Boolean, RDD[[() => Watches]], Option[ RDD[_] ])
-   *         RDD[Watches] will be used as the training input
+   *         Boolean if building DMatrix in rabit context
+   *         RDD[() => Watches] will be used as the training input to build DMatrix
    *         Option[ RDD[_] ] is the optional cached RDD
    */
   override def buildDatasetToRDD(
       estimator: Estimator[_],
       dataset: Dataset[_],
-      params: Map[String, Any]): XGBoostExecutionParams => (RDD[Watches], Option[RDD[_]]) = {
+      params: Map[String, Any]):
+    XGBoostExecutionParams => (Boolean, RDD[() => Watches], Option[RDD[_]]) = {
 
     val (Seq(labelName, weightName, marginName), feturesCols, groupName, evalSets) =
       estimator match {
@@ -166,7 +170,7 @@ object GpuPreXGBoost extends PreXGBoostProvider {
     xgbExecParams: XGBoostExecutionParams =>
       val dataMap = prepareInputData(trainingData, evalDataMap, xgbExecParams.numWorkers,
         xgbExecParams.cacheTrainingSet)
-      (buildRDDWatches(dataMap, xgbExecParams, evalDataMap.isEmpty), None)
+      (true, buildRDDWatches(dataMap, xgbExecParams, evalDataMap.isEmpty), None)
   }
 
   /**
@@ -403,14 +407,9 @@ object GpuPreXGBoost extends PreXGBoostProvider {
   }
 
   private def repartitionInputData(dataFrame: DataFrame, nWorkers: Int): DataFrame = {
-    // We can't check dataFrame.rdd.getNumPartitions == nWorkers here, since dataFrame.rdd is
+    // we can't involve any coalesce operation here, since Barrier mode will check
-    // a lazy variable. If we call it here, we will not directly extract RDD[Table] again,
+    // the RDD patterns which does not allow coalesce.
-    // instead, we will involve Columnar -> Row -> Columnar and decrease the performance
+    dataFrame.repartition(nWorkers)
-    if (nWorkers == 1) {
-      dataFrame.coalesce(1)
-    } else {
-      dataFrame.repartition(nWorkers)
-    }
   }
 
   private def repartitionForGroup(
@@ -448,7 +447,7 @@ object GpuPreXGBoost extends PreXGBoostProvider {
   private def buildRDDWatches(
       dataMap: Map[String, ColumnDataBatch],
       xgbExeParams: XGBoostExecutionParams,
-      noEvalSet: Boolean): RDD[Watches] = {
+      noEvalSet: Boolean): RDD[() => Watches] = {
 
     val sc = dataMap(TRAIN_NAME).rawDF.sparkSession.sparkContext
     val maxBin = xgbExeParams.toMap.getOrElse("max_bin", 256).asInstanceOf[Int]
@@ -459,7 +458,7 @@ object GpuPreXGBoost extends PreXGBoostProvider {
       GpuUtils.toColumnarRdd(dataMap(TRAIN_NAME).rawDF).mapPartitions({
         iter =>
           val iterColBatch = iter.map(table => new GpuColumnBatch(table, null))
-          Iterator(buildWatches(
+          Iterator(() => buildWatches(
             PreXGBoost.getCacheDirName(xgbExeParams.useExternalMemory), xgbExeParams.missing,
             colIndicesForTrain, iterColBatch, maxBin))
       })
@@ -469,7 +468,7 @@ object GpuPreXGBoost extends PreXGBoostProvider {
       val nameAndColIndices = dataMap.map(nc => (nc._1, nc._2.colIndices))
       coPartitionForGpu(dataMap, sc, xgbExeParams.numWorkers).mapPartitions {
         nameAndColumnBatchIter =>
-          Iterator(buildWatchesWithEval(
+          Iterator(() => buildWatchesWithEval(
             PreXGBoost.getCacheDirName(xgbExeParams.useExternalMemory), xgbExeParams.missing,
             nameAndColIndices, nameAndColumnBatchIter, maxBin))
       }

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

@@ -39,13 +39,8 @@ trait GpuTestSuite extends FunSuite with TmpFolderSuite {
 
   def enableCsvConf(): SparkConf = {
     new SparkConf()
-      .set(RapidsConf.ENABLE_READ_CSV_DATES.key, "true")
+      .set("spark.rapids.sql.csv.read.float.enabled", "true")
-      .set(RapidsConf.ENABLE_READ_CSV_BYTES.key, "true")
+      .set("spark.rapids.sql.csv.read.double.enabled", "true")
-      .set(RapidsConf.ENABLE_READ_CSV_SHORTS.key, "true")
-      .set(RapidsConf.ENABLE_READ_CSV_INTEGERS.key, "true")
-      .set(RapidsConf.ENABLE_READ_CSV_LONGS.key, "true")
-      .set(RapidsConf.ENABLE_READ_CSV_FLOATS.key, "true")
-      .set(RapidsConf.ENABLE_READ_CSV_DOUBLES.key, "true")
   }
 
   def withGpuSparkSession[U](conf: SparkConf = new SparkConf())(f: SparkSession => U): U = {
@@ -246,12 +241,13 @@ object SparkSessionHolder extends Logging {
     Locale.setDefault(Locale.US)
 
     val builder = SparkSession.builder()
-      .master("local[1]")
+      .master("local[2]")
       .config("spark.sql.adaptive.enabled", "false")
       .config("spark.rapids.sql.enabled", "false")
       .config("spark.rapids.sql.test.enabled", "false")
       .config("spark.plugins", "com.nvidia.spark.SQLPlugin")
       .config("spark.rapids.memory.gpu.pooling.enabled", "false") // Disable RMM for unit tests.
+      .config("spark.sql.files.maxPartitionBytes", "1000")
       .appName("XGBoost4j-Spark-Gpu unit test")
 
     builder.getOrCreate()

jvm-packages/xgboost4j-spark/pom.xml

@@ -6,7 +6,7 @@
     <parent>
         <groupId>ml.dmlc</groupId>
         <artifactId>xgboost-jvm_2.12</artifactId>
-        <version>1.6.0</version>
+        <version>1.6.2</version>
     </parent>
     <artifactId>xgboost4j-spark_2.12</artifactId>
     <build>
@@ -24,7 +24,7 @@
         <dependency>
             <groupId>ml.dmlc</groupId>
             <artifactId>xgboost4j_${scala.binary.version}</artifactId>
-            <version>1.6.0</version>
+            <version>1.6.2</version>
         </dependency>
         <dependency>
             <groupId>org.apache.spark</groupId>

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

@@ -96,19 +96,21 @@ object PreXGBoost extends PreXGBoostProvider {
   }
 
   /**
-   * Convert the Dataset[_] to RDD[Watches] which will be fed to XGBoost
+   * Convert the Dataset[_] to RDD[() => Watches] which will be fed to XGBoost
    *
    * @param estimator supports XGBoostClassifier and XGBoostRegressor
    * @param dataset the training data
    * @param params all user defined and defaulted params
-   * @return [[XGBoostExecutionParams]] => (RDD[[Watches]], Option[ RDD[_] ])
+   * @return [[XGBoostExecutionParams]] => (Boolean, RDD[[() => Watches]], Option[ RDD[_] ])
-   *         RDD[Watches] will be used as the training input
+   *         Boolean if building DMatrix in rabit context
+   *         RDD[() => Watches] will be used as the training input
    *         Option[RDD[_]\] is the optional cached RDD
    */
   override def buildDatasetToRDD(
       estimator: Estimator[_],
       dataset: Dataset[_],
-      params: Map[String, Any]): XGBoostExecutionParams => (RDD[Watches], Option[RDD[_]]) = {
+      params: Map[String, Any]): XGBoostExecutionParams =>
+    (Boolean, RDD[() => Watches], Option[RDD[_]]) = {
 
     if (optionProvider.isDefined && optionProvider.get.providerEnabled(Some(dataset))) {
       return optionProvider.get.buildDatasetToRDD(estimator, dataset, params)
@@ -170,12 +172,12 @@ object PreXGBoost extends PreXGBoostProvider {
           val cachedRDD = if (xgbExecParams.cacheTrainingSet) {
             Some(trainingData.persist(StorageLevel.MEMORY_AND_DISK))
           } else None
-          (trainForRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
+          (false, trainForRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
         case Right(trainingData) =>
           val cachedRDD = if (xgbExecParams.cacheTrainingSet) {
             Some(trainingData.persist(StorageLevel.MEMORY_AND_DISK))
           } else None
-          (trainForNonRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
+          (false, trainForNonRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
       }
 
   }
@@ -311,17 +313,18 @@ object PreXGBoost extends PreXGBoostProvider {
 
 
   /**
-   * Converting the RDD[XGBLabeledPoint] to the function to build RDD[Watches]
+   * Converting the RDD[XGBLabeledPoint] to the function to build RDD[() => Watches]
    *
    * @param trainingSet the input training RDD[XGBLabeledPoint]
    * @param evalRDDMap the eval set
    * @param hasGroup if has group
-   * @return function to build (RDD[Watches], the cached RDD)
+   * @return function to build (RDD[() => Watches], the cached RDD)
    */
   private[spark] def buildRDDLabeledPointToRDDWatches(
       trainingSet: RDD[XGBLabeledPoint],
       evalRDDMap: Map[String, RDD[XGBLabeledPoint]] = Map(),
-      hasGroup: Boolean = false): XGBoostExecutionParams => (RDD[Watches], Option[RDD[_]]) = {
+      hasGroup: Boolean = false):
+  XGBoostExecutionParams => (Boolean, RDD[() => Watches], Option[RDD[_]]) = {
 
     xgbExecParams: XGBoostExecutionParams =>
       composeInputData(trainingSet, hasGroup, xgbExecParams.numWorkers) match {
@@ -329,12 +332,12 @@ object PreXGBoost extends PreXGBoostProvider {
           val cachedRDD = if (xgbExecParams.cacheTrainingSet) {
             Some(trainingData.persist(StorageLevel.MEMORY_AND_DISK))
           } else None
-          (trainForRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
+          (false, trainForRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
         case Right(trainingData) =>
           val cachedRDD = if (xgbExecParams.cacheTrainingSet) {
             Some(trainingData.persist(StorageLevel.MEMORY_AND_DISK))
           } else None
-          (trainForNonRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
+          (false, trainForNonRanking(trainingData, xgbExecParams, evalRDDMap), cachedRDD)
       }
   }
 
@@ -374,34 +377,34 @@ object PreXGBoost extends PreXGBoostProvider {
   }
 
   /**
-   * Build RDD[Watches] for Ranking
+   * Build RDD[() => Watches] for Ranking
    * @param trainingData the training data RDD
    * @param xgbExecutionParams xgboost execution params
    * @param evalSetsMap the eval RDD
-   * @return RDD[Watches]
+   * @return RDD[() => Watches]
    */
   private def trainForRanking(
       trainingData: RDD[Array[XGBLabeledPoint]],
       xgbExecutionParam: XGBoostExecutionParams,
-      evalSetsMap: Map[String, RDD[XGBLabeledPoint]]): RDD[Watches] = {
+      evalSetsMap: Map[String, RDD[XGBLabeledPoint]]): RDD[() => Watches] = {
     if (evalSetsMap.isEmpty) {
       trainingData.mapPartitions(labeledPointGroups => {
-        val watches = Watches.buildWatchesWithGroup(xgbExecutionParam,
+        val buildWatches = () => Watches.buildWatchesWithGroup(xgbExecutionParam,
           DataUtils.processMissingValuesWithGroup(labeledPointGroups, xgbExecutionParam.missing,
             xgbExecutionParam.allowNonZeroForMissing),
           getCacheDirName(xgbExecutionParam.useExternalMemory))
-        Iterator.single(watches)
+        Iterator.single(buildWatches)
       }).cache()
     } else {
       coPartitionGroupSets(trainingData, evalSetsMap, xgbExecutionParam.numWorkers).mapPartitions(
         labeledPointGroupSets => {
-          val watches = Watches.buildWatchesWithGroup(
+          val buildWatches = () => Watches.buildWatchesWithGroup(
             labeledPointGroupSets.map {
               case (name, iter) => (name, DataUtils.processMissingValuesWithGroup(iter,
                 xgbExecutionParam.missing, xgbExecutionParam.allowNonZeroForMissing))
             },
             getCacheDirName(xgbExecutionParam.useExternalMemory))
-          Iterator.single(watches)
+          Iterator.single(buildWatches)
         }).cache()
     }
   }
@@ -462,35 +465,35 @@ object PreXGBoost extends PreXGBoostProvider {
   }
 
   /**
-   * Build RDD[Watches] for Non-Ranking
+   * Build RDD[() => Watches] for Non-Ranking
    * @param trainingData the training data RDD
    * @param xgbExecutionParams xgboost execution params
    * @param evalSetsMap the eval RDD
-   * @return RDD[Watches]
+   * @return RDD[() => Watches]
    */
   private def trainForNonRanking(
       trainingData: RDD[XGBLabeledPoint],
       xgbExecutionParams: XGBoostExecutionParams,
-      evalSetsMap: Map[String, RDD[XGBLabeledPoint]]): RDD[Watches] = {
+      evalSetsMap: Map[String, RDD[XGBLabeledPoint]]): RDD[() => Watches] = {
     if (evalSetsMap.isEmpty) {
       trainingData.mapPartitions { labeledPoints => {
-        val watches = Watches.buildWatches(xgbExecutionParams,
+        val buildWatches = () => Watches.buildWatches(xgbExecutionParams,
           DataUtils.processMissingValues(labeledPoints, xgbExecutionParams.missing,
             xgbExecutionParams.allowNonZeroForMissing),
           getCacheDirName(xgbExecutionParams.useExternalMemory))
-        Iterator.single(watches)
+        Iterator.single(buildWatches)
       }}.cache()
     } else {
       coPartitionNoGroupSets(trainingData, evalSetsMap, xgbExecutionParams.numWorkers).
         mapPartitions {
           nameAndLabeledPointSets =>
-            val watches = Watches.buildWatches(
+            val buildWatches = () => Watches.buildWatches(
               nameAndLabeledPointSets.map {
                 case (name, iter) => (name, DataUtils.processMissingValues(iter,
                   xgbExecutionParams.missing, xgbExecutionParams.allowNonZeroForMissing))
               },
               getCacheDirName(xgbExecutionParams.useExternalMemory))
-            Iterator.single(watches)
+            Iterator.single(buildWatches)
         }.cache()
     }
   }

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

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2021 by Contributors
+ Copyright (c) 2021-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.
@@ -45,19 +45,21 @@ private[scala] trait PreXGBoostProvider {
   def transformSchema(xgboostEstimator: XGBoostEstimatorCommon, schema: StructType): StructType
 
   /**
-   * Convert the Dataset[_] to RDD[Watches] which will be fed to XGBoost
+   * Convert the Dataset[_] to RDD[() => Watches] which will be fed to XGBoost
    *
    * @param estimator supports XGBoostClassifier and XGBoostRegressor
    * @param dataset the training data
    * @param params all user defined and defaulted params
-   * @return [[XGBoostExecutionParams]] => (RDD[[Watches]], Option[ RDD[_] ])
+   * @return [[XGBoostExecutionParams]] => (Boolean, RDD[[() => Watches]], Option[ RDD[_] ])
-   *         RDD[Watches] will be used as the training input
+   *         Boolean if building DMatrix in rabit context
+   *         RDD[() => Watches] will be used as the training input to build DMatrix
    *         Option[ RDD[_] ] is the optional cached RDD
    */
   def buildDatasetToRDD(
     estimator: Estimator[_],
     dataset: Dataset[_],
-    params: Map[String, Any]): XGBoostExecutionParams => (RDD[Watches], Option[RDD[_]])
+    params: Map[String, Any]):
+  XGBoostExecutionParams => (Boolean, RDD[() => Watches], Option[RDD[_]])
 
   /**
    * Transform Dataset

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

@@ -21,6 +21,7 @@ import java.io.File
 import scala.collection.mutable
 import scala.util.Random
 import scala.collection.JavaConverters._
+
 import ml.dmlc.xgboost4j.java.{IRabitTracker, Rabit, XGBoostError, RabitTracker => PyRabitTracker}
 import ml.dmlc.xgboost4j.scala.rabit.RabitTracker
 import ml.dmlc.xgboost4j.scala.spark.params.LearningTaskParams
@@ -30,8 +31,9 @@ import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
 import org.apache.commons.io.FileUtils
 import org.apache.commons.logging.LogFactory
 import org.apache.hadoop.fs.FileSystem
+
 import org.apache.spark.rdd.RDD
-import org.apache.spark.{SparkContext, SparkParallelismTracker, TaskContext}
+import org.apache.spark.{SparkContext, TaskContext}
 import org.apache.spark.sql.SparkSession
 
 /**
@@ -79,8 +81,7 @@ private[scala] case class XGBoostExecutionParams(
     earlyStoppingParams: XGBoostExecutionEarlyStoppingParams,
     cacheTrainingSet: Boolean,
     treeMethod: Option[String],
-    isLocal: Boolean,
+    isLocal: Boolean) {
-    killSparkContextOnWorkerFailure: Boolean) {
 
   private var rawParamMap: Map[String, Any] = _
 
@@ -224,9 +225,6 @@ private[this] class XGBoostExecutionParamsFactory(rawParams: Map[String, Any], s
     val cacheTrainingSet = overridedParams.getOrElse("cache_training_set", false)
       .asInstanceOf[Boolean]
 
-    val killSparkContext = overridedParams.getOrElse("kill_spark_context_on_worker_failure", true)
-      .asInstanceOf[Boolean]
-
     val xgbExecParam = XGBoostExecutionParams(nWorkers, round, useExternalMemory, obj, eval,
       missing, allowNonZeroForMissing, trackerConf,
       timeoutRequestWorkers,
@@ -235,8 +233,7 @@ private[this] class XGBoostExecutionParamsFactory(rawParams: Map[String, Any], s
       xgbExecEarlyStoppingParams,
       cacheTrainingSet,
       treeMethod,
-      isLocal,
+      isLocal)
-      killSparkContext)
     xgbExecParam.setRawParamMap(overridedParams)
     xgbExecParam
   }
@@ -283,13 +280,8 @@ object XGBoost extends Serializable {
     }
   }
 
-  private def buildDistributedBooster(
+  private def buildWatchesAndCheck(buildWatchesFun: () => Watches): Watches = {
-      watches: Watches,
+    val watches = buildWatchesFun()
-      xgbExecutionParam: XGBoostExecutionParams,
-      rabitEnv: java.util.Map[String, String],
-      obj: ObjectiveTrait,
-      eval: EvalTrait,
-      prevBooster: Booster): Iterator[(Booster, Map[String, Array[Float]])] = {
     // to workaround the empty partitions in training dataset,
     // this might not be the best efficient implementation, see
     // (https://github.com/dmlc/xgboost/issues/1277)
@@ -298,14 +290,39 @@ object XGBoost extends Serializable {
         s"detected an empty partition in the training data, partition ID:" +
           s" ${TaskContext.getPartitionId()}")
     }
+    watches
+  }
+
+  private def buildDistributedBooster(
+      buildDMatrixInRabit: Boolean,
+      buildWatches: () => Watches,
+      xgbExecutionParam: XGBoostExecutionParams,
+      rabitEnv: java.util.Map[String, String],
+      obj: ObjectiveTrait,
+      eval: EvalTrait,
+      prevBooster: Booster): Iterator[(Booster, Map[String, Array[Float]])] = {
+
+    var watches: Watches = null
+    if (!buildDMatrixInRabit) {
+      // for CPU pipeline, we need to build DMatrix out of rabit context
+      watches = buildWatchesAndCheck(buildWatches)
+    }
+
     val taskId = TaskContext.getPartitionId().toString
     val attempt = TaskContext.get().attemptNumber.toString
     rabitEnv.put("DMLC_TASK_ID", taskId)
     rabitEnv.put("DMLC_NUM_ATTEMPT", attempt)
     val numRounds = xgbExecutionParam.numRounds
     val makeCheckpoint = xgbExecutionParam.checkpointParam.isDefined && taskId.toInt == 0
+
     try {
       Rabit.init(rabitEnv)
+
+      if (buildDMatrixInRabit) {
+        // for GPU pipeline, we need to move dmatrix building into rabit context
+        watches = buildWatchesAndCheck(buildWatches)
+      }
+
       val numEarlyStoppingRounds = xgbExecutionParam.earlyStoppingParams.numEarlyStoppingRounds
       val metrics = Array.tabulate(watches.size)(_ => Array.ofDim[Float](numRounds))
       val externalCheckpointParams = xgbExecutionParam.checkpointParam
@@ -331,14 +348,18 @@ object XGBoost extends Serializable {
           watches.toMap, metrics, obj, eval,
           earlyStoppingRound = numEarlyStoppingRounds, prevBooster)
       }
-      Iterator(booster -> watches.toMap.keys.zip(metrics).toMap)
+      if (TaskContext.get().partitionId() == 0) {
+        Iterator(booster -> watches.toMap.keys.zip(metrics).toMap)
+      } else {
+        Iterator.empty
+      }
     } catch {
       case xgbException: XGBoostError =>
         logger.error(s"XGBooster worker $taskId has failed $attempt times due to ", xgbException)
         throw xgbException
     } finally {
       Rabit.shutdown()
-      watches.delete()
+      if (watches != null) watches.delete()
     }
   }
 
@@ -364,7 +385,7 @@ object XGBoost extends Serializable {
   @throws(classOf[XGBoostError])
   private[spark] def trainDistributed(
       sc: SparkContext,
-      buildTrainingData: XGBoostExecutionParams => (RDD[Watches], Option[RDD[_]]),
+      buildTrainingData: XGBoostExecutionParams => (Boolean, RDD[() => Watches], Option[RDD[_]]),
       params: Map[String, Any]):
     (Booster, Map[String, Array[Float]]) = {
 
@@ -383,50 +404,36 @@ object XGBoost extends Serializable {
     }.orNull
 
     // Get the training data RDD and the cachedRDD
-    val (trainingRDD, optionalCachedRDD) = buildTrainingData(xgbExecParams)
+    val (buildDMatrixInRabit, trainingRDD, optionalCachedRDD) = buildTrainingData(xgbExecParams)
 
     try {
       // Train for every ${savingRound} rounds and save the partially completed booster
       val tracker = startTracker(xgbExecParams.numWorkers, xgbExecParams.trackerConf)
       val (booster, metrics) = try {
-        val parallelismTracker = new SparkParallelismTracker(sc,
-          xgbExecParams.timeoutRequestWorkers,
-          xgbExecParams.numWorkers,
-          xgbExecParams.killSparkContextOnWorkerFailure)
-
         tracker.getWorkerEnvs().putAll(xgbRabitParams)
         val rabitEnv = tracker.getWorkerEnvs
 
-        val boostersAndMetrics = trainingRDD.mapPartitions { iter => {
+        val boostersAndMetrics = trainingRDD.barrier().mapPartitions { iter => {
-          var optionWatches: Option[Watches] = None
+          var optionWatches: Option[() => Watches] = None
 
           // take the first Watches to train
           if (iter.hasNext) {
             optionWatches = Some(iter.next())
           }
 
-          optionWatches.map { watches => buildDistributedBooster(watches, xgbExecParams, rabitEnv,
+          optionWatches.map { buildWatches => buildDistributedBooster(buildDMatrixInRabit,
-            xgbExecParams.obj, xgbExecParams.eval, prevBooster)}
+            buildWatches, xgbExecParams, rabitEnv, xgbExecParams.obj,
+            xgbExecParams.eval, prevBooster)}
             .getOrElse(throw new RuntimeException("No Watches to train"))
 
-        }}.cache()
+        }}
-
-        val sparkJobThread = new Thread() {
-          override def run() {
-            // force the job
-            boostersAndMetrics.foreachPartition(() => _)
-          }
-        }
-        sparkJobThread.setUncaughtExceptionHandler(tracker)
-
-        val trackerReturnVal = parallelismTracker.execute {
-          sparkJobThread.start()
-          tracker.waitFor(0L)
-        }
 
+        val (booster, metrics) = boostersAndMetrics.collect()(0)
+        val trackerReturnVal = tracker.waitFor(0L)
         logger.info(s"Rabit returns with exit code $trackerReturnVal")
-        val (booster, metrics) = postTrackerReturnProcessing(trackerReturnVal,
+        if (trackerReturnVal != 0) {
-          boostersAndMetrics, sparkJobThread)
+          throw new XGBoostError("XGBoostModel training failed.")
+        }
         (booster, metrics)
       } finally {
         tracker.stop()
@@ -446,42 +453,12 @@ object XGBoost extends Serializable {
       case t: Throwable =>
         // if the job was aborted due to an exception
         logger.error("the job was aborted due to ", t)
-        if (xgbExecParams.killSparkContextOnWorkerFailure) {
-          sc.stop()
-        }
         throw t
     } finally {
       optionalCachedRDD.foreach(_.unpersist())
     }
   }
 
-  private def postTrackerReturnProcessing(
-      trackerReturnVal: Int,
-      distributedBoostersAndMetrics: RDD[(Booster, Map[String, Array[Float]])],
-      sparkJobThread: Thread): (Booster, Map[String, Array[Float]]) = {
-    if (trackerReturnVal == 0) {
-      // Copies of the final booster and the corresponding metrics
-      // reside in each partition of the `distributedBoostersAndMetrics`.
-      // Any of them can be used to create the model.
-      // it's safe to block here forever, as the tracker has returned successfully, and the Spark
-      // job should have finished, there is no reason for the thread cannot return
-      sparkJobThread.join()
-      val (booster, metrics) = distributedBoostersAndMetrics.first()
-      distributedBoostersAndMetrics.unpersist(false)
-      (booster, metrics)
-    } else {
-      try {
-        if (sparkJobThread.isAlive) {
-          sparkJobThread.interrupt()
-        }
-      } catch {
-        case _: InterruptedException =>
-          logger.info("spark job thread is interrupted")
-      }
-      throw new XGBoostError("XGBoostModel training failed")
-    }
-  }
-
 }
 
 class Watches private[scala] (

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

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,18 +16,22 @@
 
 package ml.dmlc.xgboost4j.scala.spark.params
 
+import ml.dmlc.xgboost4j.scala.spark
+import org.apache.commons.logging.LogFactory
 import org.apache.hadoop.fs.Path
 import org.json4s.{DefaultFormats, JValue}
 import org.json4s.JsonAST.JObject
 import org.json4s.jackson.JsonMethods.{compact, parse, render}
 
 import org.apache.spark.SparkContext
-import org.apache.spark.ml.param.{Param, Params}
+import org.apache.spark.ml.param.Params
 import org.apache.spark.ml.util.MLReader
 
 // This originates from apache-spark DefaultPramsReader copy paste
 private[spark] object DefaultXGBoostParamsReader {
 
+  private val logger = LogFactory.getLog("XGBoostSpark")
+
   private val paramNameCompatibilityMap: Map[String, String] = Map("silent" -> "verbosity")
 
   private val paramValueCompatibilityMap: Map[String, Map[Any, Any]] =
@@ -126,9 +130,16 @@ private[spark] object DefaultXGBoostParamsReader {
     metadata.params match {
       case JObject(pairs) =>
         pairs.foreach { case (paramName, jsonValue) =>
-          val param = instance.getParam(handleBrokenlyChangedName(paramName))
+          val finalName = handleBrokenlyChangedName(paramName)
-          val value = param.jsonDecode(compact(render(jsonValue)))
+          // For the deleted parameters, we'd better to remove it instead of throwing an exception.
-          instance.set(param, handleBrokenlyChangedValue(paramName, value))
+          // So we need to check if the parameter exists instead of blindly setting it.
+          if (instance.hasParam(finalName)) {
+            val param = instance.getParam(finalName)
+            val value = param.jsonDecode(compact(render(jsonValue)))
+            instance.set(param, handleBrokenlyChangedValue(paramName, value))
+          } else {
+            logger.warn(s"$finalName is no longer used in ${spark.VERSION}")
+          }
         }
       case _ =>
         throw new IllegalArgumentException(

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

@@ -105,14 +105,8 @@ private[spark] trait LearningTaskParams extends Params {
 
   final def getMaximizeEvaluationMetrics: Boolean = $(maximizeEvaluationMetrics)
 
-  /**
-   * whether killing SparkContext when training task fails
-   */
-  final val killSparkContextOnWorkerFailure = new BooleanParam(this,
-    "killSparkContextOnWorkerFailure", "whether killing SparkContext when training task fails")
-
   setDefault(objective -> "reg:squarederror", baseScore -> 0.5, trainTestRatio -> 1.0,
-    numEarlyStoppingRounds -> 0, cacheTrainingSet -> false, killSparkContextOnWorkerFailure -> true)
+    numEarlyStoppingRounds -> 0, cacheTrainingSet -> false)
 }
 
 private[spark] object LearningTaskParams {

jvm-packages/xgboost4j-spark/src/main/scala/org/apache/spark/SparkParallelismTracker.scala

@@ -1,175 +0,0 @@
-/*
- Copyright (c) 2014 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 org.apache.spark
-
-import org.apache.commons.logging.LogFactory
-import org.apache.spark.scheduler._
-
-import scala.collection.mutable.{HashMap, HashSet}
-
-/**
- * A tracker that ensures enough number of executor cores are alive.
- * Throws an exception when the number of alive cores is less than nWorkers.
- *
- * @param sc The SparkContext object
- * @param timeout The maximum time to wait for enough number of workers.
- * @param numWorkers nWorkers used in an XGBoost Job
- * @param killSparkContextOnWorkerFailure kill SparkContext or not when task fails
- */
-class SparkParallelismTracker(
-    val sc: SparkContext,
-    timeout: Long,
-    numWorkers: Int,
-    killSparkContextOnWorkerFailure: Boolean = true) {
-
-  private[this] val requestedCores = numWorkers * sc.conf.getInt("spark.task.cpus", 1)
-  private[this] val logger = LogFactory.getLog("XGBoostSpark")
-
-  private[this] def numAliveCores: Int = {
-    sc.statusStore.executorList(true).map(_.totalCores).sum
-  }
-
-  private[this] def waitForCondition(
-      condition: => Boolean,
-      timeout: Long,
-      checkInterval: Long = 100L) = {
-    val waitImpl = new ((Long, Boolean) => Boolean) {
-      override def apply(waitedTime: Long, status: Boolean): Boolean = status match {
-        case s if s => true
-        case _ => waitedTime match {
-          case t if t < timeout =>
-            Thread.sleep(checkInterval)
-            apply(t + checkInterval, status = condition)
-          case _ => false
-        }
-      }
-    }
-    waitImpl(0L, condition)
-  }
-
-  private[this] def safeExecute[T](body: => T): T = {
-    val listener = new TaskFailedListener(killSparkContextOnWorkerFailure)
-    sc.addSparkListener(listener)
-    try {
-      body
-    } finally {
-      sc.removeSparkListener(listener)
-    }
-  }
-
-  /**
-   * Execute a blocking function call with two checks on enough nWorkers:
-   *  - Before the function starts, wait until there are enough executor cores.
-   *  - During the execution, throws an exception if there is any executor lost.
-   *
-   * @param body A blocking function call
-   * @tparam T Return type
-   * @return The return of body
-   */
-  def execute[T](body: => T): T = {
-    if (timeout <= 0) {
-      logger.info("starting training without setting timeout for waiting for resources")
-      safeExecute(body)
-    } else {
-      logger.info(s"starting training with timeout set as $timeout ms for waiting for resources")
-      if (!waitForCondition(numAliveCores >= requestedCores, timeout)) {
-        throw new IllegalStateException(s"Unable to get $requestedCores cores for XGBoost training")
-      }
-      safeExecute(body)
-    }
-  }
-}
-
-class TaskFailedListener(killSparkContext: Boolean = true) extends SparkListener {
-
-  private[this] val logger = LogFactory.getLog("XGBoostTaskFailedListener")
-
-  // {jobId, [stageId0, stageId1, ...] }
-  // keep track of the mapping of job id and stage ids
-  // when a task fails, find the job id and stage id the task belongs to, finally
-  // cancel the jobs
-  private val jobIdToStageIds: HashMap[Int, HashSet[Int]] = HashMap.empty
-
-  override def onJobStart(jobStart: SparkListenerJobStart): Unit = {
-    if (!killSparkContext) {
-      jobStart.stageIds.foreach(stageId => {
-        jobIdToStageIds.getOrElseUpdate(jobStart.jobId, new HashSet[Int]()) += stageId
-      })
-    }
-  }
-
-  override def onJobEnd(jobEnd: SparkListenerJobEnd): Unit = {
-    if (!killSparkContext) {
-      jobIdToStageIds.remove(jobEnd.jobId)
-    }
-  }
-
-  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
-    taskEnd.reason match {
-      case taskEndReason: TaskFailedReason =>
-        logger.error(s"Training Task Failed during XGBoost Training: " +
-            s"$taskEndReason")
-        if (killSparkContext) {
-          logger.error("killing SparkContext")
-          TaskFailedListener.startedSparkContextKiller()
-        } else {
-          val stageId = taskEnd.stageId
-          // find job ids according to stage id and then cancel the job
-
-          jobIdToStageIds.foreach {
-            case (jobId, stageIds) =>
-              if (stageIds.contains(stageId)) {
-                logger.error("Cancelling jobId:" + jobId)
-                jobIdToStageIds.remove(jobId)
-                SparkContext.getOrCreate().cancelJob(jobId)
-              }
-          }
-        }
-      case _ =>
-    }
-  }
-}
-
-object TaskFailedListener {
-
-  var killerStarted: Boolean = false
-
-  var sparkContextKiller: Thread = _
-
-  val sparkContextShutdownLock = new AnyRef
-
-  private def startedSparkContextKiller(): Unit = this.synchronized {
-    if (!killerStarted) {
-      killerStarted = true
-      // Spark does not allow ListenerThread to shutdown SparkContext so that we have to do it
-      // in a separate thread
-      sparkContextKiller = new Thread() {
-        override def run(): Unit = {
-          LiveListenerBus.withinListenerThread.withValue(false) {
-            sparkContextShutdownLock.synchronized {
-              SparkContext.getActive.foreach(_.stop())
-              killerStarted = false
-              sparkContextShutdownLock.notify()
-            }
-          }
-        }
-      }
-      sparkContextKiller.setDaemon(true)
-      sparkContextKiller.start()
-    }
-  }
-}

jvm-packages/xgboost4j-spark/src/test/resources/log4j.properties

@@ -1 +1 @@
-log4j.logger.org.apache.spark=ERROR
+log4j.logger.org.apache.spark=ERROR

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/ExternalCheckpointManagerSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -19,7 +19,7 @@ package ml.dmlc.xgboost4j.scala.spark
 import java.io.File
 
 import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, ExternalCheckpointManager, XGBoost => SXGBoost}
-import org.scalatest.{FunSuite, Ignore}
+import org.scalatest.FunSuite
 import org.apache.hadoop.fs.{FileSystem, Path}
 
 class ExternalCheckpointManagerSuite extends FunSuite with TmpFolderPerSuite with PerTest {

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/FeatureSizeValidatingSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,10 +16,8 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
 import org.apache.spark.Partitioner
 import org.apache.spark.ml.feature.VectorAssembler
-import org.apache.spark.sql.SparkSession
 import org.scalatest.FunSuite
 import org.apache.spark.sql.functions._
 
@@ -53,7 +51,7 @@ class FeatureSizeValidatingSuite extends FunSuite with PerTest {
       "objective" -> "binary:logistic",
       "num_round" -> 5, "num_workers" -> 2, "use_external_memory" -> true, "missing" -> 0)
     import DataUtils._
-    val sparkSession = SparkSession.builder().getOrCreate()
+    val sparkSession = ss
     import sparkSession.implicits._
     val repartitioned = sc.parallelize(Synthetic.trainWithDiffFeatureSize, 2)
       .map(lp => (lp.label, lp)).partitionBy(

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/MissingValueHandlingSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,14 +16,14 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
 import org.apache.spark.ml.feature.VectorAssembler
 import org.apache.spark.ml.linalg.Vectors
 import org.apache.spark.sql.DataFrame
 import org.scalatest.FunSuite
-
 import scala.util.Random
 
+import org.apache.spark.SparkException
+
 class MissingValueHandlingSuite extends FunSuite with PerTest {
   test("dense vectors containing missing value") {
     def buildDenseDataFrame(): DataFrame = {
@@ -113,7 +113,7 @@ class MissingValueHandlingSuite extends FunSuite with PerTest {
     val inputDF = vectorAssembler.transform(testDF).select("features", "label")
     val paramMap = List("eta" -> "1", "max_depth" -> "2",
       "objective" -> "binary:logistic", "missing" -> -1.0f, "num_workers" -> 1).toMap
-    intercept[XGBoostError] {
+    intercept[SparkException] {
       new XGBoostClassifier(paramMap).fit(inputDF)
     }
   }
@@ -140,7 +140,7 @@ class MissingValueHandlingSuite extends FunSuite with PerTest {
     inputDF.show()
     val paramMap = List("eta" -> "1", "max_depth" -> "2",
       "objective" -> "binary:logistic", "missing" -> -1.0f, "num_workers" -> 1).toMap
-    intercept[XGBoostError] {
+    intercept[SparkException] {
       new XGBoostClassifier(paramMap).fit(inputDF)
     }
   }

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/ParameterSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,9 +16,9 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
+import org.scalatest.{BeforeAndAfterAll, FunSuite}
-import org.scalatest.{BeforeAndAfterAll, FunSuite, Ignore}
 
+import org.apache.spark.SparkException
 import org.apache.spark.ml.param.ParamMap
 
 class ParameterSuite extends FunSuite with PerTest with BeforeAndAfterAll {
@@ -40,28 +40,16 @@ class ParameterSuite extends FunSuite with PerTest with BeforeAndAfterAll {
     assert(xgbCopy2.MLlib2XGBoostParams("eval_metric").toString === "logloss")
   }
 
-  private def waitForSparkContextShutdown(): Unit = {
-    var totalWaitedTime = 0L
-    while (!ss.sparkContext.isStopped && totalWaitedTime <= 120000) {
-      Thread.sleep(10000)
-      totalWaitedTime += 10000
-    }
-    assert(ss.sparkContext.isStopped === true)
-  }
-
   test("fail training elegantly with unsupported objective function") {
     val paramMap = Map("eta" -> "0.1", "max_depth" -> "6", "silent" -> "1",
       "objective" -> "wrong_objective_function", "num_class" -> "6", "num_round" -> 5,
       "num_workers" -> numWorkers)
     val trainingDF = buildDataFrame(MultiClassification.train)
     val xgb = new XGBoostClassifier(paramMap)
-    try {
+    intercept[SparkException] {
-      val model = xgb.fit(trainingDF)
+      xgb.fit(trainingDF)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      waitForSparkContextShutdown()
     }
+
   }
 
   test("fail training elegantly with unsupported eval metrics") {
@@ -70,12 +58,8 @@ class ParameterSuite extends FunSuite with PerTest with BeforeAndAfterAll {
       "num_workers" -> numWorkers, "eval_metric" -> "wrong_eval_metrics")
     val trainingDF = buildDataFrame(MultiClassification.train)
     val xgb = new XGBoostClassifier(paramMap)
-    try {
+    intercept[SparkException] {
-      val model = xgb.fit(trainingDF)
+      xgb.fit(trainingDF)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      waitForSparkContextShutdown()
     }
   }
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/PerTest.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -19,7 +19,7 @@ package ml.dmlc.xgboost4j.scala.spark
 import java.io.File
 
 import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
-import org.apache.spark.{SparkConf, SparkContext, TaskFailedListener}
+import org.apache.spark.SparkContext
 import org.apache.spark.sql._
 import org.scalatest.{BeforeAndAfterEach, FunSuite}
 
@@ -40,32 +40,16 @@ trait PerTest extends BeforeAndAfterEach { self: FunSuite =>
       .appName("XGBoostSuite")
       .config("spark.ui.enabled", false)
       .config("spark.driver.memory", "512m")
+      .config("spark.barrier.sync.timeout", 10)
       .config("spark.task.cpus", 1)
 
   override def beforeEach(): Unit = getOrCreateSession
 
   override def afterEach() {
-    TaskFailedListener.sparkContextShutdownLock.synchronized {
+    if (currentSession != null) {
-      if (currentSession != null) {
+      currentSession.stop()
-        // this synchronization is mostly for the tests involving SparkContext shutdown
+      cleanExternalCache(currentSession.sparkContext.appName)
-        // for unit test involving the sparkContext shutdown there are two different events sequence
+      currentSession = null
-        // 1. SparkContext killer is executed before afterEach, in this case, before SparkContext
-        // is fully stopped, afterEach() will block at the following code block
-        // 2. SparkContext killer is executed afterEach, in this case, currentSession.stop() in will
-        // block to wait for all msgs in ListenerBus get processed. Because currentSession.stop()
-        // has been called, SparkContext killer will not take effect
-        while (TaskFailedListener.killerStarted) {
-          TaskFailedListener.sparkContextShutdownLock.wait()
-        }
-        currentSession.stop()
-        cleanExternalCache(currentSession.sparkContext.appName)
-        currentSession = null
-      }
-      if (TaskFailedListener.sparkContextKiller != null) {
-        TaskFailedListener.sparkContextKiller.interrupt()
-        TaskFailedListener.sparkContextKiller = null
-      }
-      TaskFailedListener.killerStarted = false
     }
   }
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/PersistenceSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014,2021 by Contributors
+ 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.

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostConfigureSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,10 +16,8 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.Rabit
 import ml.dmlc.xgboost4j.scala.{Booster, DMatrix}
 
-import scala.collection.JavaConverters._
 import org.apache.spark.sql._
 import org.scalatest.FunSuite
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostGeneralSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,13 +16,12 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.XGBoostError
 import scala.util.Random
 
 import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
 import ml.dmlc.xgboost4j.scala.DMatrix
 
-import org.apache.spark.TaskContext
+import org.apache.spark.{SparkException, TaskContext}
 import org.scalatest.FunSuite
 
 import org.apache.spark.ml.feature.VectorAssembler
@@ -375,13 +374,14 @@ class XGBoostGeneralSuite extends FunSuite with TmpFolderPerSuite with PerTest {
 
   test("throw exception for empty partition in trainingset") {
     val paramMap = Map("eta" -> "0.1", "max_depth" -> "6", "silent" -> "1",
-      "objective" -> "multi:softmax", "num_class" -> "2", "num_round" -> 5,
+      "objective" -> "binary:logistic", "num_class" -> "2", "num_round" -> 5,
-      "num_workers" -> numWorkers, "tree_method" -> "auto")
+      "num_workers" -> numWorkers, "tree_method" -> "auto", "allow_non_zero_for_missing" -> true)
     // The Dmatrix will be empty
-    val trainingDF = buildDataFrame(Seq(XGBLabeledPoint(1.0f, 1, Array(), Array())))
+    val trainingDF = buildDataFrame(Seq(XGBLabeledPoint(1.0f, 4,
+      Array(0, 1, 2, 3), Array(0, 1, 2, 3))))
     val xgb = new XGBoostClassifier(paramMap)
-    intercept[XGBoostError] {
+    intercept[SparkException] {
-      val model = xgb.fit(trainingDF)
+      xgb.fit(trainingDF)
     }
   }
 

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRabitRegressionSuite.scala

@@ -1,5 +1,5 @@
 /*
- Copyright (c) 2014 by Contributors
+ 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.
@@ -16,14 +16,15 @@
 
 package ml.dmlc.xgboost4j.scala.spark
 
-import ml.dmlc.xgboost4j.java.{Rabit, XGBoostError}
+import ml.dmlc.xgboost4j.java.Rabit
-import ml.dmlc.xgboost4j.scala.{Booster, DMatrix}
+import ml.dmlc.xgboost4j.scala.Booster
-import org.apache.spark.TaskFailedListener
-import org.apache.spark.SparkException
 import scala.collection.JavaConverters._
+
 import org.apache.spark.sql._
 import org.scalatest.FunSuite
 
+import org.apache.spark.SparkException
+
 class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
   val predictionErrorMin = 0.00001f
   val maxFailure = 2;
@@ -33,15 +34,6 @@ class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
     .config("spark.kryo.classesToRegister", classOf[Booster].getName)
     .master(s"local[${numWorkers},${maxFailure}]")
 
-  private def waitAndCheckSparkShutdown(waitMiliSec: Int): Boolean = {
-    var totalWaitedTime = 0L
-    while (!ss.sparkContext.isStopped && totalWaitedTime <= waitMiliSec) {
-      Thread.sleep(10)
-      totalWaitedTime += 10
-    }
-    return ss.sparkContext.isStopped
-  }
-
   test("test classification prediction parity w/o ring reduce") {
     val training = buildDataFrame(Classification.train)
     val testDF = buildDataFrame(Classification.test)
@@ -91,14 +83,11 @@ class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
   }
 
   test("test rabit timeout fail handle") {
-    // disable spark kill listener to verify if rabit_timeout take effect and kill tasks
-    TaskFailedListener.killerStarted = true
-
     val training = buildDataFrame(Classification.train)
     // mock rank 0 failure during 8th allreduce synchronization
     Rabit.mockList = Array("0,8,0,0").toList.asJava
 
-    try {
+    intercept[SparkException] {
       new XGBoostClassifier(Map(
         "eta" -> "0.1",
         "max_depth" -> "10",
@@ -108,37 +97,7 @@ class XGBoostRabitRegressionSuite extends FunSuite with PerTest {
         "num_workers" -> numWorkers,
         "rabit_timeout" -> 0))
         .fit(training)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      // assume all tasks throw exception almost same time
-      // 100ms should be enough to exhaust all retries
-      assert(waitAndCheckSparkShutdown(100) == true)
-      TaskFailedListener.killerStarted = false
     }
   }
 
-  test("test SparkContext should not be killed ") {
-    val training = buildDataFrame(Classification.train)
-    // mock rank 0 failure during 8th allreduce synchronization
-    Rabit.mockList = Array("0,8,0,0").toList.asJava
-
-    try {
-      new XGBoostClassifier(Map(
-        "eta" -> "0.1",
-        "max_depth" -> "10",
-        "verbosity" -> "1",
-        "objective" -> "binary:logistic",
-        "num_round" -> 5,
-        "num_workers" -> numWorkers,
-        "kill_spark_context_on_worker_failure" -> false,
-        "rabit_timeout" -> 0))
-        .fit(training)
-    } catch {
-      case e: Throwable => // swallow anything
-    } finally {
-      // wait 3s to check if SparkContext is killed
-      assert(waitAndCheckSparkShutdown(3000) == false)
-    }
-  }
 }

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressorSuite.scala

@@ -21,7 +21,6 @@ import ml.dmlc.xgboost4j.scala.{DMatrix, XGBoost => ScalaXGBoost}
 import org.apache.spark.ml.linalg.{Vector, Vectors}
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.types._
 import org.scalatest.FunSuite
 
 import org.apache.spark.ml.feature.VectorAssembler

jvm-packages/xgboost4j-spark/src/test/scala/org/apache/spark/SparkParallelismTrackerSuite.scala

@@ -1,151 +0,0 @@
-/*
- Copyright (c) 2014 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 org.apache.spark
-
-import org.scalatest.FunSuite
-import _root_.ml.dmlc.xgboost4j.scala.spark.PerTest
-import org.apache.spark.rdd.RDD
-import org.apache.spark.sql.SparkSession
-
-import scala.math.min
-
-class SparkParallelismTrackerSuite extends FunSuite with PerTest {
-
-  val numParallelism: Int = min(Runtime.getRuntime.availableProcessors(), 4)
-
-  override protected def sparkSessionBuilder: SparkSession.Builder = SparkSession.builder()
-    .master(s"local[${numParallelism}]")
-    .appName("XGBoostSuite")
-    .config("spark.ui.enabled", true)
-    .config("spark.driver.memory", "512m")
-    .config("spark.task.cpus", 1)
-
-  private def waitAndCheckSparkShutdown(waitMiliSec: Int): Boolean = {
-    var totalWaitedTime = 0L
-    while (!ss.sparkContext.isStopped && totalWaitedTime <= waitMiliSec) {
-      Thread.sleep(100)
-      totalWaitedTime += 100
-    }
-    ss.sparkContext.isStopped
-  }
-
-  test("tracker should not affect execution result when timeout is not larger than 0") {
-    val nWorkers = numParallelism
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers)
-    val tracker = new SparkParallelismTracker(sc, 10000, nWorkers)
-    val disabledTracker = new SparkParallelismTracker(sc, 0, nWorkers)
-    assert(tracker.execute(rdd.sum()) == rdd.sum())
-    assert(disabledTracker.execute(rdd.sum()) == rdd.sum())
-  }
-
-  test("tracker should throw exception if parallelism is not sufficient") {
-    val nWorkers = numParallelism * 3
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers)
-    val tracker = new SparkParallelismTracker(sc, 1000, nWorkers)
-    intercept[IllegalStateException] {
-      tracker.execute {
-        rdd.map { i =>
-          // Test interruption
-          Thread.sleep(Long.MaxValue)
-          i
-        }.sum()
-      }
-    }
-  }
-
-  test("tracker should throw exception if parallelism is not sufficient with" +
-    " spark.task.cpus larger than 1") {
-    sc.conf.set("spark.task.cpus", "2")
-    val nWorkers = numParallelism
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers)
-    val tracker = new SparkParallelismTracker(sc, 1000, nWorkers)
-    intercept[IllegalStateException] {
-      tracker.execute {
-        rdd.map { i =>
-          // Test interruption
-          Thread.sleep(Long.MaxValue)
-          i
-        }.sum()
-      }
-    }
-  }
-
-  test("tracker should not kill SparkContext when killSparkContextOnWorkerFailure=false") {
-    val nWorkers = numParallelism
-    val tracker = new SparkParallelismTracker(sc, 0, nWorkers, false)
-    val rdd: RDD[Int] = sc.parallelize(1 to nWorkers, nWorkers)
-    try {
-      tracker.execute {
-        rdd.map { i =>
-          val partitionId = TaskContext.get().partitionId()
-          if (partitionId == 0) {
-            throw new RuntimeException("mocking task failing")
-          }
-          i
-        }.sum()
-      }
-    } catch {
-      case e: Exception => // catch the exception
-    } finally {
-      // wait 3s to check if SparkContext is killed
-      assert(waitAndCheckSparkShutdown(3000) == false)
-    }
-  }
-
-  test("tracker should cancel the correct job when killSparkContextOnWorkerFailure=false") {
-    val nWorkers = 2
-    val tracker = new SparkParallelismTracker(sc, 0, nWorkers, false)
-    val rdd: RDD[Int] = sc.parallelize(1 to 10, nWorkers)
-    val thread = new TestThread(sc)
-    thread.start()
-    try {
-      tracker.execute {
-        rdd.map { i =>
-          Thread.sleep(100)
-          val partitionId = TaskContext.get().partitionId()
-          if (partitionId == 0) {
-            throw new RuntimeException("mocking task failing")
-          }
-          i
-        }.sum()
-      }
-    } catch {
-      case e: Exception => // catch the exception
-    } finally {
-      thread.join(8000)
-      // wait 3s to check if SparkContext is killed
-      assert(waitAndCheckSparkShutdown(3000) == false)
-    }
-  }
-
-  private[this] class TestThread(sc: SparkContext) extends Thread {
-    override def run(): Unit = {
-      var sum: Double = 0.0f
-      try {
-        val rdd = sc.parallelize(1 to 4, 2)
-        sum = rdd.mapPartitions(iter => {
-          // sleep 2s to ensure task is alive when cancelling other jobs
-          Thread.sleep(2000)
-          iter
-        }).sum()
-      } finally {
-        // get the correct result
-        assert(sum.toInt == 10)
-      }
-    }
-  }
-}

jvm-packages/xgboost4j/pom.xml

@@ -6,10 +6,10 @@
     <parent>
         <groupId>ml.dmlc</groupId>
         <artifactId>xgboost-jvm_2.12</artifactId>
-        <version>1.6.0</version>
+        <version>1.6.2</version>
     </parent>
     <artifactId>xgboost4j_2.12</artifactId>
-    <version>1.6.0</version>
+    <version>1.6.2</version>
     <packaging>jar</packaging>
 
     <dependencies>

jvm-packages/xgboost4j/src/main/java/ml/dmlc/xgboost4j/java/NativeLibLoader.java

@@ -100,7 +100,7 @@ class NativeLibLoader {
         });
 
         return muslRelatedMemoryMappedFilename.isPresent();
-      } catch (IOException ignored) {
+      } catch (Exception ignored) {
         // ignored
       }
       return false;

plugin/CMakeLists.txt

@@ -2,19 +2,6 @@ if (PLUGIN_DENSE_PARSER)
   target_sources(objxgboost PRIVATE ${xgboost_SOURCE_DIR}/plugin/dense_parser/dense_libsvm.cc)
 endif (PLUGIN_DENSE_PARSER)
 
-if (PLUGIN_RMM)
-  find_path(RMM_INCLUDE "rmm" HINTS "$ENV{RMM_ROOT}/include")
-  if (NOT RMM_INCLUDE)
-    message(FATAL_ERROR "Could not locate RMM library")
-  endif ()
-
-  message(STATUS "RMM: RMM_LIBRARY set to ${RMM_LIBRARY}")
-  message(STATUS "RMM: RMM_INCLUDE set to ${RMM_INCLUDE}")
-
-  target_include_directories(objxgboost PUBLIC ${RMM_INCLUDE})
-  target_compile_definitions(objxgboost PUBLIC -DXGBOOST_USE_RMM=1)
-endif (PLUGIN_RMM)
-
 if (PLUGIN_UPDATER_ONEAPI)
   add_library(oneapi_plugin OBJECT
     ${xgboost_SOURCE_DIR}/plugin/updater_oneapi/regression_obj_oneapi.cc

python-package/MANIFEST.in

@@ -18,6 +18,8 @@ recursive-include xgboost/include *
 recursive-include xgboost/plugin *
 recursive-include xgboost/src *
 
+recursive-include xgboost/gputreeshap/GPUTreeShap *
+
 include xgboost/rabit/CMakeLists.txt
 recursive-include xgboost/rabit/include *
 recursive-include xgboost/rabit/src *

python-package/setup.py

@@ -61,6 +61,7 @@ def copy_tree(src_dir: str, target_dir: str) -> None:
     src = os.path.join(src_dir, 'src')
     inc = os.path.join(src_dir, 'include')
     dmlc_core = os.path.join(src_dir, 'dmlc-core')
+    gputreeshap = os.path.join(src_dir, "gputreeshap")
     rabit = os.path.join(src_dir, 'rabit')
     cmake = os.path.join(src_dir, 'cmake')
     plugin = os.path.join(src_dir, 'plugin')
@@ -68,6 +69,7 @@ def copy_tree(src_dir: str, target_dir: str) -> None:
     clean_copy_tree(src, os.path.join(target_dir, 'src'))
     clean_copy_tree(inc, os.path.join(target_dir, 'include'))
     clean_copy_tree(dmlc_core, os.path.join(target_dir, 'dmlc-core'))
+    clean_copy_tree(gputreeshap, os.path.join(target_dir, "gputreeshap"))
     clean_copy_tree(rabit, os.path.join(target_dir, 'rabit'))
     clean_copy_tree(cmake, os.path.join(target_dir, 'cmake'))
     clean_copy_tree(plugin, os.path.join(target_dir, 'plugin'))
@@ -97,7 +99,7 @@ class BuildExt(build_ext.build_ext):  # pylint: disable=too-many-ancestors
 
     logger = logging.getLogger('XGBoost build_ext')
 
-    # pylint: disable=too-many-arguments,no-self-use
+    # pylint: disable=too-many-arguments
     def build(
         self,
         src_dir: str,

python-package/xgboost/VERSION

@@ -1 +1 @@
-1.6.0
+1.6.2

python-package/xgboost/__init__.py

@@ -1,12 +1,16 @@
-# coding: utf-8
 """XGBoost: eXtreme Gradient Boosting library.
 
 Contributors: https://github.com/dmlc/xgboost/blob/master/CONTRIBUTORS.md
 """
 
-import os
+from .core import (
-
+    DMatrix,
-from .core import DMatrix, DeviceQuantileDMatrix, Booster, DataIter, build_info
+    DeviceQuantileDMatrix,
+    Booster,
+    DataIter,
+    build_info,
+    _py_version,
+)
 from .training import train, cv
 from . import rabit  # noqa
 from . import tracker  # noqa
@@ -21,9 +25,9 @@ try:
 except ImportError:
     pass
 
-VERSION_FILE = os.path.join(os.path.dirname(__file__), "VERSION")
+
-with open(VERSION_FILE, encoding="ascii") as f:
+__version__ = _py_version()
-    __version__ = f.read().strip()
+
 
 __all__ = [
     # core

python-package/xgboost/callback.py

@@ -1,5 +1,5 @@
 # coding: utf-8
-# pylint: disable=invalid-name, too-many-statements, no-self-use
+# pylint: disable=invalid-name, too-many-statements
 # pylint: disable=too-many-arguments
 """Callback library containing training routines.  See :doc:`Callback Functions
 </python/callbacks>` for a quick introduction.

python-package/xgboost/core.py

@@ -3,6 +3,7 @@
 """Core XGBoost Library."""
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
+import copy
 from typing import List, Optional, Any, Union, Dict, TypeVar
 from typing import Callable, Tuple, cast, Sequence, Type, Iterable
 import ctypes
@@ -136,6 +137,22 @@ def _get_log_callback_func() -> Callable:
     return c_callback(_log_callback)
 
 
+def _lib_version(lib: ctypes.CDLL) -> Tuple[int, int, int]:
+    """Get the XGBoost version from native shared object."""
+    major = ctypes.c_int()
+    minor = ctypes.c_int()
+    patch = ctypes.c_int()
+    lib.XGBoostVersion(ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch))
+    return major.value, minor.value, patch.value
+
+
+def _py_version() -> str:
+    """Get the XGBoost version from Python version file."""
+    VERSION_FILE = os.path.join(os.path.dirname(__file__), "VERSION")
+    with open(VERSION_FILE, encoding="ascii") as f:
+        return f.read().strip()
+
+
 def _load_lib() -> ctypes.CDLL:
     """Load xgboost Library."""
     lib_paths = find_lib_path()
@@ -143,7 +160,7 @@ def _load_lib() -> ctypes.CDLL:
         # This happens only when building document.
         return None  # type: ignore
     try:
-        pathBackup = os.environ['PATH'].split(os.pathsep)
+        pathBackup = os.environ["PATH"].split(os.pathsep)
     except KeyError:
         pathBackup = []
     lib_success = False
@@ -152,15 +169,16 @@ def _load_lib() -> ctypes.CDLL:
         try:
             # needed when the lib is linked with non-system-available
             # dependencies
-            os.environ['PATH'] = os.pathsep.join(
+            os.environ["PATH"] = os.pathsep.join(
-                pathBackup + [os.path.dirname(lib_path)])
+                pathBackup + [os.path.dirname(lib_path)]
+            )
             lib = ctypes.cdll.LoadLibrary(lib_path)
             lib_success = True
         except OSError as e:
             os_error_list.append(str(e))
             continue
         finally:
-            os.environ['PATH'] = os.pathsep.join(pathBackup)
+            os.environ["PATH"] = os.pathsep.join(pathBackup)
     if not lib_success:
         libname = os.path.basename(lib_paths[0])
         raise XGBoostError(
@@ -176,11 +194,36 @@ Likely causes:
   * You are running 32-bit Python on a 64-bit OS
 
 Error message(s): {os_error_list}
-""")
+"""
+        )
     lib.XGBGetLastError.restype = ctypes.c_char_p
     lib.callback = _get_log_callback_func()  # type: ignore
     if lib.XGBRegisterLogCallback(lib.callback) != 0:
         raise XGBoostError(lib.XGBGetLastError())
+
+    def parse(ver: str) -> Tuple[int, int, int]:
+        """Avoid dependency on packaging (PEP 440)."""
+        # 2.0.0-dev or 2.0.0
+        major, minor, patch = ver.split("-")[0].split(".")
+        return int(major), int(minor), int(patch)
+
+    libver = _lib_version(lib)
+    pyver = parse(_py_version())
+
+    # verify that we are loading the correct binary.
+    if pyver != libver:
+        pyver_str = ".".join((str(v) for v in pyver))
+        libver_str = ".".join((str(v) for v in libver))
+        msg = (
+            "Mismatched version between the Python package and the native shared "
+            f"""object.  Python package version: {pyver_str}. Shared object """
+            f"""version: {libver_str}. Shared object is loaded from: {lib.path}.
+Likely cause:
+  * XGBoost is first installed with anaconda then upgraded with pip. To fix it """
+            "please remove one of the installations."
+        )
+        raise ValueError(msg)
+
     return lib
 
 
@@ -1402,10 +1445,12 @@ class Booster:
         self.set_param(params_processed or {})
 
     def _transform_monotone_constrains(
-        self, value: Union[Dict[str, int], str]
+        self, value: Union[Dict[str, int], str, Tuple[int, ...]]
     ) -> Union[Tuple[int, ...], str]:
         if isinstance(value, str):
             return value
+        if isinstance(value, tuple):
+            return value
 
         constrained_features = set(value.keys())
         feature_names = self.feature_names or []
@@ -1577,7 +1622,7 @@ class Booster:
         booster: `Booster`
             a copied booster model
         """
-        return self.__copy__()
+        return copy.copy(self)
 
     def attr(self, key: str) -> Optional[str]:
         """Get attribute string from the Booster.
@@ -2309,15 +2354,15 @@ class Booster:
         ret = self.get_dump(fmap, with_stats, dump_format)
         if dump_format == 'json':
             fout_obj.write('[\n')
-            for i, _ in enumerate(ret):
+            for i, val in enumerate(ret):
-                fout_obj.write(ret[i])
+                fout_obj.write(val)
                 if i < len(ret) - 1:
                     fout_obj.write(",\n")
             fout_obj.write('\n]')
         else:
-            for i, _ in enumerate(ret):
+            for i, val in enumerate(ret):
                 fout_obj.write(f"booster[{i}]:\n")
-                fout_obj.write(ret[i])
+                fout_obj.write(val)
         if need_close:
             fout_obj.close()
 
@@ -2604,8 +2649,8 @@ class Booster:
         values = []
         # pylint: disable=consider-using-f-string
         regexp = re.compile(r"\[{0}<([\d.Ee+-]+)\]".format(feature))
-        for i, _ in enumerate(xgdump):
+        for i, val in enumerate(xgdump):
-            m = re.findall(regexp, xgdump[i])
+            m = re.findall(regexp, val)
             values.extend([float(x) for x in m])
 
         n_unique = len(np.unique(values))

python-package/xgboost/dask.py

@@ -177,9 +177,11 @@ def _try_start_tracker(
                 use_logger=False,
             )
         else:
-            assert isinstance(addrs[0], str) or addrs[0] is None
+            addr = addrs[0]
+            assert isinstance(addr, str) or addr is None
+            host_ip = get_host_ip(addr)
             rabit_context = RabitTracker(
-                host_ip=get_host_ip(addrs[0]), n_workers=n_workers, use_logger=False
+                host_ip=host_ip, n_workers=n_workers, use_logger=False, sortby="task"
             )
         env.update(rabit_context.worker_envs())
         rabit_context.start(n_workers)
@@ -229,8 +231,16 @@ class RabitContext:
     def __init__(self, args: List[bytes]) -> None:
         self.args = args
         worker = distributed.get_worker()
+        with distributed.worker_client() as client:
+            info = client.scheduler_info()
+            w = info["workers"][worker.address]
+            wid = w["id"]
+        # We use task ID for rank assignment which makes the RABIT rank consistent (but
+        # not the same as task ID is string and "10" is sorted before "2") with dask
+        # worker ID. This outsources the rank assignment to dask and prevents
+        # non-deterministic issue.
         self.args.append(
-            ("DMLC_TASK_ID=[xgboost.dask]:" + str(worker.address)).encode()
+            (f"DMLC_TASK_ID=[xgboost.dask-{wid}]:" + str(worker.address)).encode()
         )
 
     def __enter__(self) -> None:
@@ -870,6 +880,8 @@ async def _get_rabit_args(
     except Exception:  # pylint: disable=broad-except
         sched_addr = None
 
+    # make sure all workers are online so that we can obtain reliable scheduler_info
+    client.wait_for_workers(n_workers)
     env = await client.run_on_scheduler(
         _start_tracker, n_workers, sched_addr, user_addr
     )
@@ -1721,7 +1733,7 @@ class DaskScikitLearnBase(XGBModel):
     """Implementation of the Scikit-Learn API for XGBoost.""", ["estimators", "model"]
 )
 class DaskXGBRegressor(DaskScikitLearnBase, XGBRegressorBase):
-    # pylint: disable=missing-class-docstring
+    """dummy doc string to workaround pylint, replaced by the decorator."""
     async def _fit_async(
         self,
         X: _DaskCollection,

python-package/xgboost/data.py

@@ -2,7 +2,6 @@
 # pylint: disable=too-many-return-statements, import-error
 '''Data dispatching for DMatrix.'''
 import ctypes
-from distutils import version
 import json
 import warnings
 import os
@@ -506,7 +505,6 @@ def record_batch_data_iter(data_iter: Iterator) -> Callable:
     """Data iterator used to ingest Arrow columnar record batches. We are not using
     class DataIter because it is only intended for building Device DMatrix and external
     memory DMatrix.
-
     """
     from pyarrow.cffi import ffi
 
@@ -557,13 +555,7 @@ def _from_arrow(
     if enable_categorical:
         raise ValueError("categorical data in arrow is not supported yet.")
 
-    major, _, _ = version.StrictVersion(pa.__version__).version
+    rb_iter = iter(data.to_batches())
-    if major == 4:
-        rb_iter = iter(data.to_batches())
-    else:
-        # use_async=True to workaround pyarrow 6.0.1 hang,
-        # see Modin-3982 and ARROW-15362
-        rb_iter = iter(data.to_batches(use_async=True))
     it = record_batch_data_iter(rb_iter)
     next_callback = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p)(it)
     handle = ctypes.c_void_p()
@@ -714,9 +706,10 @@ def _is_cudf_ser(data):
     return isinstance(data, cudf.Series)
 
 
-def _is_cupy_array(data: Any) -> bool:
+def _is_cupy_array(data) -> bool:
-    return lazy_isinstance(data, "cupy.core.core", "ndarray") or lazy_isinstance(
+    return any(
-        data, "cupy._core.core", "ndarray"
+        lazy_isinstance(data, n, "ndarray")
+        for n in ("cupy.core.core", "cupy", "cupy._core.core")
     )
 
 

python-package/xgboost/tracker.py

@@ -32,15 +32,15 @@ class ExSocket:
             chunk = self.sock.recv(min(nbytes - nread, 1024))
             nread += len(chunk)
             res.append(chunk)
-        return b''.join(res)
+        return b"".join(res)
 
     def recvint(self) -> int:
         """Receive an integer of 32 bytes"""
-        return struct.unpack('@i', self.recvall(4))[0]
+        return struct.unpack("@i", self.recvall(4))[0]
 
     def sendint(self, value: int) -> None:
         """Send an integer of 32 bytes"""
-        self.sock.sendall(struct.pack('@i', value))
+        self.sock.sendall(struct.pack("@i", value))
 
     def sendstr(self, value: str) -> None:
         """Send a Python string"""
@@ -69,6 +69,7 @@ def get_family(addr: str) -> int:
 
 class WorkerEntry:
     """Hanlder to each worker."""
+
     def __init__(self, sock: socket.socket, s_addr: Tuple[str, int]):
         worker = ExSocket(sock)
         self.sock = worker
@@ -78,7 +79,7 @@ class WorkerEntry:
         worker.sendint(MAGIC_NUM)
         self.rank = worker.recvint()
         self.world_size = worker.recvint()
-        self.jobid = worker.recvstr()
+        self.task_id = worker.recvstr()
         self.cmd = worker.recvstr()
         self.wait_accept = 0
         self.port: Optional[int] = None
@@ -96,8 +97,8 @@ class WorkerEntry:
         """Get the rank of current entry."""
         if self.rank >= 0:
             return self.rank
-        if self.jobid != 'NULL' and self.jobid in job_map:
+        if self.task_id != "NULL" and self.task_id in job_map:
-            return job_map[self.jobid]
+            return job_map[self.task_id]
         return -1
 
     def assign_rank(
@@ -180,7 +181,12 @@ class RabitTracker:
     """
 
     def __init__(
-        self, host_ip: str, n_workers: int, port: int = 0, use_logger: bool = False
+        self,
+        host_ip: str,
+        n_workers: int,
+        port: int = 0,
+        use_logger: bool = False,
+        sortby: str = "host",
     ) -> None:
         """A Python implementation of RABIT tracker.
 
@@ -190,6 +196,13 @@ class RabitTracker:
             Use logging.info for tracker print command.  When set to False, Python print
             function is used instead.
 
+        sortby:
+            How to sort the workers for rank assignment. The default is host, but users
+            can set the `DMLC_TASK_ID` via RABIT initialization arguments and obtain
+            deterministic rank assignment. Available options are:
+              - host
+              - task
+
         """
         sock = socket.socket(get_family(host_ip), socket.SOCK_STREAM)
         sock.bind((host_ip, port))
@@ -200,6 +213,7 @@ class RabitTracker:
         self.thread: Optional[Thread] = None
         self.n_workers = n_workers
         self._use_logger = use_logger
+        self._sortby = sortby
         logging.info("start listen on %s:%d", host_ip, self.port)
 
     def __del__(self) -> None:
@@ -223,7 +237,7 @@ class RabitTracker:
         get environment variables for workers
         can be passed in as args or envs
         """
-        return {'DMLC_TRACKER_URI': self.host_ip, 'DMLC_TRACKER_PORT': self.port}
+        return {"DMLC_TRACKER_URI": self.host_ip, "DMLC_TRACKER_PORT": self.port}
 
     def _get_tree(self, n_workers: int) -> Tuple[_TreeMap, Dict[int, int]]:
         tree_map: _TreeMap = {}
@@ -296,8 +310,16 @@ class RabitTracker:
                 parent_map_[rmap[k]] = -1
         return tree_map_, parent_map_, ring_map_
 
+    def _sort_pending(self, pending: List[WorkerEntry]) -> List[WorkerEntry]:
+        if self._sortby == "host":
+            pending.sort(key=lambda s: s.host)
+        elif self._sortby == "task":
+            pending.sort(key=lambda s: s.task_id)
+        return pending
+
     def accept_workers(self, n_workers: int) -> None:
         """Wait for all workers to connect to the tracker."""
+
         # set of nodes that finishes the job
         shutdown: Dict[int, WorkerEntry] = {}
         # set of nodes that is waiting for connections
@@ -341,27 +363,32 @@ class RabitTracker:
                 assert todo_nodes
                 pending.append(s)
                 if len(pending) == len(todo_nodes):
-                    pending.sort(key=lambda x: x.host)
+                    pending = self._sort_pending(pending)
                     for s in pending:
                         rank = todo_nodes.pop(0)
-                        if s.jobid != 'NULL':
+                        if s.task_id != "NULL":
-                            job_map[s.jobid] = rank
+                            job_map[s.task_id] = rank
                         s.assign_rank(rank, wait_conn, tree_map, parent_map, ring_map)
                         if s.wait_accept > 0:
                             wait_conn[rank] = s
-                        logging.debug('Received %s signal from %s; assign rank %d',
+                        logging.debug(
-                                      s.cmd, s.host, s.rank)
+                            "Received %s signal from %s; assign rank %d",
+                            s.cmd,
+                            s.host,
+                            s.rank,
+                        )
                 if not todo_nodes:
-                    logging.info('@tracker All of %d nodes getting started', n_workers)
+                    logging.info("@tracker All of %d nodes getting started", n_workers)
             else:
                 s.assign_rank(rank, wait_conn, tree_map, parent_map, ring_map)
-                logging.debug('Received %s signal from %d', s.cmd, s.rank)
+                logging.debug("Received %s signal from %d", s.cmd, s.rank)
                 if s.wait_accept > 0:
                     wait_conn[rank] = s
-        logging.info('@tracker All nodes finishes job')
+        logging.info("@tracker All nodes finishes job")
 
     def start(self, n_workers: int) -> None:
         """Strat the tracker, it will wait for `n_workers` to connect."""
+
         def run() -> None:
             self.accept_workers(n_workers)
 

src/c_api/c_api.cc

@@ -193,9 +193,7 @@ XGB_DLL int XGBGetGlobalConfig(const char** json_str) {
   API_END();
 }
 
-XGB_DLL int XGDMatrixCreateFromFile(const char *fname,
+XGB_DLL int XGDMatrixCreateFromFile(const char *fname, int silent, DMatrixHandle *out) {
-                                    int silent,
-                                    DMatrixHandle *out) {
   API_BEGIN();
   bool load_row_split = false;
   if (rabit::IsDistributed()) {

src/c_api/c_api_utils.h

@@ -115,7 +115,7 @@ inline void CalcPredictShape(bool strict_shape, PredictionType type, size_t rows
   }
   }
   CHECK_EQ(
-      std::accumulate(shape.cbegin(), shape.cend(), 1, std::multiplies<>{}),
+      std::accumulate(shape.cbegin(), shape.cend(), static_cast<bst_ulong>(1), std::multiplies<>{}),
       chunksize * rows);
 }
 

src/common/categorical.h

@@ -12,7 +12,6 @@
 #include "xgboost/data.h"
 #include "xgboost/parameter.h"
 #include "xgboost/span.h"
-#include "xgboost/task.h"
 
 namespace xgboost {
 namespace common {
@@ -75,15 +74,20 @@ inline void InvalidCategory() {
   // values to be less than this last representable value.
   auto str = std::to_string(OutOfRangeCat());
   LOG(FATAL) << "Invalid categorical value detected.  Categorical value should be non-negative, "
-                "less than total umber of categories in training data and less than " +
+                "less than total number of categories in training data and less than " +
                     str;
 }
 
+inline void CheckMaxCat(float max_cat, size_t n_categories) {
+  CHECK_GE(max_cat + 1, n_categories)
+      << "Maximum cateogry should not be lesser than the total number of categories.";
+}
+
 /*!
  * \brief Whether should we use onehot encoding for categorical data.
  */
-XGBOOST_DEVICE inline bool UseOneHot(uint32_t n_cats, uint32_t max_cat_to_onehot, ObjInfo task) {
+XGBOOST_DEVICE inline bool UseOneHot(uint32_t n_cats, uint32_t max_cat_to_onehot) {
-  bool use_one_hot = n_cats < max_cat_to_onehot || task.UseOneHot();
+  bool use_one_hot = n_cats < max_cat_to_onehot;
   return use_one_hot;
 }
 

src/common/common.h

@@ -164,6 +164,74 @@ class Range {
   Iterator end_;
 };
 
+/**
+ * \brief Transform iterator that takes an index and calls transform operator.
+ *
+ *   This is CPU-only right now as taking host device function as operator complicates the
+ *   code.  For device side one can use `thrust::transform_iterator` instead.
+ */
+template <typename Fn>
+class IndexTransformIter {
+  size_t iter_{0};
+  Fn fn_;
+
+ public:
+  using iterator_category = std::random_access_iterator_tag;  // NOLINT
+  using value_type = std::result_of_t<Fn(size_t)>;            // NOLINT
+  using difference_type = detail::ptrdiff_t;                  // NOLINT
+  using reference = std::add_lvalue_reference_t<value_type>;  // NOLINT
+  using pointer = std::add_pointer_t<value_type>;             // NOLINT
+
+ public:
+  /**
+   * \param op Transform operator, takes a size_t index as input.
+   */
+  explicit IndexTransformIter(Fn &&op) : fn_{op} {}
+  IndexTransformIter(IndexTransformIter const &) = default;
+  IndexTransformIter& operator=(IndexTransformIter&&) = default;
+  IndexTransformIter& operator=(IndexTransformIter const& that) {
+    iter_ = that.iter_;
+    return *this;
+  }
+
+  value_type operator*() const { return fn_(iter_); }
+
+  auto operator-(IndexTransformIter const &that) const { return iter_ - that.iter_; }
+  bool operator==(IndexTransformIter const &that) const { return iter_ == that.iter_; }
+  bool operator!=(IndexTransformIter const &that) const { return !(*this == that); }
+
+  IndexTransformIter &operator++() {
+    iter_++;
+    return *this;
+  }
+  IndexTransformIter operator++(int) {
+    auto ret = *this;
+    ++(*this);
+    return ret;
+  }
+  IndexTransformIter &operator+=(difference_type n) {
+    iter_ += n;
+    return *this;
+  }
+  IndexTransformIter &operator-=(difference_type n) {
+    (*this) += -n;
+    return *this;
+  }
+  IndexTransformIter operator+(difference_type n) const {
+    auto ret = *this;
+    return ret += n;
+  }
+  IndexTransformIter operator-(difference_type n) const {
+    auto ret = *this;
+    return ret -= n;
+  }
+};
+
+template <typename Fn>
+auto MakeIndexTransformIter(Fn&& fn) {
+  return IndexTransformIter<Fn>(std::forward<Fn>(fn));
+}
+
 int AllVisibleGPUs();
 
 inline void AssertGPUSupport() {

src/common/device_helpers.cu

@@ -38,6 +38,9 @@ void AllReducer::Init(int _device_ordinal) {
 
   int32_t const rank = rabit::GetRank();
   int32_t const world = rabit::GetWorldSize();
+  if (world == 1) {
+    return;
+  }
 
   std::vector<uint64_t> uuids(world * kUuidLength, 0);
   auto s_uuid = xgboost::common::Span<uint64_t>{uuids.data(), uuids.size()};

src/common/device_helpers.cuh

@@ -775,13 +775,16 @@ class AllReducer {
    */
 
   void AllReduceSum(const double *sendbuff, double *recvbuff, int count) {
+    if (rabit::GetWorldSize() == 1) {
+      return;
+    }
 #ifdef XGBOOST_USE_NCCL
     CHECK(initialised_);
     dh::safe_cuda(cudaSetDevice(device_ordinal_));
     dh::safe_nccl(ncclAllReduce(sendbuff, recvbuff, count, ncclDouble, ncclSum, comm_, stream_));
     allreduce_bytes_ += count * sizeof(double);
     allreduce_calls_ += 1;
-#endif
+#endif  // XGBOOST_USE_NCCL
   }
 
   /**
@@ -796,9 +799,12 @@ class AllReducer {
 
   void AllGather(uint32_t const* data, size_t length,
                  dh::caching_device_vector<uint32_t>* recvbuf) {
+    size_t world = rabit::GetWorldSize();
+    if (world == 1) {
+      return;
+    }
 #ifdef XGBOOST_USE_NCCL
     CHECK(initialised_);
-    size_t world = rabit::GetWorldSize();
     recvbuf->resize(length * world);
     safe_nccl(ncclAllGather(data, recvbuf->data().get(), length, ncclUint32,
                             comm_, stream_));
@@ -813,9 +819,11 @@ class AllReducer {
    * \param recvbuff                The recvbuff.
    * \param count                   Number of elements.
    */
-
   void AllReduceSum(const float *sendbuff, float *recvbuff, int count) {
 #ifdef XGBOOST_USE_NCCL
+    if (rabit::GetWorldSize() == 1) {
+      return;
+    }
     CHECK(initialised_);
     dh::safe_cuda(cudaSetDevice(device_ordinal_));
     dh::safe_nccl(ncclAllReduce(sendbuff, recvbuff, count, ncclFloat, ncclSum, comm_, stream_));
@@ -836,6 +844,9 @@ class AllReducer {
 
   void AllReduceSum(const int64_t *sendbuff, int64_t *recvbuff, int count) {
 #ifdef XGBOOST_USE_NCCL
+    if (rabit::GetWorldSize() == 1) {
+      return;
+    }
     CHECK(initialised_);
 
     dh::safe_cuda(cudaSetDevice(device_ordinal_));
@@ -845,6 +856,9 @@ class AllReducer {
 
   void AllReduceSum(const uint32_t *sendbuff, uint32_t *recvbuff, int count) {
 #ifdef XGBOOST_USE_NCCL
+    if (rabit::GetWorldSize() == 1) {
+      return;
+    }
     CHECK(initialised_);
 
     dh::safe_cuda(cudaSetDevice(device_ordinal_));
@@ -853,6 +867,9 @@ class AllReducer {
   }
 
   void AllReduceSum(const uint64_t *sendbuff, uint64_t *recvbuff, int count) {
+    if (rabit::GetWorldSize() == 1) {
+      return;
+    }
 #ifdef XGBOOST_USE_NCCL
     CHECK(initialised_);
 
@@ -867,12 +884,15 @@ class AllReducer {
             std::enable_if_t<std::is_same<size_t, T>::value &&
                              !std::is_same<size_t, unsigned long long>::value>  // NOLINT
                 * = nullptr>
-  void AllReduceSum(const T *sendbuff, T *recvbuff, int count) { // NOLINT
+  void AllReduceSum(const T *sendbuff, T *recvbuff, int count) {  // NOLINT
 #ifdef XGBOOST_USE_NCCL
+    if (rabit::GetWorldSize() == 1) {
+      return;
+    }
     CHECK(initialised_);
 
     dh::safe_cuda(cudaSetDevice(device_ordinal_));
-    static_assert(sizeof(unsigned long long) == sizeof(uint64_t), ""); // NOLINT
+    static_assert(sizeof(unsigned long long) == sizeof(uint64_t), "");  // NOLINT
     dh::safe_nccl(ncclAllReduce(sendbuff, recvbuff, count, ncclUint64, ncclSum, comm_, stream_));
 #endif
   }

src/common/hist_util.cuh

@@ -184,8 +184,6 @@ void ProcessWeightedSlidingWindow(Batch batch, MetaInfo const& info,
   dh::safe_cuda(cudaSetDevice(device));
   info.weights_.SetDevice(device);
   auto weights = info.weights_.ConstDeviceSpan();
-  dh::caching_device_vector<bst_group_t> group_ptr(info.group_ptr_);
-  auto d_group_ptr = dh::ToSpan(group_ptr);
 
   auto batch_iter = dh::MakeTransformIterator<data::COOTuple>(
     thrust::make_counting_iterator(0llu),
@@ -205,9 +203,13 @@ void ProcessWeightedSlidingWindow(Batch batch, MetaInfo const& info,
   auto d_temp_weights = dh::ToSpan(temp_weights);
 
   if (is_ranking) {
+    if (!weights.empty()) {
+      CHECK_EQ(weights.size(), info.group_ptr_.size() - 1);
+    }
+    dh::caching_device_vector<bst_group_t> group_ptr(info.group_ptr_);
+    auto d_group_ptr = dh::ToSpan(group_ptr);
     auto const weight_iter = dh::MakeTransformIterator<float>(
-        thrust::make_constant_iterator(0lu),
+        thrust::make_counting_iterator(0lu), [=] __device__(size_t idx) -> float {
-        [=]__device__(size_t idx) -> float {
           auto ridx = batch.GetElement(idx).row_idx;
           bst_group_t group_idx = dh::SegmentId(d_group_ptr, ridx);
           return weights[group_idx];
@@ -272,7 +274,7 @@ void AdapterDeviceSketch(Batch batch, int num_bins,
   size_t num_cols = batch.NumCols();
   size_t num_cuts_per_feature = detail::RequiredSampleCutsPerColumn(num_bins, num_rows);
   int32_t device = sketch_container->DeviceIdx();
-  bool weighted = info.weights_.Size() != 0;
+  bool weighted = !info.weights_.Empty();
 
   if (weighted) {
     sketch_batch_num_elements = detail::SketchBatchNumElements(

src/common/quantile.cc

@@ -122,27 +122,6 @@ std::vector<float> MergeWeights(MetaInfo const &info, Span<float const> hessian,
   }
   return results;
 }
-
-std::vector<float> UnrollGroupWeights(MetaInfo const &info) {
-  std::vector<float> const &group_weights = info.weights_.HostVector();
-  if (group_weights.empty()) {
-    return group_weights;
-  }
-
-  size_t n_samples = info.num_row_;
-  auto const &group_ptr = info.group_ptr_;
-  std::vector<float> results(n_samples);
-  CHECK_GE(group_ptr.size(), 2);
-  CHECK_EQ(group_ptr.back(), n_samples);
-  size_t cur_group = 0;
-  for (size_t i = 0; i < n_samples; ++i) {
-    results[i] = group_weights[cur_group];
-    if (i == group_ptr[cur_group + 1]) {
-      cur_group++;
-    }
-  }
-  return results;
-}
 }  // anonymous namespace
 
 template <typename WQSketch>
@@ -156,12 +135,10 @@ void SketchContainerImpl<WQSketch>::PushRowPage(SparsePage const &page, MetaInfo
 
   // glue these conditions using ternary operator to avoid making data copies.
   auto const &weights =
-      hessian.empty()
+      hessian.empty() ? (use_group_ind_ ? detail::UnrollGroupWeights(info)  // use group weight
-          ? (use_group_ind_ ? UnrollGroupWeights(info)     // use group weight
+                                        : info.weights_.HostVector())       // use sample weight
-                            : info.weights_.HostVector())  // use sample weight
+                      : MergeWeights(info, hessian, use_group_ind_,
-          : MergeWeights(
+                                     n_threads_);  // use hessian merged with group/sample weights
-                info, hessian, use_group_ind_,
-                n_threads_);  // use hessian merged with group/sample weights
   if (!weights.empty()) {
     CHECK_EQ(weights.size(), info.num_row_);
   }
@@ -468,11 +445,17 @@ void AddCutPoint(typename SketchType::SummaryContainer const &summary, int max_b
   }
 }
 
-void AddCategories(std::set<float> const &categories, HistogramCuts *cuts) {
+auto AddCategories(std::set<float> const &categories, HistogramCuts *cuts) {
-  auto &cut_values = cuts->cut_values_.HostVector();
+  if (std::any_of(categories.cbegin(), categories.cend(), InvalidCat)) {
-  for (auto const &v : categories) {
+    InvalidCategory();
-    cut_values.push_back(AsCat(v));
   }
+  auto &cut_values = cuts->cut_values_.HostVector();
+  auto max_cat = *std::max_element(categories.cbegin(), categories.cend());
+  CheckMaxCat(max_cat, categories.size());
+  for (bst_cat_t i = 0; i <= AsCat(max_cat); ++i) {
+    cut_values.push_back(i);
+  }
+  return max_cat;
 }
 
 template <typename WQSketch>
@@ -505,11 +488,12 @@ void SketchContainerImpl<WQSketch>::MakeCuts(HistogramCuts* cuts) {
     }
   });
 
+  float max_cat{-1.f};
   for (size_t fid = 0; fid < reduced.size(); ++fid) {
     size_t max_num_bins = std::min(num_cuts[fid], max_bins_);
     typename WQSketch::SummaryContainer const& a = final_summaries[fid];
     if (IsCat(feature_types_, fid)) {
-      AddCategories(categories_.at(fid), cuts);
+      max_cat = std::max(max_cat, AddCategories(categories_.at(fid), cuts));
     } else {
       AddCutPoint<WQSketch>(a, max_num_bins, cuts);
       // push a value that is greater than anything
@@ -527,30 +511,7 @@ void SketchContainerImpl<WQSketch>::MakeCuts(HistogramCuts* cuts) {
     cuts->cut_ptrs_.HostVector().push_back(cut_size);
   }
 
-  if (has_categorical_) {
+  cuts->SetCategorical(this->has_categorical_, max_cat);
-    for (auto const &feat : categories_) {
-      if (std::any_of(feat.cbegin(), feat.cend(), InvalidCat)) {
-        InvalidCategory();
-      }
-    }
-    auto const &ptrs = cuts->Ptrs();
-    auto const &vals = cuts->Values();
-
-    float max_cat{-std::numeric_limits<float>::infinity()};
-    for (size_t i = 1; i < ptrs.size(); ++i) {
-      if (IsCat(feature_types_, i - 1)) {
-        auto beg = ptrs[i - 1];
-        auto end = ptrs[i];
-        auto feat = Span<float const>{vals}.subspan(beg, end - beg);
-        auto max_elem = *std::max_element(feat.cbegin(), feat.cend());
-        if (max_elem > max_cat) {
-          max_cat = max_elem;
-        }
-      }
-    }
-    cuts->SetCategorical(true, max_cat);
-  }
-
   monitor_.Stop(__func__);
 }
 
@@ -579,8 +540,8 @@ void SortedSketchContainer::PushColPage(SparsePage const &page, MetaInfo const &
   monitor_.Start(__func__);
   // glue these conditions using ternary operator to avoid making data copies.
   auto const &weights =
-      hessian.empty() ? (use_group_ind_ ? UnrollGroupWeights(info)     // use group weight
+      hessian.empty() ? (use_group_ind_ ? detail::UnrollGroupWeights(info)  // use group weight
-                                        : info.weights_.HostVector())  // use sample weight
+                                        : info.weights_.HostVector())       // use sample weight
                       : MergeWeights(info, hessian, use_group_ind_,
                                      n_threads_);  // use hessian merged with group/sample weights
   CHECK_EQ(weights.size(), info.num_row_);

src/common/quantile.cu

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2020 by XGBoost Contributors
+ * Copyright 2020-2022 by XGBoost Contributors
  */
 #include <thrust/binary_search.h>
 #include <thrust/execution_policy.h>
@@ -583,13 +583,13 @@ void SketchContainer::AllReduce() {
 
 namespace {
 struct InvalidCatOp {
-  Span<float const> values;
+  Span<SketchEntry const> values;
-  Span<uint32_t const> ptrs;
+  Span<size_t const> ptrs;
   Span<FeatureType const> ft;
 
   XGBOOST_DEVICE bool operator()(size_t i) const {
     auto fidx = dh::SegmentId(ptrs, i);
-    return IsCat(ft, fidx) && InvalidCat(values[i]);
+    return IsCat(ft, fidx) && InvalidCat(values[i].value);
   }
 };
 }  // anonymous namespace
@@ -611,7 +611,7 @@ void SketchContainer::MakeCuts(HistogramCuts* p_cuts) {
 
   p_cuts->min_vals_.SetDevice(device_);
   auto d_min_values = p_cuts->min_vals_.DeviceSpan();
-  auto in_cut_values = dh::ToSpan(this->Current());
+  auto const in_cut_values = dh::ToSpan(this->Current());
 
   // Set up output ptr
   p_cuts->cut_ptrs_.SetDevice(device_);
@@ -619,26 +619,70 @@ void SketchContainer::MakeCuts(HistogramCuts* p_cuts) {
   h_out_columns_ptr.clear();
   h_out_columns_ptr.push_back(0);
   auto const& h_feature_types = this->feature_types_.ConstHostSpan();
-  for (bst_feature_t i = 0; i < num_columns_; ++i) {
+
-    size_t column_size = std::max(static_cast<size_t>(1ul),
+  auto d_ft = feature_types_.ConstDeviceSpan();
-                                  this->Column(i).size());
+
-    if (IsCat(h_feature_types, i)) {
+  std::vector<SketchEntry> max_values;
-      h_out_columns_ptr.push_back(static_cast<size_t>(column_size));
+  float max_cat{-1.f};
-    } else {
+  if (has_categorical_) {
-      h_out_columns_ptr.push_back(std::min(static_cast<size_t>(column_size),
+    dh::XGBCachingDeviceAllocator<char> alloc;
-                                           static_cast<size_t>(num_bins_)));
+    auto key_it = dh::MakeTransformIterator<bst_feature_t>(
+        thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) -> bst_feature_t {
+          return dh::SegmentId(d_in_columns_ptr, i);
+        });
+    auto invalid_op = InvalidCatOp{in_cut_values, d_in_columns_ptr, d_ft};
+    auto val_it = dh::MakeTransformIterator<SketchEntry>(
+        thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) {
+          auto fidx = dh::SegmentId(d_in_columns_ptr, i);
+          auto v = in_cut_values[i];
+          if (IsCat(d_ft, fidx)) {
+            if (invalid_op(i)) {
+              // use inf to indicate invalid value, this way we can keep it as in
+              // indicator in the reduce operation as it's always the greatest value.
+              v.value = std::numeric_limits<float>::infinity();
+            }
+          }
+          return v;
+        });
+    CHECK_EQ(num_columns_, d_in_columns_ptr.size() - 1);
+    max_values.resize(d_in_columns_ptr.size() - 1);
+    dh::caching_device_vector<SketchEntry> d_max_values(d_in_columns_ptr.size() - 1);
+    thrust::reduce_by_key(thrust::cuda::par(alloc), key_it, key_it + in_cut_values.size(), val_it,
+                          thrust::make_discard_iterator(), d_max_values.begin(),
+                          thrust::equal_to<bst_feature_t>{},
+                          [] __device__(auto l, auto r) { return l.value > r.value ? l : r; });
+    dh::CopyDeviceSpanToVector(&max_values, dh::ToSpan(d_max_values));
+    auto max_it = common::MakeIndexTransformIter([&](auto i) {
+      if (IsCat(h_feature_types, i)) {
+        return max_values[i].value;
+      }
+      return -1.f;
+    });
+    max_cat = *std::max_element(max_it, max_it + max_values.size());
+    if (std::isinf(max_cat)) {
+      InvalidCategory();
     }
   }
-  std::partial_sum(h_out_columns_ptr.begin(), h_out_columns_ptr.end(),
-                   h_out_columns_ptr.begin());
-  auto d_out_columns_ptr = p_cuts->cut_ptrs_.ConstDeviceSpan();
 
   // Set up output cuts
+  for (bst_feature_t i = 0; i < num_columns_; ++i) {
+    size_t column_size = std::max(static_cast<size_t>(1ul), this->Column(i).size());
+    if (IsCat(h_feature_types, i)) {
+      // column_size is the number of unique values in that feature.
+      CheckMaxCat(max_values[i].value, column_size);
+      h_out_columns_ptr.push_back(max_values[i].value + 1);  // includes both max_cat and 0.
+    } else {
+      h_out_columns_ptr.push_back(
+          std::min(static_cast<size_t>(column_size), static_cast<size_t>(num_bins_)));
+    }
+  }
+  std::partial_sum(h_out_columns_ptr.begin(), h_out_columns_ptr.end(), h_out_columns_ptr.begin());
+  auto d_out_columns_ptr = p_cuts->cut_ptrs_.ConstDeviceSpan();
+
   size_t total_bins = h_out_columns_ptr.back();
   p_cuts->cut_values_.SetDevice(device_);
   p_cuts->cut_values_.Resize(total_bins);
   auto out_cut_values = p_cuts->cut_values_.DeviceSpan();
-  auto d_ft = feature_types_.ConstDeviceSpan();
 
   dh::LaunchN(total_bins, [=] __device__(size_t idx) {
     auto column_id = dh::SegmentId(d_out_columns_ptr, idx);
@@ -667,8 +711,7 @@ void SketchContainer::MakeCuts(HistogramCuts* p_cuts) {
     }
 
     if (IsCat(d_ft, column_id)) {
-      assert(out_column.size() == in_column.size());
+      out_column[idx] = idx;
-      out_column[idx] = in_column[idx].value;
       return;
     }
 
@@ -684,36 +727,7 @@ void SketchContainer::MakeCuts(HistogramCuts* p_cuts) {
     out_column[idx] = in_column[idx+1].value;
   });
 
-  float max_cat{-1.0f};
-  if (has_categorical_) {
-    auto invalid_op = InvalidCatOp{out_cut_values, d_out_columns_ptr, d_ft};
-    auto it = dh::MakeTransformIterator<thrust::pair<bool, float>>(
-        thrust::make_counting_iterator(0ul), [=] XGBOOST_DEVICE(size_t i) {
-          auto fidx = dh::SegmentId(d_out_columns_ptr, i);
-          if (IsCat(d_ft, fidx)) {
-            auto invalid = invalid_op(i);
-            auto v = out_cut_values[i];
-            return thrust::make_pair(invalid, v);
-          }
-          return thrust::make_pair(false, std::numeric_limits<float>::min());
-        });
-
-    bool invalid{false};
-    dh::XGBCachingDeviceAllocator<char> alloc;
-    thrust::tie(invalid, max_cat) =
-        thrust::reduce(thrust::cuda::par(alloc), it, it + out_cut_values.size(),
-                       thrust::make_pair(false, std::numeric_limits<float>::min()),
-                       [=] XGBOOST_DEVICE(thrust::pair<bool, bst_cat_t> const &l,
-                                          thrust::pair<bool, bst_cat_t> const &r) {
-                         return thrust::make_pair(l.first || r.first, std::max(l.second, r.second));
-                       });
-    if (invalid) {
-      InvalidCategory();
-    }
-  }
-
   p_cuts->SetCategorical(this->has_categorical_, max_cat);
-
   timer_.Stop(__func__);
 }
 }  // namespace common

src/common/quantile.h

@@ -697,6 +697,29 @@ class WXQuantileSketch :
       public QuantileSketchTemplate<DType, RType, WXQSummary<DType, RType> > {
 };
 
+namespace detail {
+inline std::vector<float> UnrollGroupWeights(MetaInfo const &info) {
+  std::vector<float> const &group_weights = info.weights_.HostVector();
+  if (group_weights.empty()) {
+    return group_weights;
+  }
+
+  size_t n_samples = info.num_row_;
+  auto const &group_ptr = info.group_ptr_;
+  std::vector<float> results(n_samples);
+  CHECK_GE(group_ptr.size(), 2);
+  CHECK_EQ(group_ptr.back(), n_samples);
+  size_t cur_group = 0;
+  for (size_t i = 0; i < n_samples; ++i) {
+    results[i] = group_weights[cur_group];
+    if (i == group_ptr[cur_group + 1]) {
+      cur_group++;
+    }
+  }
+  return results;
+}
+}  // namespace detail
+
 class HistogramCuts;
 
 /*!

src/data/data.cc

@@ -378,35 +378,6 @@ MetaInfo MetaInfo::Slice(common::Span<int32_t const> ridxs) const {
   return out;
 }
 
-// try to load group information from file, if exists
-inline bool MetaTryLoadGroup(const std::string& fname,
-                             std::vector<unsigned>* group) {
-  std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(fname.c_str(), "r", true));
-  if (fi == nullptr) return false;
-  dmlc::istream is(fi.get());
-  group->clear();
-  group->push_back(0);
-  unsigned nline = 0;
-  while (is >> nline) {
-    group->push_back(group->back() + nline);
-  }
-  return true;
-}
-
-// try to load weight information from file, if exists
-inline bool MetaTryLoadFloatInfo(const std::string& fname,
-                                 std::vector<bst_float>* data) {
-  std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(fname.c_str(), "r", true));
-  if (fi == nullptr) return false;
-  dmlc::istream is(fi.get());
-  data->clear();
-  bst_float value;
-  while (is >> value) {
-    data->push_back(value);
-  }
-  return true;
-}
-
 namespace {
 template <int32_t D, typename T>
 void CopyTensorInfoImpl(Context const& ctx, Json arr_interface, linalg::Tensor<T, D>* p_out) {
@@ -811,9 +782,7 @@ DMatrix *TryLoadBinary(std::string fname, bool silent) {
   return nullptr;
 }
 
-DMatrix* DMatrix::Load(const std::string& uri,
+DMatrix* DMatrix::Load(const std::string& uri, bool silent, bool load_row_split,
-                       bool silent,
-                       bool load_row_split,
                        const std::string& file_format) {
   std::string fname, cache_file;
   size_t dlm_pos = uri.find('#');
@@ -846,50 +815,47 @@ DMatrix* DMatrix::Load(const std::string& uri,
   } else {
     fname = uri;
   }
+
+  // legacy handling of binary data loading
+  if (file_format == "auto") {
+    DMatrix* loaded = TryLoadBinary(fname, silent);
+    if (loaded) {
+      return loaded;
+    }
+  }
+
   int partid = 0, npart = 1;
   if (load_row_split) {
     partid = rabit::GetRank();
     npart = rabit::GetWorldSize();
   } else {
     // test option to load in part
-    npart = dmlc::GetEnv("XGBOOST_TEST_NPART", 1);
+    npart = 1;
   }
 
   if (npart != 1) {
-    LOG(CONSOLE) << "Load part of data " << partid
+    LOG(CONSOLE) << "Load part of data " << partid << " of " << npart << " parts";
-                 << " of " << npart << " parts";
-  }
-
-  // legacy handling of binary data loading
-  if (file_format == "auto" && npart == 1) {
-    DMatrix *loaded = TryLoadBinary(fname, silent);
-    if (loaded) {
-      return loaded;
-    }
   }
 
   DMatrix* dmat {nullptr};
   try {
     if (cache_file.empty()) {
       std::unique_ptr<dmlc::Parser<uint32_t>> parser(
-          dmlc::Parser<uint32_t>::Create(fname.c_str(), partid, npart,
+          dmlc::Parser<uint32_t>::Create(fname.c_str(), partid, npart, file_format.c_str()));
-                                         file_format.c_str()));
       data::FileAdapter adapter(parser.get());
-      dmat = DMatrix::Create(&adapter, std::numeric_limits<float>::quiet_NaN(),
+      dmat = DMatrix::Create(&adapter, std::numeric_limits<float>::quiet_NaN(), 1, cache_file);
-                             1, cache_file);
     } else {
       data::FileIterator iter{fname, static_cast<uint32_t>(partid), static_cast<uint32_t>(npart),
                               file_format};
-      dmat = new data::SparsePageDMatrix{
+      dmat = new data::SparsePageDMatrix{&iter,
-          &iter,
+                                         iter.Proxy(),
-          iter.Proxy(),
+                                         data::fileiter::Reset,
-          data::fileiter::Reset,
+                                         data::fileiter::Next,
-          data::fileiter::Next,
+                                         std::numeric_limits<float>::quiet_NaN(),
-          std::numeric_limits<float>::quiet_NaN(),
+                                         1,
-          1,
+                                         cache_file};
-          cache_file};
     }
-  } catch (dmlc::Error &e) {
+  } catch (dmlc::Error& e) {
     std::vector<std::string> splited = common::Split(fname, '#');
     std::vector<std::string> args = common::Split(splited.front(), '?');
     std::string format {file_format};
@@ -917,24 +883,6 @@ DMatrix* DMatrix::Load(const std::string& uri,
    * partitioned data will fail the train/val validation check
    * since partitioned data not knowing the real number of features. */
   rabit::Allreduce<rabit::op::Max>(&dmat->Info().num_col_, 1);
-  // backward compatiblity code.
-  if (!load_row_split) {
-    MetaInfo& info = dmat->Info();
-    if (MetaTryLoadGroup(fname + ".group", &info.group_ptr_) && !silent) {
-      LOG(CONSOLE) << info.group_ptr_.size() - 1
-                   << " groups are loaded from " << fname << ".group";
-    }
-    if (MetaTryLoadFloatInfo(fname + ".base_margin", &info.base_margin_.Data()->HostVector()) &&
-        !silent) {
-      LOG(CONSOLE) << info.base_margin_.Size() << " base_margin are loaded from " << fname
-                   << ".base_margin";
-    }
-    if (MetaTryLoadFloatInfo
-        (fname + ".weight", &info.weights_.HostVector()) && !silent) {
-      LOG(CONSOLE) << info.weights_.Size()
-                   << " weights are loaded from " << fname << ".weight";
-    }
-  }
   return dmat;
 }
 template <typename DataIterHandle, typename DMatrixHandle,

src/data/data.cu

@@ -130,12 +130,12 @@ void MetaInfo::SetInfoFromCUDA(Context const&, StringView key, Json array) {
   }
   // uint info
   if (key == "group") {
-    auto array_interface{ArrayInterface<1>(array)};
+    ArrayInterface<1> array_interface{array};
     CopyGroupInfoImpl(array_interface, &group_ptr_);
     data::ValidateQueryGroup(group_ptr_);
     return;
   } else if (key == "qid") {
-    auto array_interface{ArrayInterface<1>(array)};
+    ArrayInterface<1> array_interface{array};
     CopyQidImpl(array_interface, &group_ptr_);
     data::ValidateQueryGroup(group_ptr_);
     return;

src/metric/metric_common.h

@@ -1,5 +1,5 @@
 /*!
- * Copyright 2018-2020 by Contributors
+ * Copyright 2018-2022 by Contributors
  * \file metric_common.h
  */
 #ifndef XGBOOST_METRIC_METRIC_COMMON_H_
@@ -9,6 +9,7 @@
 #include <string>
 
 #include "../common/common.h"
+#include "xgboost/metric.h"
 
 namespace xgboost {
 

src/metric/rank_metric.cu

@@ -27,7 +27,7 @@ DMLC_REGISTRY_FILE_TAG(rank_metric_gpu);
 
 /*! \brief Evaluate rank list on GPU */
 template <typename EvalMetricT>
-struct EvalRankGpu : public Metric, public EvalRankConfig {
+struct EvalRankGpu : public GPUMetric, public EvalRankConfig {
  public:
   double Eval(const HostDeviceVector<bst_float> &preds, const MetaInfo &info,
               bool distributed) override {

src/tree/gpu_hist/evaluate_splits.cu

@@ -199,13 +199,11 @@ __device__ void EvaluateFeature(
 }
 
 template <int BLOCK_THREADS, typename GradientSumT>
-__global__ void EvaluateSplitsKernel(
+__global__ void EvaluateSplitsKernel(EvaluateSplitInputs<GradientSumT> left,
-    EvaluateSplitInputs<GradientSumT> left,
+                                     EvaluateSplitInputs<GradientSumT> right,
-    EvaluateSplitInputs<GradientSumT> right,
+                                     common::Span<bst_feature_t> sorted_idx,
-    ObjInfo task,
+                                     TreeEvaluator::SplitEvaluator<GPUTrainingParam> evaluator,
-    common::Span<bst_feature_t> sorted_idx,
+                                     common::Span<DeviceSplitCandidate> out_candidates) {
-    TreeEvaluator::SplitEvaluator<GPUTrainingParam> evaluator,
-    common::Span<DeviceSplitCandidate> out_candidates) {
   // KeyValuePair here used as threadIdx.x -> gain_value
   using ArgMaxT = cub::KeyValuePair<int, float>;
   using BlockScanT =
@@ -241,7 +239,7 @@ __global__ void EvaluateSplitsKernel(
 
   if (common::IsCat(inputs.feature_types, fidx)) {
     auto n_bins_in_feat = inputs.feature_segments[fidx + 1] - inputs.feature_segments[fidx];
-    if (common::UseOneHot(n_bins_in_feat, inputs.param.max_cat_to_onehot, task)) {
+    if (common::UseOneHot(n_bins_in_feat, inputs.param.max_cat_to_onehot)) {
       EvaluateFeature<BLOCK_THREADS, SumReduceT, BlockScanT, MaxReduceT, TempStorage, GradientSumT,
                       kOneHot>(fidx, inputs, evaluator, sorted_idx, 0, &best_split, &temp_storage);
     } else {
@@ -310,7 +308,7 @@ __device__ void SortBasedSplit(EvaluateSplitInputs<GradientSumT> const &input,
 
 template <typename GradientSumT>
 void GPUHistEvaluator<GradientSumT>::EvaluateSplits(
-    EvaluateSplitInputs<GradientSumT> left, EvaluateSplitInputs<GradientSumT> right, ObjInfo task,
+    EvaluateSplitInputs<GradientSumT> left, EvaluateSplitInputs<GradientSumT> right,
     TreeEvaluator::SplitEvaluator<GPUTrainingParam> evaluator,
     common::Span<DeviceSplitCandidate> out_splits) {
   if (!split_cats_.empty()) {
@@ -323,7 +321,7 @@ void GPUHistEvaluator<GradientSumT>::EvaluateSplits(
   // One block for each feature
   uint32_t constexpr kBlockThreads = 256;
   dh::LaunchKernel {static_cast<uint32_t>(combined_num_features), kBlockThreads, 0}(
-      EvaluateSplitsKernel<kBlockThreads, GradientSumT>, left, right, task, this->SortedIdx(left),
+      EvaluateSplitsKernel<kBlockThreads, GradientSumT>, left, right, this->SortedIdx(left),
       evaluator, dh::ToSpan(feature_best_splits));
 
   // Reduce to get best candidate for left and right child over all features
@@ -365,7 +363,7 @@ void GPUHistEvaluator<GradientSumT>::CopyToHost(EvaluateSplitInputs<GradientSumT
 }
 
 template <typename GradientSumT>
-void GPUHistEvaluator<GradientSumT>::EvaluateSplits(GPUExpandEntry candidate, ObjInfo task,
+void GPUHistEvaluator<GradientSumT>::EvaluateSplits(GPUExpandEntry candidate,
                                                     EvaluateSplitInputs<GradientSumT> left,
                                                     EvaluateSplitInputs<GradientSumT> right,
                                                     common::Span<GPUExpandEntry> out_entries) {
@@ -373,7 +371,7 @@ void GPUHistEvaluator<GradientSumT>::EvaluateSplits(GPUExpandEntry candidate, Ob
 
   dh::TemporaryArray<DeviceSplitCandidate> splits_out_storage(2);
   auto out_splits = dh::ToSpan(splits_out_storage);
-  this->EvaluateSplits(left, right, task, evaluator, out_splits);
+  this->EvaluateSplits(left, right, evaluator, out_splits);
 
   auto d_sorted_idx = this->SortedIdx(left);
   auto d_entries = out_entries;
@@ -385,7 +383,7 @@ void GPUHistEvaluator<GradientSumT>::EvaluateSplits(GPUExpandEntry candidate, Ob
     auto fidx = out_splits[i].findex;
 
     if (split.is_cat &&
-        !common::UseOneHot(input.FeatureBins(fidx), input.param.max_cat_to_onehot, task)) {
+        !common::UseOneHot(input.FeatureBins(fidx), input.param.max_cat_to_onehot)) {
       bool is_left = i == 0;
       auto out = is_left ? cats_out.first(cats_out.size() / 2) : cats_out.last(cats_out.size() / 2);
       SortBasedSplit(input, d_sorted_idx, fidx, is_left, out, &out_splits[i]);
@@ -405,11 +403,11 @@ void GPUHistEvaluator<GradientSumT>::EvaluateSplits(GPUExpandEntry candidate, Ob
 
 template <typename GradientSumT>
 GPUExpandEntry GPUHistEvaluator<GradientSumT>::EvaluateSingleSplit(
-    EvaluateSplitInputs<GradientSumT> input, float weight, ObjInfo task) {
+    EvaluateSplitInputs<GradientSumT> input, float weight) {
   dh::TemporaryArray<DeviceSplitCandidate> splits_out(1);
   auto out_split = dh::ToSpan(splits_out);
   auto evaluator = tree_evaluator_.GetEvaluator<GPUTrainingParam>();
-  this->EvaluateSplits(input, {}, task, evaluator, out_split);
+  this->EvaluateSplits(input, {}, evaluator, out_split);
 
   auto cats_out = this->DeviceCatStorage(input.nidx);
   auto d_sorted_idx = this->SortedIdx(input);
@@ -421,7 +419,7 @@ GPUExpandEntry GPUHistEvaluator<GradientSumT>::EvaluateSingleSplit(
     auto fidx = out_split[i].findex;
 
     if (split.is_cat &&
-        !common::UseOneHot(input.FeatureBins(fidx), input.param.max_cat_to_onehot, task)) {
+        !common::UseOneHot(input.FeatureBins(fidx), input.param.max_cat_to_onehot)) {
       SortBasedSplit(input, d_sorted_idx, fidx, true, cats_out, &out_split[i]);
     }
 

src/tree/gpu_hist/evaluate_splits.cuh

@@ -114,7 +114,7 @@ class GPUHistEvaluator {
   /**
    * \brief Reset the evaluator, should be called before any use.
    */
-  void Reset(common::HistogramCuts const &cuts, common::Span<FeatureType const> ft, ObjInfo task,
+  void Reset(common::HistogramCuts const &cuts, common::Span<FeatureType const> ft,
              bst_feature_t n_features, TrainParam const &param, int32_t device);
 
   /**
@@ -150,21 +150,20 @@ class GPUHistEvaluator {
 
   // impl of evaluate splits, contains CUDA kernels so it's public
   void EvaluateSplits(EvaluateSplitInputs<GradientSumT> left,
-                      EvaluateSplitInputs<GradientSumT> right, ObjInfo task,
+                      EvaluateSplitInputs<GradientSumT> right,
                       TreeEvaluator::SplitEvaluator<GPUTrainingParam> evaluator,
                       common::Span<DeviceSplitCandidate> out_splits);
   /**
    * \brief Evaluate splits for left and right nodes.
    */
-  void EvaluateSplits(GPUExpandEntry candidate, ObjInfo task,
+  void EvaluateSplits(GPUExpandEntry candidate,
                       EvaluateSplitInputs<GradientSumT> left,
                       EvaluateSplitInputs<GradientSumT> right,
                       common::Span<GPUExpandEntry> out_splits);
   /**
    * \brief Evaluate splits for root node.
    */
-  GPUExpandEntry EvaluateSingleSplit(EvaluateSplitInputs<GradientSumT> input, float weight,
+  GPUExpandEntry EvaluateSingleSplit(EvaluateSplitInputs<GradientSumT> input, float weight);
-                                     ObjInfo task);
 };
 }  // namespace tree
 }  // namespace xgboost

src/tree/gpu_hist/evaluator.cu

@@ -16,12 +16,12 @@ namespace xgboost {
 namespace tree {
 template <typename GradientSumT>
 void GPUHistEvaluator<GradientSumT>::Reset(common::HistogramCuts const &cuts,
-                                           common::Span<FeatureType const> ft, ObjInfo task,
+                                           common::Span<FeatureType const> ft,
                                            bst_feature_t n_features, TrainParam const &param,
                                            int32_t device) {
   param_ = param;
   tree_evaluator_ = TreeEvaluator{param, n_features, device};
-  if (cuts.HasCategorical() && !task.UseOneHot()) {
+  if (cuts.HasCategorical()) {
     dh::XGBCachingDeviceAllocator<char> alloc;
     auto ptrs = cuts.cut_ptrs_.ConstDeviceSpan();
     auto beg = thrust::make_counting_iterator<size_t>(1ul);
@@ -34,7 +34,7 @@ void GPUHistEvaluator<GradientSumT>::Reset(common::HistogramCuts const &cuts,
       auto idx = i - 1;
       if (common::IsCat(ft, idx)) {
         auto n_bins = ptrs[i] - ptrs[idx];
-        bool use_sort = !common::UseOneHot(n_bins, to_onehot, task);
+        bool use_sort = !common::UseOneHot(n_bins, to_onehot);
         return use_sort;
       }
       return false;

src/tree/hist/evaluate_splits.h

@@ -11,7 +11,6 @@
 #include <utility>
 #include <vector>
 
-#include "xgboost/task.h"
 #include "../param.h"
 #include "../constraints.h"
 #include "../split_evaluator.h"
@@ -39,7 +38,6 @@ template <typename GradientSumT, typename ExpandEntry> class HistEvaluator {
   int32_t n_threads_ {0};
   FeatureInteractionConstraintHost interaction_constraints_;
   std::vector<NodeEntry> snode_;
-  ObjInfo task_;
 
   // if sum of statistics for non-missing values in the node
   // is equal to sum of statistics for all values:
@@ -244,7 +242,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistEvaluator {
         }
         if (is_cat) {
           auto n_bins = cut_ptrs.at(fidx + 1) - cut_ptrs[fidx];
-          if (common::UseOneHot(n_bins, param_.max_cat_to_onehot, task_)) {
+          if (common::UseOneHot(n_bins, param_.max_cat_to_onehot)) {
             EnumerateSplit<+1, kOneHot>(cut, {}, histogram, fidx, nidx, evaluator, best);
             EnumerateSplit<-1, kOneHot>(cut, {}, histogram, fidx, nidx, evaluator, best);
           } else {
@@ -345,7 +343,6 @@ template <typename GradientSumT, typename ExpandEntry> class HistEvaluator {
 
   auto Evaluator() const { return tree_evaluator_.GetEvaluator(); }
   auto const& Stats() const { return snode_; }
-  auto Task() const { return task_; }
 
   float InitRoot(GradStats const& root_sum) {
     snode_.resize(1);
@@ -363,12 +360,11 @@ template <typename GradientSumT, typename ExpandEntry> class HistEvaluator {
   // The column sampler must be constructed by caller since we need to preserve the rng
   // for the entire training session.
   explicit HistEvaluator(TrainParam const &param, MetaInfo const &info, int32_t n_threads,
-                         std::shared_ptr<common::ColumnSampler> sampler, ObjInfo task)
+                         std::shared_ptr<common::ColumnSampler> sampler)
       : param_{param},
         column_sampler_{std::move(sampler)},
         tree_evaluator_{param, static_cast<bst_feature_t>(info.num_col_), GenericParameter::kCpuId},
-        n_threads_{n_threads},
+        n_threads_{n_threads} {
-        task_{task} {
     interaction_constraints_.Configure(param, info.num_col_);
     column_sampler_->Init(info.num_col_, info.feature_weights.HostVector(), param_.colsample_bynode,
                           param_.colsample_bylevel, param_.colsample_bytree);

src/tree/param.h

@@ -211,12 +211,13 @@ struct TrainParam : public XGBoostParameter<TrainParam> {
       n_nodes = this->max_leaves * 2 - 1;
     } else {
       // bst_node_t will overflow.
-      CHECK_LE(this->max_depth, 31)
+      CHECK_LE(this->max_depth, 30)
-          << "max_depth can not be greater than 31 as that might generate 2 ** "
+          << "max_depth can not be greater than 30 as that might generate 2^31 - 1"
-             "32 - 1 nodes.";
+             "nodes.";
-      n_nodes = (1 << (this->max_depth + 1)) - 1;
+      // same as: (1 << (max_depth + 1)) - 1, but avoids 1 << 31, which overflows.
+      n_nodes = (1 << this->max_depth) + ((1 << this->max_depth) - 1);
     }
-    CHECK_NE(n_nodes, 0);
+    CHECK_GT(n_nodes, 0);
     return n_nodes;
   }
 };

src/tree/updater_approx.cc

@@ -28,10 +28,8 @@ DMLC_REGISTRY_FILE_TAG(updater_approx);
 
 namespace {
 // Return the BatchParam used by DMatrix.
-template <typename GradientSumT>
+auto BatchSpec(TrainParam const &p, common::Span<float> hess, ObjInfo const task) {
-auto BatchSpec(TrainParam const &p, common::Span<float> hess,
+  return BatchParam{p.max_bin, hess, !task.const_hess};
-               HistEvaluator<GradientSumT, CPUExpandEntry> const &evaluator) {
-  return BatchParam{p.max_bin, hess, !evaluator.Task().const_hess};
 }
 
 auto BatchSpec(TrainParam const &p, common::Span<float> hess) {
@@ -46,7 +44,8 @@ class GloablApproxBuilder {
   std::shared_ptr<common::ColumnSampler> col_sampler_;
   HistEvaluator<GradientSumT, CPUExpandEntry> evaluator_;
   HistogramBuilder<GradientSumT, CPUExpandEntry> histogram_builder_;
-  GenericParameter const *ctx_;
+  Context const *ctx_;
+  ObjInfo const task_;
 
   std::vector<ApproxRowPartitioner> partitioner_;
   // Pointer to last updated tree, used for update prediction cache.
@@ -64,8 +63,7 @@ class GloablApproxBuilder {
     int32_t n_total_bins = 0;
     partitioner_.clear();
     // Generating the GHistIndexMatrix is quite slow, is there a way to speed it up?
-    for (auto const &page :
+    for (auto const &page : p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess, task_))) {
-         p_fmat->GetBatches<GHistIndexMatrix>(BatchSpec(param_, hess, evaluator_))) {
       if (n_total_bins == 0) {
         n_total_bins = page.cut.TotalBins();
         feature_values_ = page.cut;
@@ -160,8 +158,9 @@ class GloablApproxBuilder {
                                common::Monitor *monitor)
       : param_{std::move(param)},
         col_sampler_{std::move(column_sampler)},
-        evaluator_{param_, info, ctx->Threads(), col_sampler_, task},
+        evaluator_{param_, info, ctx->Threads(), col_sampler_},
         ctx_{ctx},
+        task_{task},
         monitor_{monitor} {}
 
   void UpdateTree(RegTree *p_tree, std::vector<GradientPair> const &gpair, common::Span<float> hess,

src/tree/updater_gpu_hist.cu

@@ -229,16 +229,14 @@ struct GPUHistMakerDevice {
   // Reset values for each update iteration
   // Note that the column sampler must be passed by value because it is not
   // thread safe
-  void Reset(HostDeviceVector<GradientPair>* dh_gpair, DMatrix* dmat, int64_t num_columns,
+  void Reset(HostDeviceVector<GradientPair>* dh_gpair, DMatrix* dmat, int64_t num_columns) {
-             ObjInfo task) {
     auto const& info = dmat->Info();
     this->column_sampler.Init(num_columns, info.feature_weights.HostVector(),
                               param.colsample_bynode, param.colsample_bylevel,
                               param.colsample_bytree);
     dh::safe_cuda(cudaSetDevice(device_id));
 
-    this->evaluator_.Reset(page->Cuts(), feature_types, task, dmat->Info().num_col_, param,
+    this->evaluator_.Reset(page->Cuts(), feature_types, dmat->Info().num_col_, param, device_id);
-                           device_id);
 
     this->interaction_constraints.Reset();
     std::fill(node_sum_gradients.begin(), node_sum_gradients.end(), GradientPairPrecise{});
@@ -260,7 +258,7 @@ struct GPUHistMakerDevice {
     hist.Reset();
   }
 
-  GPUExpandEntry EvaluateRootSplit(GradientPairPrecise root_sum, float weight, ObjInfo task) {
+  GPUExpandEntry EvaluateRootSplit(GradientPairPrecise root_sum, float weight) {
     int nidx = RegTree::kRoot;
     GPUTrainingParam gpu_param(param);
     auto sampled_features = column_sampler.GetFeatureSet(0);
@@ -277,12 +275,12 @@ struct GPUHistMakerDevice {
                                              matrix.gidx_fvalue_map,
                                              matrix.min_fvalue,
                                              hist.GetNodeHistogram(nidx)};
-    auto split = this->evaluator_.EvaluateSingleSplit(inputs, weight, task);
+    auto split = this->evaluator_.EvaluateSingleSplit(inputs, weight);
     return split;
   }
 
-  void EvaluateLeftRightSplits(GPUExpandEntry candidate, ObjInfo task, int left_nidx,
+  void EvaluateLeftRightSplits(GPUExpandEntry candidate, int left_nidx, int right_nidx,
-                               int right_nidx, const RegTree& tree,
+                               const RegTree& tree,
                                common::Span<GPUExpandEntry> pinned_candidates_out) {
     dh::TemporaryArray<DeviceSplitCandidate> splits_out(2);
     GPUTrainingParam gpu_param(param);
@@ -316,7 +314,7 @@ struct GPUHistMakerDevice {
                                             hist.GetNodeHistogram(right_nidx)};
 
     dh::TemporaryArray<GPUExpandEntry> entries(2);
-    this->evaluator_.EvaluateSplits(candidate, task, left, right, dh::ToSpan(entries));
+    this->evaluator_.EvaluateSplits(candidate, left, right, dh::ToSpan(entries));
     dh::safe_cuda(cudaMemcpyAsync(pinned_candidates_out.data(), entries.data().get(),
                                   sizeof(GPUExpandEntry) * entries.size(), cudaMemcpyDeviceToHost));
   }
@@ -584,7 +582,7 @@ struct GPUHistMakerDevice {
                                   tree[candidate.nid].RightChild());
   }
 
-  GPUExpandEntry InitRoot(RegTree* p_tree, ObjInfo task, dh::AllReducer* reducer) {
+  GPUExpandEntry InitRoot(RegTree* p_tree, dh::AllReducer* reducer) {
     constexpr bst_node_t kRootNIdx = 0;
     dh::XGBCachingDeviceAllocator<char> alloc;
     auto gpair_it = dh::MakeTransformIterator<GradientPairPrecise>(
@@ -605,7 +603,7 @@ struct GPUHistMakerDevice {
     (*p_tree)[kRootNIdx].SetLeaf(param.learning_rate * weight);
 
     // Generate first split
-    auto root_entry = this->EvaluateRootSplit(root_sum, weight, task);
+    auto root_entry = this->EvaluateRootSplit(root_sum, weight);
     return root_entry;
   }
 
@@ -615,11 +613,11 @@ struct GPUHistMakerDevice {
     Driver<GPUExpandEntry> driver(static_cast<TrainParam::TreeGrowPolicy>(param.grow_policy));
 
     monitor.Start("Reset");
-    this->Reset(gpair_all, p_fmat, p_fmat->Info().num_col_, task);
+    this->Reset(gpair_all, p_fmat, p_fmat->Info().num_col_);
     monitor.Stop("Reset");
 
     monitor.Start("InitRoot");
-    driver.Push({ this->InitRoot(p_tree, task, reducer) });
+    driver.Push({ this->InitRoot(p_tree, reducer) });
     monitor.Stop("InitRoot");
 
     auto num_leaves = 1;
@@ -656,7 +654,7 @@ struct GPUHistMakerDevice {
           monitor.Stop("BuildHist");
 
           monitor.Start("EvaluateSplits");
-          this->EvaluateLeftRightSplits(candidate, task, left_child_nidx, right_child_nidx, *p_tree,
+          this->EvaluateLeftRightSplits(candidate, left_child_nidx, right_child_nidx, *p_tree,
                                         new_candidates.subspan(i * 2, 2));
           monitor.Stop("EvaluateSplits");
         } else {

src/tree/updater_quantile_hist.cc

@@ -342,7 +342,7 @@ void QuantileHistMaker::Builder<GradientSumT>::InitData(DMatrix *fmat, const Reg
   // store a pointer to the tree
   p_last_tree_ = &tree;
   evaluator_.reset(new HistEvaluator<GradientSumT, CPUExpandEntry>{
-      param_, info, this->ctx_->Threads(), column_sampler_, task_});
+      param_, info, this->ctx_->Threads(), column_sampler_});
 
   monitor_->Stop(__func__);
 }

tests/ci_build/Dockerfile.clang_tidy

@@ -7,6 +7,7 @@ ENV DEBIAN_FRONTEND noninteractive
 
 # Install all basic requirements
 RUN \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub && \
     apt-get update && \
     apt-get install -y tar unzip wget git build-essential python3 python3-pip software-properties-common \
                        apt-transport-https ca-certificates gnupg-agent && \

tests/ci_build/Dockerfile.cpu

@@ -10,13 +10,13 @@ RUN \
     apt-get install -y software-properties-common && \
     add-apt-repository ppa:ubuntu-toolchain-r/test && \
     apt-get update && \
-    apt-get install -y tar unzip wget git build-essential doxygen graphviz llvm libasan2 libidn11 ninja-build gcc-8 g++-8 && \
+    apt-get install -y tar unzip wget git build-essential doxygen graphviz llvm libasan2 libidn11 ninja-build gcc-8 g++-8 openjdk-8-jdk-headless && \
     # CMake
     wget -nv -nc https://cmake.org/files/v3.14/cmake-3.14.0-Linux-x86_64.sh --no-check-certificate && \
     bash cmake-3.14.0-Linux-x86_64.sh --skip-license --prefix=/usr && \
     # Python
-    wget -O Miniconda3.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    wget https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh && \
-    bash Miniconda3.sh -b -p /opt/python
+    bash Mambaforge-Linux-x86_64.sh -b -p /opt/python
 
 ENV PATH=/opt/python/bin:$PATH
 ENV CC=gcc-8
@@ -24,10 +24,11 @@ ENV CXX=g++-8
 ENV CPP=cpp-8
 
 ENV GOSU_VERSION 1.10
+ENV JAVA_HOME /usr/lib/jvm/java-8-openjdk-amd64/
 
 # Create new Conda environment
 COPY conda_env/cpu_test.yml /scripts/
-RUN conda env create -n cpu_test --file=/scripts/cpu_test.yml
+RUN mamba env create -n cpu_test --file=/scripts/cpu_test.yml
 
 # Install lightweight sudo (not bound to TTY)
 RUN set -ex; \

tests/ci_build/Dockerfile.gpu

@@ -8,8 +8,9 @@ SHELL ["/bin/bash", "-c"]   # Use Bash as shell
 
 # Install all basic requirements
 RUN \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub && \
     apt-get update && \
-    apt-get install -y wget unzip bzip2 libgomp1 build-essential && \
+    apt-get install -y wget unzip bzip2 libgomp1 build-essential openjdk-8-jdk-headless && \
     # Python
     wget -O Miniconda3.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
     bash Miniconda3.sh -b -p /opt/python
@@ -18,11 +19,14 @@ ENV PATH=/opt/python/bin:$PATH
 
 # Create new Conda environment with cuDF, Dask, and cuPy
 RUN \
-    conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
+    conda install -c conda-forge mamba && \
-        python=3.8 cudf=21.10* rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=21.10* dask-cudf=21.10* cupy=9.1* \
+    mamba create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis
+        python=3.8 cudf=22.04* rmm=22.04* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=22.04* dask-cudf=22.04* cupy \
+        numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis \
+        pyspark cloudpickle cuda-python=11.7.0
 
 ENV GOSU_VERSION 1.10
+ENV JAVA_HOME /usr/lib/jvm/java-8-openjdk-amd64/
 
 # Install lightweight sudo (not bound to TTY)
 RUN set -ex; \

tests/ci_build/Dockerfile.gpu_build

@@ -8,6 +8,7 @@ SHELL ["/bin/bash", "-c"]   # Use Bash as shell
 
 # Install all basic requirements
 RUN \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/3bf863cc.pub && \
     apt-get update && \
     apt-get install -y software-properties-common && \
     add-apt-repository ppa:ubuntu-toolchain-r/test && \
@@ -23,7 +24,7 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.7.5-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
     apt-get update && \
     apt-get install -y --allow-downgrades --allow-change-held-packages libnccl2=${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-dev=${NCCL_VERSION}+cuda${CUDA_SHORT}
 

tests/ci_build/Dockerfile.gpu_build_centos7

@@ -4,6 +4,8 @@ ARG CUDA_VERSION_ARG
 
 # Install all basic requirements
 RUN \
+    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
+        > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
     yum install -y epel-release centos-release-scl && \
     yum-config-manager --enable centos-sclo-rh-testing && \
     yum -y update && \
@@ -19,7 +21,7 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.7.3-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
     wget -nv -nc https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     rpm -i nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     yum -y update && \

tests/ci_build/Dockerfile.gpu_build_r_centos7

@@ -4,6 +4,8 @@ ARG CUDA_VERSION_ARG
 
 # Install all basic requirements
 RUN \
+    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
+        > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
     yum install -y epel-release centos-release-scl && \
     yum-config-manager --enable centos-sclo-rh-testing && \
     yum -y update && \

tests/ci_build/Dockerfile.gpu_jvm

@@ -9,6 +9,7 @@ ENV DEBIAN_FRONTEND noninteractive
 
 # Install all basic requirements
 RUN \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/3bf863cc.pub && \
     apt-get update && \
     apt-get install -y software-properties-common && \
     add-apt-repository ppa:openjdk-r/ppa && \

tests/ci_build/Dockerfile.jvm_gpu_build

@@ -4,6 +4,8 @@ ARG CUDA_VERSION_ARG
 
 # Install all basic requirements
 RUN \
+    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
+        > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
     yum install -y epel-release centos-release-scl && \
     yum-config-manager --enable centos-sclo-rh-testing && \
     yum -y update && \
@@ -22,12 +24,10 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.8.3-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
-    wget -nv -nc https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
+    yum-config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo && \
-    rpm -i nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     yum -y update && \
-    yum install -y libnccl-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-devel-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-static-${NCCL_VERSION}+cuda${CUDA_SHORT} && \
+    yum install -y libnccl-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-devel-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-static-${NCCL_VERSION}+cuda${CUDA_SHORT}
-    rm -f nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm;
 
 ENV PATH=/opt/python/bin:/opt/maven/bin:$PATH
 ENV CC=/opt/rh/devtoolset-8/root/usr/bin/gcc

tests/ci_build/Dockerfile.rmm

@@ -8,19 +8,17 @@ SHELL ["/bin/bash", "-c"]   # Use Bash as shell
 
 # Install all basic requirements
 RUN \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub && \
     apt-get update && \
     apt-get install -y wget unzip bzip2 libgomp1 build-essential ninja-build git && \
     # Python
     wget -O Miniconda3.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
-    bash Miniconda3.sh -b -p /opt/python && \
+    bash Miniconda3.sh -b -p /opt/python
-    # CMake
-    wget -nv -nc https://cmake.org/files/v3.14/cmake-3.14.0-Linux-x86_64.sh --no-check-certificate && \
-    bash cmake-3.14.0-Linux-x86_64.sh --skip-license --prefix=/usr
 
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.7.5-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
     apt-get update && \
     apt-get install -y --allow-downgrades --allow-change-held-packages libnccl2=${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-dev=${NCCL_VERSION}+cuda${CUDA_SHORT}
 
@@ -29,7 +27,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with RMM
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.8 rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG
+        python=3.9 rmm=22.04* cudatoolkit=$CUDA_VERSION_ARG cmake
 
 ENV GOSU_VERSION 1.10
 

tests/ci_build/build_python_wheels.sh

@@ -29,13 +29,15 @@ if [[ "$platform_id" == macosx_* ]]; then
         setup_env_var='CIBW_TARGET_OSX_ARM64=1'  # extra flag to be passed to setup.py
         export PYTHON_CROSSENV=1
         export MACOSX_DEPLOYMENT_TARGET=12.0
-        OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-arm64/llvm-openmp-11.1.0-hf3c4609_1.tar.bz2"
+        #OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-arm64/llvm-openmp-11.1.0-hf3c4609_1.tar.bz2"
+        OPENMP_URL="https://xgboost-ci-jenkins-artifacts.s3.us-west-2.amazonaws.com/llvm-openmp-11.1.0-hf3c4609_1-osx-arm64.tar.bz2"
     elif [[ "$platform_id" == macosx_x86_64 ]]; then
         # MacOS, Intel
         wheel_tag=macosx_10_15_x86_64.macosx_11_0_x86_64.macosx_12_0_x86_64
         cpython_ver=37
         export MACOSX_DEPLOYMENT_TARGET=10.13
-        OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-64/llvm-openmp-11.1.0-hda6cdc1_1.tar.bz2"
+        #OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-64/llvm-openmp-11.1.0-hda6cdc1_1.tar.bz2"
+        OPENMP_URL="https://xgboost-ci-jenkins-artifacts.s3.us-west-2.amazonaws.com/llvm-openmp-11.1.0-hda6cdc1_1-osx-64.tar.bz2"
     else
         echo "Platform not supported: $platform_id"
         exit 3

tests/ci_build/conda_env/cpu_test.yml

@@ -30,14 +30,12 @@ dependencies:
 - jsonschema
 - boto3
 - awscli
-- numba
-- llvmlite
 - py-ubjson
 - cffi
 - pyarrow
+- protobuf
+- pyspark>=3.3.0
+- cloudpickle
+- shap
 - pip:
-  - shap
-  - ipython                     # required by shap at import time.
-  - sphinx_rtd_theme
   - datatable
-  - modin[all]

tests/ci_build/entrypoint.sh

@@ -20,9 +20,9 @@ else
 fi
 
 if [[ -n $CI_BUILD_UID ]] && [[ -n $CI_BUILD_GID ]]; then
-    groupadd -o -g "${CI_BUILD_GID}" "${CI_BUILD_GROUP}"
+    groupadd -o -g "${CI_BUILD_GID}" "${CI_BUILD_GROUP}" || true
     useradd -o -m -g "${CI_BUILD_GID}" -u "${CI_BUILD_UID}" \
-        "${CI_BUILD_USER}"
+        "${CI_BUILD_USER}" || true
     export HOME="/home/${CI_BUILD_USER}"
     shopt -s dotglob
     cp -r /root/* "$HOME/"

tests/ci_build/rename_whl.py

@@ -42,4 +42,4 @@ with cd(dirname):
 
     filesize = os.path.getsize(new_name) / 1024 / 1024  # MB
     msg = f"Limit of wheel size set by PyPI is exceeded. {new_name}: {filesize}"
-    assert filesize <= 200, msg
+    assert filesize <= 300, msg

tests/cpp/common/test_hist_util.cc

@@ -381,6 +381,7 @@ void TestSketchFromWeights(bool with_group) {
   ValidateCuts(cuts, m.get(), kBins);
 
   if (with_group) {
+    m->Info().weights_ = decltype(m->Info().weights_)();  // remove weight
     HistogramCuts non_weighted = SketchOnDMatrix(m.get(), kBins, common::OmpGetNumThreads(0));
     for (size_t i = 0; i < cuts.Values().size(); ++i) {
       EXPECT_EQ(cuts.Values()[i], non_weighted.Values()[i]);
@@ -392,6 +393,17 @@ void TestSketchFromWeights(bool with_group) {
       ASSERT_EQ(cuts.Ptrs().at(i), non_weighted.Ptrs().at(i));
     }
   }
+
+  if (with_group) {
+    auto& h_weights = info.weights_.HostVector();
+    h_weights.resize(kGroups);
+    // Generate different weight.
+    for (size_t i = 0; i < h_weights.size(); ++i) {
+      h_weights[i] = static_cast<float>(i + 1) / static_cast<float>(kGroups);
+    }
+    HistogramCuts weighted = SketchOnDMatrix(m.get(), kBins, common::OmpGetNumThreads(0));
+    ValidateCuts(weighted, m.get(), kBins);
+  }
 }
 
 TEST(HistUtil, SketchFromWeights) {

tests/cpp/common/test_hist_util.cu

@@ -593,9 +593,10 @@ void TestAdapterSketchFromWeights(bool with_group) {
   ValidateCuts(cuts, dmat.get(), kBins);
 
   if (with_group) {
+    dmat->Info().weights_ = decltype(dmat->Info().weights_)();  // remove weight
     HistogramCuts non_weighted = DeviceSketch(0, dmat.get(), kBins, 0);
     for (size_t i = 0; i < cuts.Values().size(); ++i) {
-      EXPECT_EQ(cuts.Values()[i], non_weighted.Values()[i]);
+      ASSERT_EQ(cuts.Values()[i], non_weighted.Values()[i]);
     }
     for (size_t i = 0; i < cuts.MinValues().size(); ++i) {
       ASSERT_EQ(cuts.MinValues()[i], non_weighted.MinValues()[i]);
@@ -604,6 +605,24 @@ void TestAdapterSketchFromWeights(bool with_group) {
       ASSERT_EQ(cuts.Ptrs().at(i), non_weighted.Ptrs().at(i));
     }
   }
+
+  if (with_group) {
+    common::HistogramCuts weighted;
+    auto& h_weights = info.weights_.HostVector();
+    h_weights.resize(kGroups);
+    // Generate different weight.
+    for (size_t i = 0; i < h_weights.size(); ++i) {
+      // FIXME(jiamingy): Some entries generated GPU test cannot pass the validate cuts if
+      // we use more diverse weights, partially caused by
+      // https://github.com/dmlc/xgboost/issues/7946
+      h_weights[i] = (i % 2 == 0 ? 1 : 2) / static_cast<float>(kGroups);
+    }
+    SketchContainer sketch_container(ft, kBins, kCols, kRows, 0);
+    AdapterDeviceSketch(adapter.Value(), kBins, info, std::numeric_limits<float>::quiet_NaN(),
+                        &sketch_container);
+    sketch_container.MakeCuts(&weighted);
+    ValidateCuts(weighted, dmat.get(), kBins);
+  }
 }
 
 TEST(HistUtil, AdapterSketchFromWeights) {

tests/cpp/common/test_hist_util.h

@@ -98,7 +98,11 @@ inline void TestBinDistribution(const HistogramCuts &cuts, int column_idx,
                                 int num_bins) {
   std::map<int, int> bin_weights;
   for (auto i = 0ull; i < sorted_column.size(); i++) {
-    bin_weights[cuts.SearchBin(sorted_column[i], column_idx)] += sorted_weights[i];
+    auto bin_idx = cuts.SearchBin(sorted_column[i], column_idx);
+    if (bin_weights.find(bin_idx) == bin_weights.cend()) {
+      bin_weights[bin_idx] = 0;
+    }
+    bin_weights.at(bin_idx) += sorted_weights[i];
   }
   int local_num_bins = cuts.Ptrs()[column_idx + 1] - cuts.Ptrs()[column_idx];
   auto total_weight = std::accumulate(sorted_weights.begin(), sorted_weights.end(),0);
@@ -176,8 +180,7 @@ inline void ValidateColumn(const HistogramCuts& cuts, int column_idx,
   }
 }
 
-inline void ValidateCuts(const HistogramCuts& cuts, DMatrix* dmat,
+inline void ValidateCuts(const HistogramCuts& cuts, DMatrix* dmat, int num_bins) {
-                         int num_bins) {
   // Collect data into columns
   std::vector<std::vector<float>> columns(dmat->Info().num_col_);
   for (auto& batch : dmat->GetBatches<SparsePage>()) {
@@ -189,17 +192,22 @@ inline void ValidateCuts(const HistogramCuts& cuts, DMatrix* dmat,
       }
     }
   }
+
+  // construct weights.
+  std::vector<float> w = dmat->Info().group_ptr_.empty() ? dmat->Info().weights_.HostVector()
+                                                         : detail::UnrollGroupWeights(dmat->Info());
+
   // Sort
   for (auto i = 0ull; i < columns.size(); i++) {
     auto& col = columns.at(i);
-    const auto& w = dmat->Info().weights_.HostVector();
+    std::vector<size_t> index(col.size());
-    std::vector<size_t > index(col.size());
     std::iota(index.begin(), index.end(), 0);
-    std::sort(index.begin(), index.end(),
+    std::sort(index.begin(), index.end(), [=](size_t a, size_t b) { return col[a] < col[b]; });
-              [=](size_t a, size_t b) { return col[a] < col[b]; });
 
     std::vector<float> sorted_column(col.size());
     std::vector<float> sorted_weights(col.size(), 1.0);
+    const auto& w = dmat->Info().weights_.HostVector();
+
     for (auto j = 0ull; j < col.size(); j++) {
       sorted_column[j] = col[index[j]];
       if (w.size() == col.size()) {

tests/cpp/common/test_span.cu

@@ -252,7 +252,7 @@ __global__ void TestLastStaticKernel(Span<float> _span) {
   _span.last(static_cast<Span<float>::index_type>(-1));
 }
 
-TEST(GPUSpan, FirstLast) {
+TEST(GPUSpanDeathTest, FirstLast) {
   // We construct vectors multiple times since thrust can not recover from
   // death test.
   auto lambda_first_dy = []() {
@@ -312,40 +312,37 @@ TEST(GPUSpan, FirstLast) {
   output = testing::internal::GetCapturedStdout();
 }
 
-__global__ void TestFrontKernel(Span<float> _span)  {
+namespace {
-  _span.front();
+void TestFrontBack() {
-}
-
-__global__ void TestBackKernel(Span<float> _span)  {
-  _span.back();
-}
-
-TEST(GPUSpan, FrontBack) {
-  dh::safe_cuda(cudaSetDevice(0));
-
   Span<float> s;
-  auto lambda_test_front = [=]() {
+  EXPECT_DEATH(
-    // make sure the termination happens inside this test.
+      {
-    try {
+        // make sure the termination happens inside this test.
-      TestFrontKernel<<<1, 1>>>(s);
+        try {
-      dh::safe_cuda(cudaDeviceSynchronize());
+          dh::LaunchN(1, [=] __device__(size_t) { s.front(); });
-      dh::safe_cuda(cudaGetLastError());
+          dh::safe_cuda(cudaDeviceSynchronize());
-    } catch (dmlc::Error const& e) {
+          dh::safe_cuda(cudaGetLastError());
-      std::terminate();
+        } catch (dmlc::Error const& e) {
-    }
+          std::terminate();
-  };
+        }
-  EXPECT_DEATH(lambda_test_front(), "");
+      },
+      "");
+  EXPECT_DEATH(
+      {
+        try {
+          dh::LaunchN(1, [=] __device__(size_t) { s.back(); });
+          dh::safe_cuda(cudaDeviceSynchronize());
+          dh::safe_cuda(cudaGetLastError());
+        } catch (dmlc::Error const& e) {
+          std::terminate();
+        }
+      },
+      "");
+}
+}  // namespace
 
-  auto lambda_test_back = [=]() {
+TEST(GPUSpanDeathTest, FrontBack) {
-    try {
+  TestFrontBack();
-      TestBackKernel<<<1, 1>>>(s);
-      dh::safe_cuda(cudaDeviceSynchronize());
-      dh::safe_cuda(cudaGetLastError());
-    } catch (dmlc::Error const& e) {
-      std::terminate();
-    }
-  };
-  EXPECT_DEATH(lambda_test_back(), "");
 }
 
 __global__ void TestSubspanDynamicKernel(Span<float> _span) {
@@ -354,7 +351,7 @@ __global__ void TestSubspanDynamicKernel(Span<float> _span) {
 __global__ void TestSubspanStaticKernel(Span<float> _span) {
   _span.subspan<16>();
 }
-TEST(GPUSpan, Subspan) {
+TEST(GPUSpanDeathTest, Subspan) {
   auto lambda_subspan_dynamic = []() {
     thrust::host_vector<float> h_vec (4);
     InitializeRange(h_vec.begin(), h_vec.end());

tests/cpp/tree/gpu_hist/test_evaluate_splits.cu

@@ -57,8 +57,7 @@ void TestEvaluateSingleSplit(bool is_categorical) {
   GPUHistEvaluator<GradientPair> evaluator{
       tparam, static_cast<bst_feature_t>(feature_min_values.size()), 0};
   dh::device_vector<common::CatBitField::value_type> out_cats;
-  DeviceSplitCandidate result =
+  DeviceSplitCandidate result = evaluator.EvaluateSingleSplit(input, 0).split;
-      evaluator.EvaluateSingleSplit(input, 0, ObjInfo{ObjInfo::kRegression}).split;
 
   EXPECT_EQ(result.findex, 1);
   EXPECT_EQ(result.fvalue, 11.0);
@@ -101,8 +100,7 @@ TEST(GpuHist, EvaluateSingleSplitMissing) {
                                           dh::ToSpan(feature_histogram)};
 
   GPUHistEvaluator<GradientPair> evaluator(tparam, feature_set.size(), 0);
-  DeviceSplitCandidate result =
+  DeviceSplitCandidate result = evaluator.EvaluateSingleSplit(input, 0).split;
-      evaluator.EvaluateSingleSplit(input, 0, ObjInfo{ObjInfo::kRegression}).split;
 
   EXPECT_EQ(result.findex, 0);
   EXPECT_EQ(result.fvalue, 1.0);
@@ -114,10 +112,8 @@ TEST(GpuHist, EvaluateSingleSplitMissing) {
 TEST(GpuHist, EvaluateSingleSplitEmpty) {
   TrainParam tparam = ZeroParam();
   GPUHistEvaluator<GradientPair> evaluator(tparam, 1, 0);
-  DeviceSplitCandidate result = evaluator
+  DeviceSplitCandidate result =
-                                    .EvaluateSingleSplit(EvaluateSplitInputs<GradientPair>{}, 0,
+      evaluator.EvaluateSingleSplit(EvaluateSplitInputs<GradientPair>{}, 0).split;
-                                                         ObjInfo{ObjInfo::kRegression})
-                                    .split;
   EXPECT_EQ(result.findex, -1);
   EXPECT_LT(result.loss_chg, 0.0f);
 }
@@ -152,8 +148,7 @@ TEST(GpuHist, EvaluateSingleSplitFeatureSampling) {
                                           dh::ToSpan(feature_histogram)};
 
   GPUHistEvaluator<GradientPair> evaluator(tparam, feature_min_values.size(), 0);
-  DeviceSplitCandidate result =
+  DeviceSplitCandidate result = evaluator.EvaluateSingleSplit(input, 0).split;
-      evaluator.EvaluateSingleSplit(input, 0, ObjInfo{ObjInfo::kRegression}).split;
 
   EXPECT_EQ(result.findex, 1);
   EXPECT_EQ(result.fvalue, 11.0);
@@ -191,8 +186,7 @@ TEST(GpuHist, EvaluateSingleSplitBreakTies) {
                                           dh::ToSpan(feature_histogram)};
 
   GPUHistEvaluator<GradientPair> evaluator(tparam, feature_min_values.size(), 0);
-  DeviceSplitCandidate result =
+  DeviceSplitCandidate result = evaluator.EvaluateSingleSplit(input, 0).split;
-      evaluator.EvaluateSingleSplit(input, 0, ObjInfo{ObjInfo::kRegression}).split;
 
   EXPECT_EQ(result.findex, 0);
   EXPECT_EQ(result.fvalue, 1.0);
@@ -243,8 +237,8 @@ TEST(GpuHist, EvaluateSplits) {
 
   GPUHistEvaluator<GradientPair> evaluator{
       tparam, static_cast<bst_feature_t>(feature_min_values.size()), 0};
-  evaluator.EvaluateSplits(input_left, input_right, ObjInfo{ObjInfo::kRegression},
+  evaluator.EvaluateSplits(input_left, input_right, evaluator.GetEvaluator(),
-                           evaluator.GetEvaluator(), dh::ToSpan(out_splits));
+                           dh::ToSpan(out_splits));
 
   DeviceSplitCandidate result_left = out_splits[0];
   EXPECT_EQ(result_left.findex, 1);
@@ -264,8 +258,7 @@ TEST_F(TestPartitionBasedSplit, GpuHist) {
   cuts_.cut_values_.SetDevice(0);
   cuts_.min_vals_.SetDevice(0);
 
-  ObjInfo task{ObjInfo::kRegression};
+  evaluator.Reset(cuts_, dh::ToSpan(ft), info_.num_col_, param_, 0);
-  evaluator.Reset(cuts_, dh::ToSpan(ft), task, info_.num_col_, param_, 0);
 
   dh::device_vector<GradientPairPrecise> d_hist(hist_[0].size());
   auto node_hist = hist_[0];
@@ -282,7 +275,7 @@ TEST_F(TestPartitionBasedSplit, GpuHist) {
                                                  cuts_.cut_values_.ConstDeviceSpan(),
                                                  cuts_.min_vals_.ConstDeviceSpan(),
                                                  dh::ToSpan(d_hist)};
-  auto split = evaluator.EvaluateSingleSplit(input, 0, ObjInfo{ObjInfo::kRegression}).split;
+  auto split = evaluator.EvaluateSingleSplit(input, 0).split;
   ASSERT_NEAR(split.loss_chg, best_score_, 1e-16);
 }
 }  // namespace tree

tests/cpp/tree/hist/test_evaluate_splits.cc

@@ -24,8 +24,8 @@ template <typename GradientSumT> void TestEvaluateSplits() {
 
   auto dmat = RandomDataGenerator(kRows, kCols, 0).Seed(3).GenerateDMatrix();
 
-  auto evaluator = HistEvaluator<GradientSumT, CPUExpandEntry>{
+  auto evaluator =
-      param, dmat->Info(), n_threads, sampler, ObjInfo{ObjInfo::kRegression}};
+      HistEvaluator<GradientSumT, CPUExpandEntry>{param, dmat->Info(), n_threads, sampler};
   common::HistCollection<GradientSumT> hist;
   std::vector<GradientPair> row_gpairs = {
       {1.23f, 0.24f}, {0.24f, 0.25f}, {0.26f, 0.27f},  {2.27f, 0.28f},
@@ -97,8 +97,7 @@ TEST(HistEvaluator, Apply) {
   param.UpdateAllowUnknown(Args{{"min_child_weight", "0"}, {"reg_lambda", "0.0"}});
   auto dmat = RandomDataGenerator(kNRows, kNCols, 0).Seed(3).GenerateDMatrix();
   auto sampler = std::make_shared<common::ColumnSampler>();
-  auto evaluator_ = HistEvaluator<float, CPUExpandEntry>{param, dmat->Info(), 4, sampler,
+  auto evaluator_ = HistEvaluator<float, CPUExpandEntry>{param, dmat->Info(), 4, sampler};
-                                                         ObjInfo{ObjInfo::kRegression}};
 
   CPUExpandEntry entry{0, 0, 10.0f};
   entry.split.left_sum = GradStats{0.4, 0.6f};
@@ -125,7 +124,7 @@ TEST_F(TestPartitionBasedSplit, CPUHist) {
   std::vector<FeatureType> ft{FeatureType::kCategorical};
   auto sampler = std::make_shared<common::ColumnSampler>();
   HistEvaluator<double, CPUExpandEntry> evaluator{param_, info_, common::OmpGetNumThreads(0),
-                                                  sampler, ObjInfo{ObjInfo::kRegression}};
+                                                  sampler};
   evaluator.InitRoot(GradStats{total_gpair_});
   RegTree tree;
   std::vector<CPUExpandEntry> entries(1);
@@ -156,8 +155,8 @@ auto CompareOneHotAndPartition(bool onehot) {
 
   int32_t n_threads = 16;
   auto sampler = std::make_shared<common::ColumnSampler>();
-  auto evaluator = HistEvaluator<GradientSumT, CPUExpandEntry>{
+  auto evaluator =
-      param, dmat->Info(), n_threads, sampler, ObjInfo{ObjInfo::kRegression}};
+      HistEvaluator<GradientSumT, CPUExpandEntry>{param, dmat->Info(), n_threads, sampler};
   std::vector<CPUExpandEntry> entries(1);
 
   for (auto const &gmat : dmat->GetBatches<GHistIndexMatrix>({32, param.sparse_threshold})) {

tests/cpp/tree/test_gpu_hist.cu

@@ -262,7 +262,7 @@ TEST(GpuHist, EvaluateRootSplit) {
   info.num_col_ = kNCols;
 
   DeviceSplitCandidate res =
-      maker.EvaluateRootSplit({6.4f, 12.8f}, 0, ObjInfo{ObjInfo::kRegression}).split;
+      maker.EvaluateRootSplit({6.4f, 12.8f}, 0).split;
 
   ASSERT_EQ(res.findex, 7);
   ASSERT_NEAR(res.fvalue, 0.26, xgboost::kRtEps);
@@ -300,11 +300,11 @@ void TestHistogramIndexImpl() {
   const auto &maker = hist_maker.maker;
   auto grad = GenerateRandomGradients(kNRows);
   grad.SetDevice(0);
-  maker->Reset(&grad, hist_maker_dmat.get(), kNCols, ObjInfo{ObjInfo::kRegression});
+  maker->Reset(&grad, hist_maker_dmat.get(), kNCols);
   std::vector<common::CompressedByteT> h_gidx_buffer(maker->page->gidx_buffer.HostVector());
 
   const auto &maker_ext = hist_maker_ext.maker;
-  maker_ext->Reset(&grad, hist_maker_ext_dmat.get(), kNCols, ObjInfo{ObjInfo::kRegression});
+  maker_ext->Reset(&grad, hist_maker_ext_dmat.get(), kNCols);
   std::vector<common::CompressedByteT> h_gidx_buffer_ext(maker_ext->page->gidx_buffer.HostVector());
 
   ASSERT_EQ(maker->page->Cuts().TotalBins(), maker_ext->page->Cuts().TotalBins());

tests/python-gpu/test_gpu_updaters.py

@@ -61,6 +61,9 @@ class TestGPUUpdaters:
     def test_categorical(self, rows, cols, rounds, cats):
         self.cputest.run_categorical_basic(rows, cols, rounds, cats, "gpu_hist")
 
+    def test_max_cat(self) -> None:
+        self.cputest.run_max_cat("gpu_hist")
+
     def test_categorical_32_cat(self):
         '''32 hits the bound of integer bitset, so special test'''
         rows = 1000

tests/python-gpu/test_gpu_with_dask.py

@@ -82,7 +82,7 @@ def run_with_dask_dataframe(DMatrixT: Type, client: Client) -> None:
 
     cp.testing.assert_allclose(single_node, predictions.compute())
     np.testing.assert_allclose(single_node,
-                               series_predictions.compute().to_array())
+                               series_predictions.compute().to_numpy())
 
     predt = dxgb.predict(client, out, X)
     assert isinstance(predt, dd.Series)

tests/python-gpu/test_large_input.py

@@ -18,4 +18,6 @@ def test_large_input():
     X = cp.ones((m, n), dtype=np.float32)
     y = cp.ones(m)
     dmat = xgb.DeviceQuantileDMatrix(X, y)
-    xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)
+    booster = xgb.train({"tree_method": "gpu_hist", "max_depth": 1}, dmat, 1)
+    del y
+    booster.inplace_predict(X)

tests/python/test_monotone_constraints.py

@@ -93,6 +93,11 @@ class TestMonotoneConstraints:
         constrained = xgb.train(params_for_constrained, training_dset)
         assert is_correctly_constrained(constrained)
 
+    def test_monotone_constraints_tuple(self) -> None:
+        params_for_constrained = {"monotone_constraints": (1, -1)}
+        constrained = xgb.train(params_for_constrained, training_dset)
+        assert is_correctly_constrained(constrained)
+
     @pytest.mark.parametrize('format', [dict, list])
     def test_monotone_constraints_feature_names(self, format):
 

tests/python/test_tracker.py

@@ -4,6 +4,7 @@ import pytest
 import testing as tm
 import numpy as np
 import sys
+import re
 
 if sys.platform.startswith("win"):
     pytest.skip("Skipping dask tests on Windows", allow_module_level=True)
@@ -59,3 +60,34 @@ def test_rabit_ops():
     with LocalCluster(n_workers=n_workers) as cluster:
         with Client(cluster) as client:
             run_rabit_ops(client, n_workers)
+
+
+def test_rank_assignment() -> None:
+    from distributed import Client, LocalCluster
+    from test_with_dask import _get_client_workers
+
+    def local_test(worker_id):
+        with xgb.dask.RabitContext(args):
+            for val in args:
+                sval = val.decode("utf-8")
+                if sval.startswith("DMLC_TASK_ID"):
+                    task_id = sval
+                    break
+            matched = re.search(".*-([0-9]).*", task_id)
+            rank = xgb.rabit.get_rank()
+            # As long as the number of workers is lesser than 10, rank and worker id
+            # should be the same
+            assert rank == int(matched.group(1))
+
+    with LocalCluster(n_workers=8) as cluster:
+        with Client(cluster) as client:
+            workers = _get_client_workers(client)
+            args = client.sync(
+                xgb.dask._get_rabit_args,
+                len(workers),
+                None,
+                client,
+            )
+
+            futures = client.map(local_test, range(len(workers)), workers=workers)
+            client.gather(futures)

tests/python/test_updaters.py

@@ -1,3 +1,5 @@
+from random import choice
+from string import ascii_lowercase
 import testing as tm
 import pytest
 import xgboost as xgb
@@ -8,10 +10,10 @@ exact_parameter_strategy = strategies.fixed_dictionaries({
     'nthread': strategies.integers(1, 4),
     'max_depth': strategies.integers(1, 11),
     'min_child_weight': strategies.floats(0.5, 2.0),
-    'alpha': strategies.floats(0.0, 2.0),
+    'alpha': strategies.floats(1e-5, 2.0),
     'lambda': strategies.floats(1e-5, 2.0),
     'eta': strategies.floats(0.01, 0.5),
-    'gamma': strategies.floats(0.0, 2.0),
+    'gamma': strategies.floats(1e-5, 2.0),
     'seed': strategies.integers(0, 10),
     # We cannot enable subsampling as the training loss can increase
     # 'subsample': strategies.floats(0.5, 1.0),
@@ -167,6 +169,30 @@ class TestTreeMethod:
 
     def test_invalid_category(self) -> None:
         self.run_invalid_category("approx")
+        self.run_invalid_category("hist")
+
+    def run_max_cat(self, tree_method: str) -> None:
+        """Test data with size smaller than number of categories."""
+        import pandas as pd
+        n_cat = 100
+        n = 5
+        X = pd.Series(
+            ["".join(choice(ascii_lowercase) for i in range(3)) for i in range(n_cat)],
+            dtype="category",
+        )[:n].to_frame()
+
+        reg = xgb.XGBRegressor(
+            enable_categorical=True,
+            tree_method=tree_method,
+            n_estimators=10,
+        )
+        y = pd.Series(range(n))
+        reg.fit(X=X, y=y, eval_set=[(X, y)])
+        assert tm.non_increasing(reg.evals_result()["validation_0"]["rmse"])
+
+    @pytest.mark.parametrize("tree_method", ["hist", "approx"])
+    def test_max_cat(self, tree_method) -> None:
+        self.run_max_cat(tree_method)
 
     def run_categorical_basic(self, rows, cols, rounds, cats, tree_method):
         onehot, label = tm.make_categorical(rows, cols, cats, True)

tests/python/test_with_dask.py

@@ -1,6 +1,7 @@
 """Copyright 2019-2022 XGBoost contributors"""
 from pathlib import Path
 import pickle
+import socket
 import testing as tm
 import pytest
 import xgboost as xgb
@@ -110,9 +111,10 @@ def make_categorical(
 
 
 def generate_array(
-        with_weights: bool = False
+    with_weights: bool = False,
-) -> Tuple[xgb.dask._DaskCollection, xgb.dask._DaskCollection,
+) -> Tuple[
-           Optional[xgb.dask._DaskCollection]]:
+    xgb.dask._DaskCollection, xgb.dask._DaskCollection, Optional[xgb.dask._DaskCollection]
+]:
     chunk_size = 20
     rng = da.random.RandomState(1994)
     X = rng.random_sample((kRows, kCols), chunks=(chunk_size, -1))
@@ -1189,6 +1191,50 @@ def test_dask_iteration_range(client: "Client"):
 
 
 class TestWithDask:
+    def test_dmatrix_binary(self, client: "Client") -> None:
+        def save_dmatrix(rabit_args: List[bytes], tmpdir: str) -> None:
+            with xgb.dask.RabitContext(rabit_args):
+                rank = xgb.rabit.get_rank()
+                X, y = tm.make_categorical(100, 4, 4, False)
+                Xy = xgb.DMatrix(X, y, enable_categorical=True)
+                path = os.path.join(tmpdir, f"{rank}.bin")
+                Xy.save_binary(path)
+
+        def load_dmatrix(rabit_args: List[bytes], tmpdir: str) -> None:
+            with xgb.dask.RabitContext(rabit_args):
+                rank = xgb.rabit.get_rank()
+                path = os.path.join(tmpdir, f"{rank}.bin")
+                Xy = xgb.DMatrix(path)
+                assert Xy.num_row() == 100
+                assert Xy.num_col() == 4
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            workers = _get_client_workers(client)
+            rabit_args = client.sync(
+                xgb.dask._get_rabit_args, len(workers), None, client
+            )
+            futures = []
+            for w in workers:
+                # same argument for each worker, must set pure to False otherwise dask
+                # will try to reuse the result from the first worker and hang waiting
+                # for it.
+                f = client.submit(
+                    save_dmatrix, rabit_args, tmpdir, workers=[w], pure=False
+                )
+                futures.append(f)
+            client.gather(futures)
+
+            rabit_args = client.sync(
+                xgb.dask._get_rabit_args, len(workers), None, client
+            )
+            futures = []
+            for w in workers:
+                f = client.submit(
+                    load_dmatrix, rabit_args, tmpdir, workers=[w], pure=False
+                )
+                futures.append(f)
+            client.gather(futures)
+
     @pytest.mark.parametrize('config_key,config_value', [('verbosity', 0), ('use_rmm', True)])
     def test_global_config(
             self,
@@ -1240,11 +1286,11 @@ class TestWithDask:
         os.remove(after_fname)
 
         with dask.config.set({'xgboost.foo': "bar"}):
-            with pytest.raises(ValueError):
+            with pytest.raises(ValueError, match=r"Unknown configuration.*"):
                 xgb.dask.train(client, {}, dtrain, num_boost_round=4)
 
-        with dask.config.set({'xgboost.scheduler_address': "127.0.0.1:22"}):
+        with dask.config.set({'xgboost.scheduler_address': "127.0.0.1:foo"}):
-            with pytest.raises(PermissionError):
+            with pytest.raises(socket.gaierror, match=r".*not known.*"):
                 xgb.dask.train(client, {}, dtrain, num_boost_round=1)
 
     def run_updater_test(