Revamp the rabit implementation. (#10112)

This PR replaces the original RABIT implementation with a new one, which has already been partially merged into XGBoost. The new one features:
- Federated learning for both CPU and GPU.
- NCCL.
- More data types.
- A unified interface for all the underlying implementations.
- Improved timeout handling for both tracker and workers.
- Exhausted tests with metrics (fixed a couple of bugs along the way).
- A reusable tracker for Python and JVM packages.

tests/cpp/metric/test_auc.cc

@@ -1,68 +0,0 @@
-#include "test_auc.h"
-
-#include <xgboost/metric.h>
-
-namespace xgboost {
-namespace metric {
-
-TEST(Metric, DeclareUnifiedTest(BinaryAUC)) { VerifyBinaryAUC(); }
-
-TEST(Metric, DeclareUnifiedTest(MultiClassAUC)) { VerifyMultiClassAUC(); }
-
-TEST(Metric, DeclareUnifiedTest(RankingAUC)) { VerifyRankingAUC(); }
-
-TEST(Metric, DeclareUnifiedTest(PRAUC)) { VerifyPRAUC(); }
-
-TEST(Metric, DeclareUnifiedTest(MultiClassPRAUC)) { VerifyMultiClassPRAUC(); }
-
-TEST(Metric, DeclareUnifiedTest(RankingPRAUC)) { VerifyRankingPRAUC(); }
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), BinaryAUCRowSplit) {
-  DoTest(VerifyBinaryAUC, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), BinaryAUCColumnSplit) {
-  DoTest(VerifyBinaryAUC, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassAUCRowSplit) {
-  DoTest(VerifyMultiClassAUC, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassAUCColumnSplit) {
-  DoTest(VerifyMultiClassAUC, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RankingAUCRowSplit) {
-  DoTest(VerifyRankingAUC, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RankingAUCColumnSplit) {
-  DoTest(VerifyRankingAUC, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), PRAUCRowSplit) {
-  DoTest(VerifyPRAUC, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), PRAUCColumnSplit) {
-  DoTest(VerifyPRAUC, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassPRAUCRowSplit) {
-  DoTest(VerifyMultiClassPRAUC, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassPRAUCColumnSplit) {
-  DoTest(VerifyMultiClassPRAUC, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RankingPRAUCRowSplit) {
-  DoTest(VerifyRankingPRAUC, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RankingPRAUCColumnSplit) {
-  DoTest(VerifyRankingPRAUC, DataSplitMode::kCol);
-}
-}  // namespace metric
-}  // namespace xgboost

tests/cpp/metric/test_auc.cu

@@ -1,5 +0,0 @@
-/*!
- * Copyright 2021 XGBoost contributors
- */
-// Dummy file to keep the CUDA conditional compile trick.
-#include "test_auc.cc"

tests/cpp/metric/test_auc.h

@@ -7,11 +7,9 @@
 
 #include "../helpers.h"
 
-namespace xgboost {
-namespace metric {
-
-inline void VerifyBinaryAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+namespace xgboost::metric {
+inline void VerifyBinaryAUC(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   std::unique_ptr<Metric> uni_ptr{Metric::Create("auc", &ctx)};
   Metric* metric = uni_ptr.get();
   ASSERT_STREQ(metric->Name(), "auc");
@@ -53,8 +51,8 @@ inline void VerifyBinaryAUC(DataSplitMode data_split_mode = DataSplitMode::kRow)
               0.5, 1e-10);
 }
 
-inline void VerifyMultiClassAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMultiClassAUC(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   std::unique_ptr<Metric> uni_ptr{Metric::Create("auc", &ctx)};
   auto metric = uni_ptr.get();
 
@@ -114,8 +112,8 @@ inline void VerifyMultiClassAUC(DataSplitMode data_split_mode = DataSplitMode::k
   ASSERT_GT(auc, 0.714);
 }
 
-inline void VerifyRankingAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyRankingAUC(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   std::unique_ptr<Metric> metric{Metric::Create("auc", &ctx)};
 
   // single group
@@ -148,8 +146,8 @@ inline void VerifyRankingAUC(DataSplitMode data_split_mode = DataSplitMode::kRow
               0.769841f, 1e-6);
 }
 
-inline void VerifyPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyPRAUC(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
 
   xgboost::Metric* metric = xgboost::Metric::Create("aucpr", &ctx);
   ASSERT_STREQ(metric->Name(), "aucpr");
@@ -185,8 +183,8 @@ inline void VerifyPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
   delete metric;
 }
 
-inline void VerifyMultiClassPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMultiClassPRAUC(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
 
   std::unique_ptr<Metric> metric{Metric::Create("aucpr", &ctx)};
 
@@ -209,8 +207,8 @@ inline void VerifyMultiClassPRAUC(DataSplitMode data_split_mode = DataSplitMode:
   ASSERT_GT(auc, 0.699);
 }
 
-inline void VerifyRankingPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyRankingPRAUC(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
 
   std::unique_ptr<Metric> metric{Metric::Create("aucpr", &ctx)};
 
@@ -245,5 +243,4 @@ inline void VerifyRankingPRAUC(DataSplitMode data_split_mode = DataSplitMode::kR
                     data_split_mode),
       0.556021f, 0.001f);
 }
-}  // namespace metric
-}  // namespace xgboost
+}  // namespace xgboost::metric

tests/cpp/metric/test_distributed_metric.cc

@@ -0,0 +1,192 @@
+/**
+ * Copyright 2023, XGBoost contributors
+ */
+#include <gtest/gtest.h>
+#include <xgboost/context.h>  // for DeviceOrd
+#include <xgboost/data.h>     // for DataSplitMode
+
+#include <algorithm>   // for min
+#include <cstdint>     // for int32_t
+#include <functional>  // for function
+#include <string>      // for string
+#include <thread>      // for thread
+
+#include "../collective/test_worker.h"  // for TestDistributedGlobal
+#include "test_auc.h"
+#include "test_elementwise_metric.h"
+#include "test_multiclass_metric.h"
+#include "test_rank_metric.h"
+#include "test_survival_metric.h"
+
+#if defined(XGBOOST_USE_FEDERATED)
+
+#include "../plugin/federated/test_worker.h"  // for TestFederatedGlobal
+
+#endif  // defined(XGBOOST_USE_FEDERATED)
+
+namespace xgboost::metric {
+namespace {
+using Verifier = std::function<void(DataSplitMode, DeviceOrd)>;
+struct Param {
+  bool is_dist;         // is distributed
+  bool is_fed;          // is federated learning
+  DataSplitMode split;  // how to split data
+  Verifier v;           // test function
+  std::string name;     // metric name
+  DeviceOrd device;     // device to run
+};
+
+class TestDistributedMetric : public ::testing::TestWithParam<Param> {
+ protected:
+  template <typename Fn>
+  void Run(bool is_dist, bool is_fed, DataSplitMode split_mode, Fn fn, DeviceOrd device) {
+    if (!is_dist) {
+      fn(split_mode, device);
+      return;
+    }
+
+    std::int32_t n_workers{0};
+    if (device.IsCUDA()) {
+      n_workers = common::AllVisibleGPUs();
+    } else {
+      n_workers = std::min(static_cast<std::int32_t>(std::thread::hardware_concurrency()), 3);
+    }
+    auto fn1 = [&]() {
+      auto r = collective::GetRank();
+      if (device.IsCPU()) {
+        fn(split_mode, DeviceOrd::CPU());
+      } else {
+        fn(split_mode, DeviceOrd::CUDA(r));
+      }
+    };
+    if (is_fed) {
+#if defined(XGBOOST_USE_FEDERATED)
+      collective::TestFederatedGlobal(n_workers, fn1);
+#endif  // defined(XGBOOST_USE_FEDERATED)
+    } else {
+      collective::TestDistributedGlobal(n_workers, fn1);
+    }
+  }
+};
+}  // anonymous namespace
+
+TEST_P(TestDistributedMetric, BinaryAUCRowSplit) {
+  auto p = GetParam();
+  this->Run(p.is_dist, p.is_fed, p.split, p.v, p.device);
+}
+
+constexpr bool UseNCCL() {
+#if defined(XGBOOST_USE_NCCL)
+  return true;
+#else
+  return false;
+#endif  // defined(XGBOOST_USE_NCCL)
+}
+
+constexpr bool UseCUDA() {
+#if defined(XGBOOST_USE_CUDA)
+  return true;
+#else
+  return false;
+#endif  // defined(XGBOOST_USE_CUDA)
+}
+
+constexpr bool UseFederated() {
+#if defined(XGBOOST_USE_FEDERATED)
+  return true;
+#else
+  return false;
+#endif
+}
+
+auto MakeParamsForTest() {
+  std::vector<Param> cases;
+
+  auto push = [&](std::string name, auto fn) {
+    for (bool is_federated : {false, true}) {
+      for (DataSplitMode m : {DataSplitMode::kCol, DataSplitMode::kRow}) {
+        for (auto d : {DeviceOrd::CPU(), DeviceOrd::CUDA(0)}) {
+          if (!is_federated && !UseNCCL() && d.IsCUDA()) {
+            // Federated doesn't use nccl.
+            continue;
+          }
+          if (!UseCUDA() && d.IsCUDA()) {
+            // skip CUDA tests
+            continue;
+          }
+          if (!UseFederated() && is_federated) {
+            // skip GRPC tests
+            continue;
+          }
+
+          auto p = Param{true, is_federated, m, fn, name, d};
+          cases.push_back(p);
+          if (!is_federated) {
+            // Add a local test.
+            p.is_dist = false;
+            cases.push_back(p);
+          }
+        }
+      }
+    }
+  };
+
+#define REFLECT_NAME(name) push(#name, Verify##name)
+  // AUC
+  REFLECT_NAME(BinaryAUC);
+  REFLECT_NAME(MultiClassAUC);
+  REFLECT_NAME(RankingAUC);
+  REFLECT_NAME(PRAUC);
+  REFLECT_NAME(MultiClassPRAUC);
+  REFLECT_NAME(RankingPRAUC);
+  // Elementwise
+  REFLECT_NAME(RMSE);
+  REFLECT_NAME(RMSLE);
+  REFLECT_NAME(MAE);
+  REFLECT_NAME(MAPE);
+  REFLECT_NAME(MPHE);
+  REFLECT_NAME(LogLoss);
+  REFLECT_NAME(Error);
+  REFLECT_NAME(PoissonNegLogLik);
+  REFLECT_NAME(MultiRMSE);
+  REFLECT_NAME(Quantile);
+  // Multi-Class
+  REFLECT_NAME(MultiClassError);
+  REFLECT_NAME(MultiClassLogLoss);
+  // Ranking
+  REFLECT_NAME(Precision);
+  REFLECT_NAME(NDCG);
+  REFLECT_NAME(MAP);
+  REFLECT_NAME(NDCGExpGain);
+  // AFT
+  using namespace xgboost::common;  // NOLINT
+  REFLECT_NAME(AFTNegLogLik);
+  REFLECT_NAME(IntervalRegressionAccuracy);
+
+#undef REFLECT_NAME
+
+  return cases;
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    DistributedMetric, TestDistributedMetric, ::testing::ValuesIn(MakeParamsForTest()),
+    [](const ::testing::TestParamInfo<TestDistributedMetric::ParamType>& info) {
+      std::string result;
+      if (info.param.is_dist) {
+        result += "Dist_";
+      }
+      if (info.param.is_fed) {
+        result += "Federated_";
+      }
+      if (info.param.split == DataSplitMode::kRow) {
+        result += "RowSplit";
+      } else {
+        result += "ColSplit";
+      }
+      result += "_";
+      result += info.param.device.IsCPU() ? "CPU" : "CUDA";
+      result += "_";
+      result += info.param.name;
+      return result;
+    });
+}  // namespace xgboost::metric

tests/cpp/metric/test_elementwise_metric.cc

@@ -1,106 +0,0 @@
-/**
- * Copyright 2018-2023 by XGBoost contributors
- */
-#include "test_elementwise_metric.h"
-
-namespace xgboost::metric {
-TEST(Metric, DeclareUnifiedTest(RMSE)) { VerifyRMSE(); }
-
-TEST(Metric, DeclareUnifiedTest(RMSLE)) { VerifyRMSLE(); }
-
-TEST(Metric, DeclareUnifiedTest(MAE)) { VerifyMAE(); }
-
-TEST(Metric, DeclareUnifiedTest(MAPE)) { VerifyMAPE(); }
-
-TEST(Metric, DeclareUnifiedTest(MPHE)) { VerifyMPHE(); }
-
-TEST(Metric, DeclareUnifiedTest(LogLoss)) { VerifyLogLoss(); }
-
-TEST(Metric, DeclareUnifiedTest(Error)) { VerifyError(); }
-
-TEST(Metric, DeclareUnifiedTest(PoissonNegLogLik)) { VerifyPoissonNegLogLik(); }
-
-TEST(Metric, DeclareUnifiedTest(MultiRMSE)) { VerifyMultiRMSE(); }
-
-TEST(Metric, DeclareUnifiedTest(Quantile)) { VerifyQuantile(); }
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RMSERowSplit) {
-  DoTest(VerifyRMSE, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RMSEColumnSplit) {
-  DoTest(VerifyRMSE, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RMSLERowSplit) {
-  DoTest(VerifyRMSLE, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), RMSLEColumnSplit) {
-  DoTest(VerifyRMSLE, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MAERowSplit) {
-  DoTest(VerifyMAE, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MAEColumnSplit) {
-  DoTest(VerifyMAE, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MAPERowSplit) {
-  DoTest(VerifyMAPE, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MAPEColumnSplit) {
-  DoTest(VerifyMAPE, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MPHERowSplit) {
-  DoTest(VerifyMPHE, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MPHEColumnSplit) {
-  DoTest(VerifyMPHE, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), LogLossRowSplit) {
-  DoTest(VerifyLogLoss, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), LogLossColumnSplit) {
-  DoTest(VerifyLogLoss, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), ErrorRowSplit) {
-  DoTest(VerifyError, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), ErrorColumnSplit) {
-  DoTest(VerifyError, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), PoissonNegLogLikRowSplit) {
-  DoTest(VerifyPoissonNegLogLik, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), PoissonNegLogLikColumnSplit) {
-  DoTest(VerifyPoissonNegLogLik, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiRMSERowSplit) {
-  DoTest(VerifyMultiRMSE, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiRMSEColumnSplit) {
-  DoTest(VerifyMultiRMSE, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), QuantileRowSplit) {
-  DoTest(VerifyQuantile, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), QuantileColumnSplit) {
-  DoTest(VerifyQuantile, DataSplitMode::kCol);
-}
-}  // namespace xgboost::metric

tests/cpp/metric/test_elementwise_metric.cu

@@ -1,5 +0,0 @@
-/*!
- * Copyright 2018 XGBoost contributors
- */
-// Dummy file to keep the CUDA conditional compile trick.
-#include "test_elementwise_metric.cc"

tests/cpp/metric/test_elementwise_metric.h

@@ -42,8 +42,8 @@ inline void CheckDeterministicMetricElementWise(StringView name, int32_t device)
   }
 }
 
-inline void VerifyRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyRMSE(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("rmse", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "rmse");
@@ -68,11 +68,11 @@ inline void VerifyRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               0.6708f, 0.001f);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"rmse"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"rmse"}, device.ordinal);
 }
 
-inline void VerifyRMSLE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyRMSLE(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("rmsle", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "rmsle");
@@ -97,11 +97,11 @@ inline void VerifyRMSLE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               0.2415f, 1e-4);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"rmsle"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"rmsle"}, device.ordinal);
 }
 
-inline void VerifyMAE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMAE(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("mae", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mae");
@@ -126,11 +126,11 @@ inline void VerifyMAE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               0.54f, 0.001f);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"mae"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"mae"}, device.ordinal);
 }
 
-inline void VerifyMAPE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMAPE(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("mape", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mape");
@@ -155,11 +155,11 @@ inline void VerifyMAPE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               1.3250f, 0.001f);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"mape"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"mape"}, device.ordinal);
 }
 
-inline void VerifyMPHE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMPHE(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   std::unique_ptr<xgboost::Metric> metric{xgboost::Metric::Create("mphe", &ctx)};
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mphe");
@@ -183,7 +183,7 @@ inline void VerifyMPHE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
                             {  1,   2,   9,   8}, {}, data_split_mode),
               0.1922f, 1e-4);
 
-  CheckDeterministicMetricElementWise(StringView{"mphe"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"mphe"}, device.ordinal);
 
   metric->Configure({{"huber_slope", "0.1"}});
   EXPECT_NEAR(GetMetricEval(metric.get(),
@@ -193,8 +193,8 @@ inline void VerifyMPHE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               0.0461686f, 1e-4);
 }
 
-inline void VerifyLogLoss(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyLogLoss(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("logloss", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "logloss");
@@ -223,11 +223,11 @@ inline void VerifyLogLoss(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               1.3138f, 0.001f);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"logloss"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"logloss"}, device.ordinal);
 }
 
-inline void VerifyError(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyError(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("error", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "error");
@@ -285,11 +285,11 @@ inline void VerifyError(DataSplitMode data_split_mode = DataSplitMode::kRow) {
               0.45f, 0.001f);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"error@0.5"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"error@0.5"}, device.ordinal);
 }
 
-inline void VerifyPoissonNegLogLik(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyPoissonNegLogLik(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("poisson-nloglik", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "poisson-nloglik");
@@ -318,11 +318,11 @@ inline void VerifyPoissonNegLogLik(DataSplitMode data_split_mode = DataSplitMode
               1.5783f, 0.001f);
   delete metric;
 
-  CheckDeterministicMetricElementWise(StringView{"poisson-nloglik"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"poisson-nloglik"}, device.ordinal);
 }
 
-inline void VerifyMultiRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMultiRMSE(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   size_t n_samples = 32, n_targets = 8;
   linalg::Tensor<float, 2> y{{n_samples, n_targets}, ctx.Device()};
   auto &h_y = y.Data()->HostVector();
@@ -343,8 +343,8 @@ inline void VerifyMultiRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow)
   ASSERT_FLOAT_EQ(ret, loss_w);
 }
 
-inline void VerifyQuantile(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyQuantile(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   std::unique_ptr<Metric> metric{Metric::Create("quantile", &ctx)};
 
   HostDeviceVector<float> predts{0.1f, 0.9f, 0.1f, 0.9f};

tests/cpp/metric/test_multiclass_metric.cc

@@ -1,29 +0,0 @@
-// Copyright by Contributors
-#include "test_multiclass_metric.h"
-
-#include <string>
-
-namespace xgboost {
-namespace metric {
-
-TEST(Metric, DeclareUnifiedTest(MultiClassError)) { VerifyMultiClassError(); }
-
-TEST(Metric, DeclareUnifiedTest(MultiClassLogLoss)) { VerifyMultiClassLogLoss(); }
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassErrorRowSplit) {
-  DoTest(VerifyMultiClassError, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassErrorColumnSplit) {
-  DoTest(VerifyMultiClassError, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassLogLossRowSplit) {
-  DoTest(VerifyMultiClassLogLoss, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MultiClassLogLossColumnSplit) {
-  DoTest(VerifyMultiClassLogLoss, DataSplitMode::kCol);
-}
-}  // namespace metric
-}  // namespace xgboost

tests/cpp/metric/test_multiclass_metric.cu

@@ -1,5 +0,0 @@
-/*!
- * Copyright 2019 XGBoost contributors
- */
-// Dummy file to keep the CUDA conditional compile trick.
-#include "test_multiclass_metric.cc"

tests/cpp/metric/test_multiclass_metric.h

@@ -44,8 +44,8 @@ inline void CheckDeterministicMetricMultiClass(StringView name, int32_t device)
   }
 }
 
-inline void TestMultiClassError(int device, DataSplitMode data_split_mode) {
-  auto ctx = MakeCUDACtx(device);
+inline void TestMultiClassError(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("merror", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "merror");
@@ -59,13 +59,13 @@ inline void TestMultiClassError(int device, DataSplitMode data_split_mode) {
   delete metric;
 }
 
-inline void VerifyMultiClassError(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  TestMultiClassError(GPUIDX, data_split_mode);
-  CheckDeterministicMetricMultiClass(StringView{"merror"}, GPUIDX);
+inline void VerifyMultiClassError(DataSplitMode data_split_mode, DeviceOrd device) {
+  TestMultiClassError(data_split_mode, device);
+  CheckDeterministicMetricMultiClass(StringView{"merror"}, device.ordinal);
 }
 
-inline void TestMultiClassLogLoss(int device, DataSplitMode data_split_mode) {
-  auto ctx = MakeCUDACtx(device);
+inline void TestMultiClassLogLoss(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   xgboost::Metric * metric = xgboost::Metric::Create("mlogloss", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mlogloss");
@@ -80,9 +80,9 @@ inline void TestMultiClassLogLoss(int device, DataSplitMode data_split_mode) {
   delete metric;
 }
 
-inline void VerifyMultiClassLogLoss(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  TestMultiClassLogLoss(GPUIDX, data_split_mode);
-  CheckDeterministicMetricMultiClass(StringView{"mlogloss"}, GPUIDX);
+inline void VerifyMultiClassLogLoss(DataSplitMode data_split_mode, DeviceOrd device) {
+  TestMultiClassLogLoss(data_split_mode, device);
+  CheckDeterministicMetricMultiClass(StringView{"mlogloss"}, device.ordinal);
 }
 
 }  // namespace metric

tests/cpp/metric/test_rank_metric.cc

@@ -1,84 +1,29 @@
 /**
- * Copyright 2016-2023 by XGBoost Contributors
+ * Copyright 2016-2023, XGBoost Contributors
  */
-#include <gtest/gtest.h>                 // for Test, EXPECT_NEAR, ASSERT_STREQ
-#include <xgboost/context.h>             // for Context
-#include <xgboost/data.h>                // for MetaInfo, DMatrix
-#include <xgboost/linalg.h>              // for Matrix
-#include <xgboost/metric.h>              // for Metric
-
-#include <algorithm>                     // for max
-#include <memory>                        // for unique_ptr
-#include <vector>                        // for vector
-
 #include "test_rank_metric.h"
-#include "../helpers.h"                  // for GetMetricEval, CreateEmptyGe...
-#include "xgboost/base.h"                // for bst_float, kRtEps
-#include "xgboost/host_device_vector.h"  // for HostDeviceVector
-#include "xgboost/json.h"                // for Json, String, Object
 
-namespace xgboost {
-namespace metric {
+#include <gtest/gtest.h>      // for Test, EXPECT_NEAR, ASSERT_STREQ
+#include <xgboost/context.h>  // for Context
+#include <xgboost/metric.h>   // for Metric
 
-#if !defined(__CUDACC__)
+#include <memory>  // for unique_ptr
+
+#include "../helpers.h"    // for GetMetricEval, CreateEmptyGe...
+#include "xgboost/base.h"  // for bst_float, kRtEps
+
+namespace xgboost::metric {
 TEST(Metric, AMS) {
   auto ctx = MakeCUDACtx(GPUIDX);
   EXPECT_ANY_THROW(Metric::Create("ams", &ctx));
-  Metric* metric = Metric::Create("ams@0.5f", &ctx);
+  std::unique_ptr<Metric> metric{Metric::Create("ams@0.5f", &ctx)};
   ASSERT_STREQ(metric->Name(), "ams@0.5");
-  EXPECT_NEAR(GetMetricEval(metric, {0, 1}, {0, 1}), 0.311f, 0.001f);
-  EXPECT_NEAR(GetMetricEval(metric,
-                            {0.1f, 0.9f, 0.1f, 0.9f},
-                            {  0,   0,   1,   1}),
-              0.29710f, 0.001f);
+  EXPECT_NEAR(GetMetricEval(metric.get(), {0, 1}, {0, 1}), 0.311f, 0.001f);
+  EXPECT_NEAR(GetMetricEval(metric.get(), {0.1f, 0.9f, 0.1f, 0.9f}, {0, 0, 1, 1}), 0.29710f,
+              0.001f);
 
-  delete metric;
-  metric = Metric::Create("ams@0", &ctx);
+  metric.reset(Metric::Create("ams@0", &ctx));
   ASSERT_STREQ(metric->Name(), "ams@0");
-  EXPECT_NEAR(GetMetricEval(metric, {0, 1}, {0, 1}), 0.311f, 0.001f);
-
-  delete metric;
+  EXPECT_NEAR(GetMetricEval(metric.get(), {0, 1}, {0, 1}), 0.311f, 0.001f);
 }
-#endif
-
-TEST(Metric, DeclareUnifiedTest(Precision)) { VerifyPrecision(); }
-
-TEST(Metric, DeclareUnifiedTest(NDCG)) { VerifyNDCG(); }
-
-TEST(Metric, DeclareUnifiedTest(MAP)) { VerifyMAP(); }
-
-TEST(Metric, DeclareUnifiedTest(NDCGExpGain)) { VerifyNDCGExpGain(); }
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), PrecisionRowSplit) {
-  DoTest(VerifyPrecision, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), PrecisionColumnSplit) {
-  DoTest(VerifyPrecision, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), NDCGRowSplit) {
-  DoTest(VerifyNDCG, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), NDCGColumnSplit) {
-  DoTest(VerifyNDCG, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MAPRowSplit) {
-  DoTest(VerifyMAP, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), MAPColumnSplit) {
-  DoTest(VerifyMAP, DataSplitMode::kCol);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), NDCGExpGainRowSplit) {
-  DoTest(VerifyNDCGExpGain, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), NDCGExpGainColumnSplit) {
-  DoTest(VerifyNDCGExpGain, DataSplitMode::kCol);
-}
-}  // namespace metric
-}  // namespace xgboost
+}  // namespace xgboost::metric

tests/cpp/metric/test_rank_metric.cu

@@ -1,5 +0,0 @@
-/*!
- * Copyright 2019 XGBoost contributors
- */
-// Dummy file to keep the CUDA conditional compile trick.
-#include "test_rank_metric.cc"

tests/cpp/metric/test_rank_metric.h

@@ -19,8 +19,8 @@
 
 namespace xgboost::metric {
 
-inline void VerifyPrecision(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyPrecision(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   std::unique_ptr<xgboost::Metric> metric{Metric::Create("pre", &ctx)};
   ASSERT_STREQ(metric->Name(), "pre");
   EXPECT_NEAR(GetMetricEval(metric.get(), {0, 1}, {0, 1}, {}, {}, data_split_mode), 0.5, 1e-7);
@@ -43,8 +43,8 @@ inline void VerifyPrecision(DataSplitMode data_split_mode = DataSplitMode::kRow)
               0.5f, 1e-7);
 }
 
-inline void VerifyNDCG(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyNDCG(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   Metric * metric = xgboost::Metric::Create("ndcg", &ctx);
   ASSERT_STREQ(metric->Name(), "ndcg");
   EXPECT_ANY_THROW(GetMetricEval(metric, {0, 1}, {}, {}, {}, data_split_mode));
@@ -101,8 +101,8 @@ inline void VerifyNDCG(DataSplitMode data_split_mode = DataSplitMode::kRow) {
   delete metric;
 }
 
-inline void VerifyMAP(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyMAP(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
   Metric * metric = xgboost::Metric::Create("map", &ctx);
   ASSERT_STREQ(metric->Name(), "map");
   EXPECT_NEAR(GetMetricEval(metric, {0, 1}, {0, 1}, {}, {}, data_split_mode), 1, kRtEps);
@@ -149,8 +149,8 @@ inline void VerifyMAP(DataSplitMode data_split_mode = DataSplitMode::kRow) {
   delete metric;
 }
 
-inline void VerifyNDCGExpGain(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  Context ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyNDCGExpGain(DataSplitMode data_split_mode, DeviceOrd device) {
+  Context ctx = MakeCUDACtx(device.ordinal);
 
   auto p_fmat = xgboost::RandomDataGenerator{0, 0, 0}.GenerateDMatrix();
   MetaInfo& info = p_fmat->Info();

tests/cpp/metric/test_survival_metric.cc

@@ -1,5 +1,5 @@
-/*!
- * Copyright (c) by Contributors 2020
+/**
+ * Copyright 2020-2023, XGBoost Contributors
  */
 #include <gtest/gtest.h>
 #include <memory>
@@ -16,8 +16,7 @@
 // CUDA conditional compile trick.
 #include "test_survival_metric.cu"
 
-namespace xgboost {
-namespace common {
+namespace xgboost::common {
 
 /** Tests for Survival metrics that should run only on CPU **/
 
@@ -113,6 +112,4 @@ TEST(AFTLoss, IntervalCensored) {
     { 8.0000, 4.8004, 2.8805, 1.7284, 1.0372, 0.6231, 0.3872, 0.3031, 0.3740, 0.5839, 0.8995,
       1.2878, 1.7231, 2.1878, 2.6707, 3.1647, 3.6653, 4.1699, 4.6770, 5.1856 });
 }
-
-}  // namespace common
-}  // namespace xgboost
+}  // namespace xgboost::common

tests/cpp/metric/test_survival_metric.cu

@@ -7,28 +7,7 @@
 
 /** Tests for Survival metrics that should run both on CPU and GPU **/
 
-namespace xgboost {
-namespace common {
-TEST(Metric, DeclareUnifiedTest(AFTNegLogLik)) { VerifyAFTNegLogLik(); }
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), AFTNegLogLikRowSplit) {
-  DoTest(VerifyAFTNegLogLik, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), AFTNegLogLikColumnSplit) {
-  DoTest(VerifyAFTNegLogLik, DataSplitMode::kCol);
-}
-
-TEST(Metric, DeclareUnifiedTest(IntervalRegressionAccuracy)) { VerifyIntervalRegressionAccuracy(); }
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), IntervalRegressionAccuracyRowSplit) {
-  DoTest(VerifyIntervalRegressionAccuracy, DataSplitMode::kRow);
-}
-
-TEST_F(DeclareUnifiedDistributedTest(MetricTest), IntervalRegressionAccuracyColumnSplit) {
-  DoTest(VerifyIntervalRegressionAccuracy, DataSplitMode::kCol);
-}
-
+namespace xgboost::common {
 // Test configuration of AFT metric
 TEST(AFTNegLogLikMetric, DeclareUnifiedTest(Configuration)) {
   auto ctx = MakeCUDACtx(GPUIDX);
@@ -44,5 +23,4 @@ TEST(AFTNegLogLikMetric, DeclareUnifiedTest(Configuration)) {
 
   CheckDeterministicMetricElementWise(StringView{"aft-nloglik"}, GPUIDX);
 }
-}  // namespace common
-}  // namespace xgboost
+}  // namespace xgboost::common

tests/cpp/metric/test_survival_metric.h

@@ -47,8 +47,8 @@ inline void CheckDeterministicMetricElementWise(StringView name, int32_t device)
   }
 }
 
-inline void VerifyAFTNegLogLik(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyAFTNegLogLik(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
 
   /**
    * Test aggregate output from the AFT metric over a small test data set.
@@ -78,8 +78,8 @@ inline void VerifyAFTNegLogLik(DataSplitMode data_split_mode = DataSplitMode::kR
   }
 }
 
-inline void VerifyIntervalRegressionAccuracy(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = MakeCUDACtx(GPUIDX);
+inline void VerifyIntervalRegressionAccuracy(DataSplitMode data_split_mode, DeviceOrd device) {
+  auto ctx = MakeCUDACtx(device.ordinal);
 
   auto p_fmat = EmptyDMatrix();
   MetaInfo& info = p_fmat->Info();
@@ -101,7 +101,7 @@ inline void VerifyIntervalRegressionAccuracy(DataSplitMode data_split_mode = Dat
   info.labels_lower_bound_.HostVector()[0] = 70.0f;
   EXPECT_FLOAT_EQ(metric->Evaluate(preds, p_fmat), 0.25f);
 
-  CheckDeterministicMetricElementWise(StringView{"interval-regression-accuracy"}, GPUIDX);
+  CheckDeterministicMetricElementWise(StringView{"interval-regression-accuracy"}, device.ordinal);
 }
 }  // namespace common
 }  // namespace xgboost