[coll] Implement a new tracker and a communicator.

R-package/src/Makevars.in

@@ -98,9 +98,13 @@ OBJECTS= \
     $(PKGROOT)/src/context.o \
     $(PKGROOT)/src/logging.o \
     $(PKGROOT)/src/global_config.o \
+    $(PKGROOT)/src/collective/allgather.o \
+    $(PKGROOT)/src/collective/comm.o \
+    $(PKGROOT)/src/collective/tracker.o \
     $(PKGROOT)/src/collective/communicator.o \
     $(PKGROOT)/src/collective/in_memory_communicator.o \
     $(PKGROOT)/src/collective/in_memory_handler.o \
+    $(PKGROOT)/src/collective/loop.o \
     $(PKGROOT)/src/collective/socket.o \
     $(PKGROOT)/src/common/charconv.o \
     $(PKGROOT)/src/common/column_matrix.o \

R-package/src/Makevars.win

@@ -98,9 +98,13 @@ OBJECTS= \
     $(PKGROOT)/src/context.o \
     $(PKGROOT)/src/logging.o \
     $(PKGROOT)/src/global_config.o \
+    $(PKGROOT)/src/collective/allgather.o \
+    $(PKGROOT)/src/collective/comm.o \
+    $(PKGROOT)/src/collective/tracker.o \
     $(PKGROOT)/src/collective/communicator.o \
     $(PKGROOT)/src/collective/in_memory_communicator.o \
     $(PKGROOT)/src/collective/in_memory_handler.o \
+    $(PKGROOT)/src/collective/loop.o \
     $(PKGROOT)/src/collective/socket.o \
     $(PKGROOT)/src/common/charconv.o \
     $(PKGROOT)/src/common/column_matrix.o \

include/xgboost/collective/result.h

@@ -157,4 +157,13 @@ struct Result {
 [[nodiscard]] inline auto Fail(std::string msg, std::error_code errc, Result&& prev) {
   return Result{std::move(msg), std::move(errc), std::forward<Result>(prev)};
 }
+
+// We don't have monad, a simple helper would do.
+template <typename Fn>
+Result operator<<(Result&& r, Fn&& fn) {
+  if (!r.OK()) {
+    return std::forward<Result>(r);
+  }
+  return fn();
+}
 }  // namespace xgboost::collective

include/xgboost/collective/socket.h

