Common interface for collective communication

amalgamation/xgboost-all0.cc

@@ -71,6 +71,9 @@
 #include "../src/logging.cc"
 #include "../src/global_config.cc"
 
+// collective
+#include "../src/collective/communicator.cc"
+
 // common
 #include "../src/common/common.cc"
 #include "../src/common/column_matrix.cc"

include/xgboost/c_api.h

@@ -9,10 +9,12 @@
 
 #ifdef __cplusplus
 #define XGB_EXTERN_C extern "C"
+#include <cstddef>
 #include <cstdio>
 #include <cstdint>
 #else
 #define XGB_EXTERN_C
+#include <stddef.h>
 #include <stdio.h>
 #include <stdint.h>
 #endif  // __cplusplus
@@ -1386,4 +1388,118 @@ XGB_DLL int XGBoosterFeatureScore(BoosterHandle handle, const char *json_config,
                                   bst_ulong *out_dim,
                                   bst_ulong const **out_shape,
                                   float const **out_scores);
+
+/*!
+ * \brief Initialize the collective communicator.
+ *
+ *  Call this once before using anything.
+ *
+ *  The additional configuration is not required. Usually the communicator will detect settings
+ *  from environment variables.
+ *
+ * \param json_config JSON encoded configuration. Accepted JSON keys are:
+ *   - xgboost_communicator: The type of the communicator. Can be set as an environment variable.
+ *     * rabit: Use Rabit. This is the default if the type is unspecified.
+ *     * mpi: Use MPI.
+ *     * federated: Use the gRPC interface for Federated Learning.
+ * Only applicable to the Rabit communicator (these are case-sensitive):
+ *   - rabit_tracker_uri: Hostname of the tracker.
+ *   - rabit_tracker_port: Port number of the tracker.
+ *   - rabit_task_id: ID of the current task, can be used to obtain deterministic rank assignment.
+ *   - rabit_world_size: Total number of workers.
+ *   - rabit_hadoop_mode: Enable Hadoop support.
+ *   - rabit_tree_reduce_minsize: Minimal size for tree reduce.
+ *   - rabit_reduce_ring_mincount: Minimal count to perform ring reduce.
+ *   - rabit_reduce_buffer: Size of the reduce buffer.
+ *   - rabit_bootstrap_cache: Size of the bootstrap cache.
+ *   - rabit_debug: Enable debugging.
+ *   - rabit_timeout: Enable timeout.
+ *   - rabit_timeout_sec: Timeout in seconds.
+ *   - rabit_enable_tcp_no_delay: Enable TCP no delay on Unix platforms.
+ * Only applicable to the Rabit communicator (these are case-sensitive, and can be set as
+ * environment variables):
+ *   - DMLC_TRACKER_URI: Hostname of the tracker.
+ *   - DMLC_TRACKER_PORT: Port number of the tracker.
+ *   - DMLC_TASK_ID: ID of the current task, can be used to obtain deterministic rank assignment.
+ *   - DMLC_ROLE: Role of the current task, "worker" or "server".
+ *   - DMLC_NUM_ATTEMPT: Number of attempts after task failure.
+ *   - DMLC_WORKER_CONNECT_RETRY: Number of retries to connect to the tracker.
+ * Only applicable to the Federated communicator (these are not case-sensitive):
+ *   - federated_server_address: Address of the federated server.
+ *   - federated_world_size: Number of federated workers.
+ *   - federated_rank: Rank of the current worker.
+ *   - federated_server_cert: Server certificate file path. Only needed for the SSL mode.
+ *   - federated_client_key: Client key file path. Only needed for the SSL mode.
+ *   - federated_client_cert: Client certificate file path. Only needed for the SSL mode.
+ */
+XGB_DLL int XGCommunicatorInit(char const* json_config);
+
+/*!
+ * \brief finalize the collective communicator,
+ * call this function after you finished all jobs.
+ * \return true if the communicator is finalized successfully otherwise false
+ */
+XGB_DLL int XGCommunicatorFinalize(void);
+
+/*!
+ * \brief get rank of current process
+ * \return rank number of worker
+ * */
+XGB_DLL int XGCommunicatorGetRank(void);
+
+/*!
+ * \brief get total number of process
+ * \return total world size
+ * */
+XGB_DLL int XGCommunicatorGetWorldSize(void);
+
+/*!
+ * \brief get if the communicator is distributed
+ * \return if the communicator is distributed
+ * */
+XGB_DLL int XGCommunicatorIsDistributed(void);
+
+/*!
+ * \brief print the msg to the communicator,
+ *    this function can be used to communicate the information of the progress to
+ *    the user who monitors the communicator
+ * \param message the message to be printed
+ */
+XGB_DLL int XGCommunicatorPrint(char const *message);
+
+/*!
+ * \brief get name of processor
+ * \param name_str pointer to received returned processor name.
+ * \return 0 for success, -1 for failure
+ */
+XGB_DLL int XGCommunicatorGetProcessorName(const char** name_str);
+
+/*!
+ * \brief broadcast an memory region to all others from root
+ *
+ *     Example: int a = 1; Broadcast(&a, sizeof(a), root);
+ * \param send_receive_buffer the pointer to send or receive buffer,
+ * \param size the size of the data
+ * \param root the root of process
+ */
+XGB_DLL int XGCommunicatorBroadcast(void *send_receive_buffer, size_t size, int root);
+
+/*!
+ * \brief perform in-place allreduce, on sendrecvbuf
+ *        this function is NOT thread-safe
+ *
+ * Example Usage: the following code gives sum of the result
+ *     vector<int> data(10);
+ *     ...
+ *     Allreduce<op::Sum>(&data[0], data.size());
+ *     ...
+ * \param send_receive_buffer buffer for both sending and receiving data
+ * \param count number of elements to be reduced
+ * \param enum_dtype the enumeration of data type, see xgboost::collective::DataType in communicator.h
+ * \param enum_op the enumeration of operation type, see xgboost::collective::Operation in communicator.h
+ */
+XGB_DLL int XGCommunicatorAllreduce(void *send_receive_buffer, size_t count, int enum_dtype,
+                                    int enum_op);
+
+
 #endif  // XGBOOST_C_API_H_

