diff --git a/tests/cpp/common/test_quantile.cc b/tests/cpp/common/test_quantile.cc index 1759c6abeb13..dade961950e4 100644 --- a/tests/cpp/common/test_quantile.cc +++ b/tests/cpp/common/test_quantile.cc @@ -13,7 +13,7 @@ #include "../../../src/data/simple_dmatrix.h" // SimpleDMatrix #include "../collective/test_worker.h" // for TestDistributedGlobal #if defined(XGBOOST_USE_FEDERATED) -#include "../plugin/federated/test_worker.h" +#include "../plugin/federated/test_worker.h" // for TestEncryptedGlobal #endif // defined(XGBOOST_USE_FEDERATED) #include "xgboost/context.h" @@ -314,6 +314,7 @@ void DoTestColSplitQuantileSecure() { Context ctx; auto const world = collective::GetWorldSize(); auto const rank = collective::GetRank(); + ASSERT_TRUE(collective::IsEncrypted()); size_t cols = 2; size_t rows = 3; @@ -395,7 +396,7 @@ void DoTestColSplitQuantileSecure() { template void TestColSplitQuantileSecure() { auto constexpr kWorkers = 2; - collective::TestFederatedGlobal(kWorkers, [&] { DoTestColSplitQuantileSecure(); }); + collective::TestEncryptedGlobal(kWorkers, [&] { DoTestColSplitQuantileSecure(); }); } } // anonymous namespace diff --git a/tests/cpp/plugin/federated/test_worker.h b/tests/cpp/plugin/federated/test_worker.h index 8ec76237df5c..682cf021b12c 100644 --- a/tests/cpp/plugin/federated/test_worker.h +++ b/tests/cpp/plugin/federated/test_worker.h @@ -83,4 +83,16 @@ void TestFederatedGlobal(std::int32_t n_workers, WorkerFn&& fn) { collective::Finalize(); }); } + +template +void TestEncryptedGlobal(std::int32_t n_workers, WorkerFn&& fn) { + TestFederatedImpl(n_workers, [&](std::int32_t port, std::int32_t i) { + auto config = FederatedTestConfig(n_workers, port, i); + config["federated_plugin"] = Object{}; + config["federated_plugin"]["name"] = String{"mock"}; + collective::Init(config); + fn(); + collective::Finalize(); + }); +} } // namespace xgboost::collective diff --git a/tests/cpp/plugin/test_federated_learner.cc b/tests/cpp/plugin/test_federated_learner.cc index 10aa8fadbaea..4d150e27311f 100644 --- a/tests/cpp/plugin/test_federated_learner.cc +++ b/tests/cpp/plugin/test_federated_learner.cc @@ -51,21 +51,8 @@ void VerifyObjective(std::size_t rows, std::size_t cols, float expected_base_sco auto model = MakeModel(tree_method, device, objective, sliced); auto base_score = GetBaseScore(model); - - std::unique_ptr expected{Learner::Create({})}; - expected->LoadModel(expected_model); - - std::unique_ptr got{Learner::Create({})}; - got->LoadModel(model); - - if (rank == 0) { - ASSERT_EQ(base_score, expected_base_score) << " rank " << rank; - HostDeviceVector expected_predt; - expected->Predict(dmat, false, &expected_predt, 0, 0); - HostDeviceVector got_predt; - expected->Predict(dmat, false, &got_predt, 0, 0); - ASSERT_EQ(expected_predt.HostVector(), got_predt.HostVector()); - } + ASSERT_EQ(base_score, expected_base_score) << " rank " << rank; + ASSERT_EQ(model, expected_model) << " rank " << rank; } } // namespace @@ -74,19 +61,6 @@ class VerticalFederatedLearnerTest : public ::testing::TestWithParamRun("approx", DeviceSym::CUDA(), objective); } TEST_P(VerticalFederatedLearnerTest, GPUHist) { std::string objective = GetParam(); - // Not yet supported by the plugin system - GTEST_SKIP(); this->Run("hist", DeviceSym::CUDA(), objective); } #endif // defined(XGBOOST_USE_CUDA) diff --git a/tests/cpp/test_learner.cc b/tests/cpp/test_learner.cc index 67db2ee26775..976ae2147a06 100644 --- a/tests/cpp/test_learner.cc +++ b/tests/cpp/test_learner.cc @@ -756,9 +756,6 @@ void TestColumnSplitWithArgs(std::string const& tree_method, bool use_gpu, Args auto model = GetModelWithArgs(p_fmat, tree_method, device, args); auto world_size{3}; - if (federated && use_gpu) { - GTEST_SKIP(); - } if (use_gpu) { world_size = common::AllVisibleGPUs(); // Simulate MPU on a single GPU. Federated doesn't use nccl, can run multiple @@ -791,26 +788,17 @@ class ColumnSplitTrainingTest public: static void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu, bool federated) { - if (federated) { - GTEST_SKIP(); - } Args args{ {"colsample_bytree", "0.5"}, {"colsample_bylevel", "0.6"}, {"colsample_bynode", "0.7"}}; TestColumnSplitWithArgs(tree_method, use_gpu, args, federated); } static void TestColumnSplitInteractionConstraints(std::string const& tree_method, bool use_gpu, bool federated) { - if (federated) { - GTEST_SKIP(); - } Args args{{"interaction_constraints", "[[0, 5, 7], [2, 8, 9], [1, 3, 6]]"}}; TestColumnSplitWithArgs(tree_method, use_gpu, args, federated); } static void TestColumnSplitMonotoneConstraints(std::string const& tree_method, bool use_gpu, bool federated) { - if (federated) { - GTEST_SKIP(); - } Args args{{"monotone_constraints", "(1,-1,0,1,1,-1,-1,0,0,1)"}}; TestColumnSplitWithArgs(tree_method, use_gpu, args, federated); }