From d6d14d0fb9490e3910015c8f9ef30a39db6757c7 Mon Sep 17 00:00:00 2001
From: Rong Ou <rong.ou@gmail.com>
Date: Tue, 26 Sep 2023 17:27:43 -0700
Subject: [PATCH] Integration tests for interaction constraints with
 column-wise data split (#9611)

---
 tests/cpp/test_learner.cc | 99 +++++++++++++++++++++++----------------
 1 file changed, 59 insertions(+), 40 deletions(-)

diff --git a/tests/cpp/test_learner.cc b/tests/cpp/test_learner.cc
index 992716db5..0526222f3 100644
--- a/tests/cpp/test_learner.cc
+++ b/tests/cpp/test_learner.cc
@@ -720,48 +720,39 @@ INSTANTIATE_TEST_SUITE_P(ColumnSplitObjective, TestColumnSplit,
                          });
 
 namespace {
-void VerifyColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu,
-                                    Json const& expected_model) {
+Json GetModelWithArgs(std::shared_ptr<DMatrix> dmat, std::string const& tree_method,
+                      std::string const& device, Args const& args) {
+  std::unique_ptr<Learner> learner{Learner::Create({dmat})};
+  learner->SetParam("tree_method", tree_method);
+  learner->SetParam("device", device);
+  learner->SetParam("objective", "reg:logistic");
+  learner->SetParams(args);
+  learner->UpdateOneIter(0, dmat);
   Json model{Object{}};
-  {
-    auto const world_size = collective::GetWorldSize();
-    auto const rank = collective::GetRank();
-    auto const objective = "reg:logistic";
-    auto p_fmat = MakeFmatForObjTest(objective);
-    std::shared_ptr<DMatrix> sliced{p_fmat->SliceCol(world_size, rank)};
-    std::unique_ptr<Learner> learner{Learner::Create({sliced})};
-    learner->SetParam("tree_method", tree_method);
-    if (use_gpu) {
-      auto gpu_id = common::AllVisibleGPUs() == 1 ? 0 : rank;
-      learner->SetParam("device", "cuda:" + std::to_string(gpu_id));
-    }
-    learner->SetParam("objective", objective);
-    learner->SetParam("colsample_bytree", "0.5");
-    learner->SetParam("colsample_bylevel", "0.6");
-    learner->SetParam("colsample_bynode", "0.7");
-    learner->UpdateOneIter(0, sliced);
-    learner->SaveModel(&model);
+  learner->SaveModel(&model);
+  return model;
+}
+
+void VerifyColumnSplitWithArgs(std::string const& tree_method, bool use_gpu, Args const& args,
+                               Json const& expected_model) {
+  auto const world_size = collective::GetWorldSize();
+  auto const rank = collective::GetRank();
+  auto p_fmat = MakeFmatForObjTest("");
+  std::shared_ptr<DMatrix> sliced{p_fmat->SliceCol(world_size, rank)};
+  std::string device = "cpu";
+  if (use_gpu) {
+    auto gpu_id = common::AllVisibleGPUs() == 1 ? 0 : rank;
+    device = "cuda:" + std::to_string(gpu_id);
   }
+  auto model = GetModelWithArgs(sliced, tree_method, device, args);
   ASSERT_EQ(model, expected_model);
 }
 
-void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu) {
-  Json model{Object{}};
-  {
-    auto objective = "reg:logistic";
-    auto p_fmat = MakeFmatForObjTest(objective);
-    std::unique_ptr<Learner> learner{Learner::Create({p_fmat})};
-    learner->SetParam("tree_method", tree_method);
-    if (use_gpu) {
-      learner->SetParam("device", "cuda:0");
-    }
-    learner->SetParam("objective", objective);
-    learner->SetParam("colsample_bytree", "0.5");
-    learner->SetParam("colsample_bylevel", "0.6");
-    learner->SetParam("colsample_bynode", "0.7");
-    learner->UpdateOneIter(0, p_fmat);
-    learner->SaveModel(&model);
-  }
+void TestColumnSplitWithArgs(std::string const& tree_method, bool use_gpu, Args const& args) {
+  auto p_fmat = MakeFmatForObjTest("");
+  std::string device = use_gpu ? "cuda:0" : "cpu";
+  auto model = GetModelWithArgs(p_fmat, tree_method, device, args);
+
   auto world_size{3};
   if (use_gpu) {
     world_size = common::AllVisibleGPUs();
@@ -770,9 +761,19 @@ void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu)
       world_size = 3;
     }
   }
-  RunWithInMemoryCommunicator(world_size, VerifyColumnSplitColumnSampler, tree_method, use_gpu,
+  RunWithInMemoryCommunicator(world_size, VerifyColumnSplitWithArgs, tree_method, use_gpu, args,
                               model);
 }
+
+void TestColumnSplitColumnSampler(std::string const& tree_method, bool use_gpu) {
+  Args args{{"colsample_bytree", "0.5"}, {"colsample_bylevel", "0.6"}, {"colsample_bynode", "0.7"}};
+  TestColumnSplitWithArgs(tree_method, use_gpu, args);
+}
+
+void TestColumnSplitInteractionConstraints(std::string const& tree_method, bool use_gpu) {
+  Args args{{"interaction_constraints", "[[0, 5, 7], [2, 8, 9], [1, 3, 6]]"}};
+  TestColumnSplitWithArgs(tree_method, use_gpu, args);
+}
 }  // anonymous namespace
 
 TEST(ColumnSplitColumnSampler, Approx) { TestColumnSplitColumnSampler("approx", false); }
@@ -780,8 +781,26 @@ TEST(ColumnSplitColumnSampler, Approx) { TestColumnSplitColumnSampler("approx",
 TEST(ColumnSplitColumnSampler, Hist) { TestColumnSplitColumnSampler("hist", false); }
 
 #if defined(XGBOOST_USE_CUDA)
-TEST(ColumnSplitColumnSampler, GPUApprox) { TestColumnSplitColumnSampler("approx", true); }
+TEST(MGPUColumnSplitColumnSampler, GPUApprox) { TestColumnSplitColumnSampler("approx", true); }
 
-TEST(ColumnSplitColumnSampler, GPUHist) { TestColumnSplitColumnSampler("hist", true); }
+TEST(MGPUColumnSplitColumnSampler, GPUHist) { TestColumnSplitColumnSampler("hist", true); }
+#endif  // defined(XGBOOST_USE_CUDA)
+
+TEST(ColumnSplitInteractionConstraints, Approx) {
+  TestColumnSplitInteractionConstraints("approx", false);
+}
+
+TEST(ColumnSplitInteractionConstraints, Hist) {
+  TestColumnSplitInteractionConstraints("hist", false);
+}
+
+#if defined(XGBOOST_USE_CUDA)
+TEST(MGPUColumnSplitInteractionConstraints, GPUApprox) {
+  TestColumnSplitInteractionConstraints("approx", true);
+}
+
+TEST(MGPUColumnSplitInteractionConstraints, GPUHist) {
+  TestColumnSplitInteractionConstraints("hist", true);
+}
 #endif  // defined(XGBOOST_USE_CUDA)
 }  // namespace xgboost