@@ -380,11 +380,18 @@ class TCPSocket {
   }
   [[nodiscard]] bool NonBlocking() const { return non_blocking_; }
   [[nodiscard]] Result RecvTimeout(std::chrono::seconds timeout) {
-    timeval tv;
+    // https://stackoverflow.com/questions/2876024/linux-is-there-a-read-or-recv-from-socket-with-timeout
+#if defined(_WIN32)
+    DWORD tv = timeout.count() * 1000;
+    auto rc =
+        setsockopt(Handle(), SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<char *>(&tv), sizeof(tv));
+#else
+    struct timeval tv;
     tv.tv_sec = timeout.count();
     tv.tv_usec = 0;
     auto rc = setsockopt(Handle(), SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<char const *>(&tv),
                          sizeof(tv));
+#endif
     if (rc != 0) {
       return system::FailWithCode("Failed to set timeout on recv.");
     }
@@ -425,7 +432,12 @@ class TCPSocket {
    */
   TCPSocket Accept() {
     HandleT newfd = accept(Handle(), nullptr, nullptr);
-    if (newfd == InvalidSocket()) {
+#if defined(_WIN32)
+    auto interrupt = WSAEINTR;
+#else
+    auto interrupt = EINTR;
+#endif
+    if (newfd == InvalidSocket() && system::LastError() != interrupt) {
       system::ThrowAtError("accept");
     }
     TCPSocket newsock{newfd};
@@ -468,7 +480,7 @@ class TCPSocket {
   /**
    * \brief Bind socket to INADDR_ANY, return the port selected by the OS.
    */
-  in_port_t BindHost() {
+  [[nodiscard]] in_port_t BindHost() {
     if (Domain() == SockDomain::kV6) {
       auto addr = SockAddrV6::InaddrAny();
       auto handle = reinterpret_cast<sockaddr const *>(&addr.Handle());
@@ -539,7 +551,7 @@ class TCPSocket {
   /**
    * \brief Send data, without error then all data should be sent.
    */
-  auto SendAll(void const *buf, std::size_t len) {
+  [[nodiscard]] auto SendAll(void const *buf, std::size_t len) {
     char const *_buf = reinterpret_cast<const char *>(buf);
     std::size_t ndone = 0;
     while (ndone < len) {
@@ -558,7 +570,7 @@ class TCPSocket {
   /**
    * \brief Receive data, without error then all data should be received.
    */
-  auto RecvAll(void *buf, std::size_t len) {
+  [[nodiscard]] auto RecvAll(void *buf, std::size_t len) {
     char *_buf = reinterpret_cast<char *>(buf);
     std::size_t ndone = 0;
     while (ndone < len) {
@@ -612,7 +624,15 @@ class TCPSocket {
    */
   void Close() {
     if (InvalidSocket() != handle_) {
+#if defined(_WIN32)
+      auto rc = system::CloseSocket(handle_);
+      // it's possible that we close TCP sockets after finalizing WSA due to detached thread.
+      if (rc != 0 && system::LastError() != WSANOTINITIALISED) {
+        system::ThrowAtError("close", rc);
+      }
+#else
       xgboost_CHECK_SYS_CALL(system::CloseSocket(handle_), 0);
+#endif
       handle_ = InvalidSocket();
     }
   }
@@ -634,6 +654,24 @@ class TCPSocket {
     socket.domain_ = domain;
 #endif  // defined(__APPLE__)
     return socket;
+#endif  // defined(xgboost_IS_MINGW)
+  }
+
+  static TCPSocket *CreatePtr(SockDomain domain) {
+#if defined(xgboost_IS_MINGW)
+    MingWError();
+    return nullptr;
+#else
+    auto fd = socket(static_cast<std::int32_t>(domain), SOCK_STREAM, 0);
+    if (fd == InvalidSocket()) {
+      system::ThrowAtError("socket");
+    }
+    auto socket = new TCPSocket{fd};
+
+#if defined(__APPLE__)
+    socket->domain_ = domain;
+#endif  // defined(__APPLE__)
+    return socket;
 #endif  // defined(xgboost_IS_MINGW)
   }
 };

src/collective/allgather.cc

@@ -0,0 +1,42 @@
+/**
+ * Copyright 2023, XGBoost Contributors
+ */
+#include "allgather.h"
+
+#include <algorithm>  // for min
+#include <cstddef>    // for size_t
+#include <cstdint>    // for int8_t
+#include <memory>     // for shared_ptr
+
+#include "comm.h"          // for Comm, Channel
+#include "xgboost/span.h"  // for Span
+
+namespace xgboost::collective::cpu_impl {
+Result RingAllgather(Comm const& comm, common::Span<std::int8_t> data, std::size_t segment_size,
+                     std::int32_t worker_off, std::shared_ptr<Channel> prev_ch,
+                     std::shared_ptr<Channel> next_ch) {
+  auto world = comm.World();
+  auto rank = comm.Rank();
+  CHECK_LT(worker_off, world);
+
+  for (std::int32_t r = 0; r < world; ++r) {
+    auto send_rank = (rank + world - r + worker_off) % world;
+    auto send_off = send_rank * segment_size;
+    send_off = std::min(send_off, data.size_bytes());
+    auto send_seg = data.subspan(send_off, std::min(segment_size, data.size_bytes() - send_off));
+    next_ch->SendAll(send_seg.data(), send_seg.size_bytes());
+
+    auto recv_rank = (rank + world - r - 1 + worker_off) % world;
+    auto recv_off = recv_rank * segment_size;
+    recv_off = std::min(recv_off, data.size_bytes());
+    auto recv_seg = data.subspan(recv_off, std::min(segment_size, data.size_bytes() - recv_off));
+    prev_ch->RecvAll(recv_seg.data(), recv_seg.size_bytes());
+    auto rc = prev_ch->Block();
+    if (!rc.OK()) {
+      return rc;
+    }
+  }
+
+  return Success();
+}
+}  // namespace xgboost::collective::cpu_impl

src/collective/allgather.h

@@ -0,0 +1,23 @@
+/**
+ * Copyright 2023, XGBoost Contributors
+ */
+#pragma once
+#include <cstddef>  // for size_t
+#include <cstdint>  // for int32_t
+#include <memory>   // for shared_ptr
+
+#include "comm.h"          // for Comm, Channel
+#include "xgboost/span.h"  // for Span
+
+namespace xgboost::collective {
+namespace cpu_impl {
+/**
+ * @param worker_off Segment offset. For example, if the rank 2 worker specifis worker_off
+ *                   = 1, then it owns the third segment.
+ */
+[[nodiscard]] Result RingAllgather(Comm const& comm, common::Span<std::int8_t> data,
+                                   std::size_t segment_size, std::int32_t worker_off,
+                                   std::shared_ptr<Channel> prev_ch,
+                                   std::shared_ptr<Channel> next_ch);
+}  // namespace cpu_impl
+}  // namespace xgboost::collective