Single worker.

dmlc · Sep 23, 2024 · 4ec9b0e · 4ec9b0e
1 parent b2ca2fa
commit 4ec9b0e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/src/collective/comm.cc b/src/collective/comm.cc
@@ -146,9 +146,12 @@ Result ConnectTrackerImpl(proto::PeerInfo info, std::chrono::seconds timeout, st
   workers.resize(comm.World());
   workers[BootstrapNext(comm.Rank(), comm.World())] = next;
   if (BootstrapNext(comm.Rank(), comm.World()) == BootstrapPrev(comm.Rank(), comm.World())) {
-    CHECK_EQ(comm.World(), 2);
     if (comm.Rank() == 0) {
-      workers[BootstrapNext(comm.Rank(), comm.World())] = prev;
+      if (comm.World() == 2) {
+        workers[BootstrapNext(comm.Rank(), comm.World())] = prev;
+      } else {
+        CHECK_EQ(comm.World(), 1);
+      }
     }
   } else {
     workers[BootstrapPrev(comm.Rank(), comm.World())] = prev;

diff --git a/tests/cpp/collective/test_broadcast.cc b/tests/cpp/collective/test_broadcast.cc
@@ -40,11 +40,24 @@ class BroadcastTest : public SocketTest {};
 }  // namespace
 
 TEST_F(BroadcastTest, Basic) {
-  std::int32_t n_workers = std::min(2u, std::thread::hardware_concurrency());
-  TestDistributed(n_workers, [=](std::string host, std::int32_t port, std::chrono::seconds timeout,
-                                 std::int32_t r) {
-    Worker worker{host, port, timeout, n_workers, r};
-    worker.Run();
-  });
+  auto test_with = [](std::int32_t n_workers) {
+    TestDistributed(n_workers, [=](std::string host, std::int32_t port,
+                                   std::chrono::seconds timeout, std::int32_t r) {
+      Worker worker{host, port, timeout, n_workers, r};
+      worker.Run();
+    });
+  };
+  {
+    std::int32_t n_workers = 1u;
+    test_with(n_workers);
+  }
+  {
+    std::int32_t n_workers = std::min(2u, std::thread::hardware_concurrency());
+    test_with(n_workers);
+  }
+  {
+    std::int32_t n_workers = std::min(3u, std::thread::hardware_concurrency());
+    test_with(n_workers);
+  }
 }
 }  // namespace xgboost::collective