plugin/federated/engine_federated.cc

@@ -71,7 +71,9 @@ class FederatedEngine : public IEngine {
   void Allreduce(void *sendrecvbuf, size_t size, mpi::DataType dtype, mpi::OpType op) {
     auto *buffer = reinterpret_cast<char *>(sendrecvbuf);
     std::string const send_buffer(buffer, size);
-    auto const receive_buffer = client_->Allreduce(send_buffer, GetDataType(dtype), GetOp(op));
+    auto const receive_buffer =
+        client_->Allreduce(send_buffer, static_cast<xgboost::federated::DataType>(dtype),
+                           static_cast<xgboost::federated::ReduceOperation>(op));
     receive_buffer.copy(buffer, size);
   }
 
@@ -113,51 +115,6 @@ class FederatedEngine : public IEngine {
   }
 
  private:
-  /** @brief Transform mpi::DataType to xgboost::federated::DataType. */
-  static xgboost::federated::DataType GetDataType(mpi::DataType data_type) {
-    switch (data_type) {
-      case mpi::kChar:
-        return xgboost::federated::CHAR;
-      case mpi::kUChar:
-        return xgboost::federated::UCHAR;
-      case mpi::kInt:
-        return xgboost::federated::INT;
-      case mpi::kUInt:
-        return xgboost::federated::UINT;
-      case mpi::kLong:
-        return xgboost::federated::LONG;
-      case mpi::kULong:
-        return xgboost::federated::ULONG;
-      case mpi::kFloat:
-        return xgboost::federated::FLOAT;
-      case mpi::kDouble:
-        return xgboost::federated::DOUBLE;
-      case mpi::kLongLong:
-        return xgboost::federated::LONGLONG;
-      case mpi::kULongLong:
-        return xgboost::federated::ULONGLONG;
-    }
-    utils::Error("unknown mpi::DataType");
-    return xgboost::federated::CHAR;
-  }
-
-  /** @brief Transform mpi::OpType to enum to MPI OP */
-  static xgboost::federated::ReduceOperation GetOp(mpi::OpType op_type) {
-    switch (op_type) {
-      case mpi::kMax:
-        return xgboost::federated::MAX;
-      case mpi::kMin:
-        return xgboost::federated::MIN;
-      case mpi::kSum:
-        return xgboost::federated::SUM;
-      case mpi::kBitwiseOR:
-        utils::Error("Bitwise OR is not supported");
-        return xgboost::federated::MAX;
-    }
-    utils::Error("unknown mpi::OpType");
-    return xgboost::federated::MAX;
-  }
-
   void SetParam(std::string const &name, std::string const &val) {
     if (!strcasecmp(name.c_str(), "FEDERATED_SERVER_ADDRESS")) {
       server_address_ = val;

plugin/federated/federated.proto

@@ -12,16 +12,14 @@ service Federated {
 }
 
 enum DataType {
-  CHAR = 0;
-  UCHAR = 1;
-  INT = 2;
-  UINT = 3;
-  LONG = 4;
-  ULONG = 5;
+  INT8 = 0;
+  UINT8 = 1;
+  INT32 = 2;
+  UINT32 = 3;
+  INT64 = 4;
+  UINT64 = 5;
   FLOAT = 6;
   DOUBLE = 7;
-  LONGLONG = 8;
-  ULONGLONG = 9;
 }
 
 enum ReduceOperation {

plugin/federated/federated_communicator.h

@@ -0,0 +1,150 @@
+/*!
+ * Copyright 2022 XGBoost contributors
+ */
+#pragma once
+#include <xgboost/json.h>
+
+#include "../../src/collective/communicator.h"
+#include "../../src/common/io.h"
+#include "federated_client.h"
+
+namespace xgboost {
+namespace collective {
+
+/**
+ * @brief A federated learning communicator class that handles collective communication.
+ */
+class FederatedCommunicator : public Communicator {
+ public:
+  static Communicator *Create(Json const &config) {
+    std::string server_address{};
+    int world_size{0};
+    int rank{-1};
+    std::string server_cert{};
+    std::string client_key{};
+    std::string client_cert{};
+
+    // Parse environment variables first.
+    auto *value = getenv("FEDERATED_SERVER_ADDRESS");
+    if (value != nullptr) {
+      server_address = value;
+    }
+    value = getenv("FEDERATED_WORLD_SIZE");
+    if (value != nullptr) {
+      world_size = std::stoi(value);
+    }
+    value = getenv("FEDERATED_RANK");
+    if (value != nullptr) {
+      rank = std::stoi(value);
+    }
+    value = getenv("FEDERATED_SERVER_CERT");
+    if (value != nullptr) {
+      server_cert = value;
+    }
+    value = getenv("FEDERATED_CLIENT_KEY");
+    if (value != nullptr) {
+      client_key = value;
+    }
+    value = getenv("FEDERATED_CLIENT_CERT");
+    if (value != nullptr) {
+      client_cert = value;
+    }
+
+    // Runtime configuration overrides.
+    auto const &j_server_address = config["federated_server_address"];
+    if (IsA<String const>(j_server_address)) {
+      server_address = get<String const>(j_server_address);
+    }
+    auto const &j_world_size = config["federated_world_size"];
+    if (IsA<Integer const>(j_world_size)) {
+      world_size = static_cast<int>(get<Integer const>(j_world_size));
+    }
+    auto const &j_rank = config["federated_rank"];
+    if (IsA<Integer const>(j_rank)) {
+      rank = static_cast<int>(get<Integer const>(j_rank));
+    }
+    auto const &j_server_cert = config["federated_server_cert"];
+    if (IsA<String const>(j_server_cert)) {
+      server_cert = get<String const>(j_server_cert);
+    }
+    auto const &j_client_key = config["federated_client_key"];
+    if (IsA<String const>(j_client_key)) {
+      client_key = get<String const>(j_client_key);
+    }
+    auto const &j_client_cert = config["federated_client_cert"];
+    if (IsA<String const>(j_client_cert)) {
+      client_cert = get<String const>(j_client_cert);
+    }
+
+    if (server_address.empty()) {
+      LOG(FATAL) << "Federated server address must be set.";
+    }
+    if (world_size == 0) {
+      LOG(FATAL) << "Federated world size must be set.";
+    }
+    if (rank == -1) {
+      LOG(FATAL) << "Federated rank must be set.";
+    }
+    return new FederatedCommunicator(world_size, rank, server_address, server_cert, client_key,
+                                     client_cert);
+  }
+
+  /**
+   * @brief Construct a new federated communicator.
+   *
+   * @param world_size Total number of processes.
+   * @param rank       Rank of the current process.
+   */
+  FederatedCommunicator(int world_size, int rank, std::string const &server_address,
+                        std::string const &server_cert_path, std::string const &client_key_path,
+                        std::string const &client_cert_path)
+      : Communicator{world_size, rank} {
+    if (server_cert_path.empty() || client_key_path.empty() || client_cert_path.empty()) {
+      client_.reset(new xgboost::federated::FederatedClient(server_address, rank));
+    } else {
+      client_.reset(new xgboost::federated::FederatedClient(
+          server_address, rank, xgboost::common::ReadAll(server_cert_path),
+          xgboost::common::ReadAll(client_key_path), xgboost::common::ReadAll(client_cert_path)));
+    }
+  }
+
+  /** @brief Insecure communicator for testing only. */
+  FederatedCommunicator(int world_size, int rank, std::string const &server_address)
+      : Communicator{world_size, rank} {
+    client_.reset(new xgboost::federated::FederatedClient(server_address, rank));
+  }
+
+  ~FederatedCommunicator() override { client_.reset(); }
+
+  bool IsDistributed() const override { return true; }
+
+  void AllReduce(void *send_receive_buffer, std::size_t count, DataType data_type,
+                 Operation op) override {
+    std::string const send_buffer(reinterpret_cast<char const *>(send_receive_buffer),
+                                  count * GetTypeSize(data_type));
+    auto const received =
+        client_->Allreduce(send_buffer, static_cast<xgboost::federated::DataType>(data_type),
+                           static_cast<xgboost::federated::ReduceOperation>(op));
+    received.copy(reinterpret_cast<char *>(send_receive_buffer), count * GetTypeSize(data_type));
+  }
+
+  void Broadcast(void *send_receive_buffer, std::size_t size, int root) override {
+    if (GetWorldSize() == 1) return;
+    if (GetRank() == root) {
+      std::string const send_buffer(reinterpret_cast<char const *>(send_receive_buffer), size);
+      client_->Broadcast(send_buffer, root);
+    } else {
+      auto const received = client_->Broadcast("", root);
+      received.copy(reinterpret_cast<char *>(send_receive_buffer), size);
+    }
+  }
+
+  std::string GetProcessorName() override { return "rank" + std::to_string(GetRank()); }
+
+  void Print(const std::string &message) override { LOG(CONSOLE) << message; }
+
+ private:
+  std::unique_ptr<xgboost::federated::FederatedClient> client_{};
+};
+}  // namespace collective
+}  // namespace xgboost