Common interface for collective communication

plugin/federated/federated_communicator.h

@@ -0,0 +1,180 @@
+/*!
+ * Copyright 2022 XGBoost contributors
+ */
+#pragma once
+#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:
+  /**
+   * @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} {
+    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(); }
+
+  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, ConvertDataType(data_type), ConvertOperation(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);
+    }
+  }
+
+ private:
+  static xgboost::federated::DataType ConvertDataType(DataType data_type) {
+    xgboost::federated::DataType result{};
+    switch (data_type) {
+      case DataType::kInt:
+        result = xgboost::federated::DataType::INT;
+        break;
+      case DataType::kFloat:
+        result = xgboost::federated::DataType::FLOAT;
+        break;
+      case DataType::kDouble:
+        result = xgboost::federated::DataType::DOUBLE;
+        break;
+    }
+    return result;
+  }
+
+  static xgboost::federated::ReduceOperation ConvertOperation(Operation operation) {
+    xgboost::federated::ReduceOperation result{};
+    switch (operation) {
+      case Operation::kMax:
+        result = xgboost::federated::ReduceOperation::MAX;
+        break;
+      case Operation::kSum:
+        result = xgboost::federated::ReduceOperation::SUM;
+        break;
+    }
+    return result;
+  }
+
+  std::unique_ptr<xgboost::federated::FederatedClient> client_{};
+};
+
+class FederatedCommunicatorFactory {
+ public:
+  FederatedCommunicatorFactory(int argc, char *argv[]) {
+    // Parse environment variables first.
+    for (auto const &env_var : env_vars_) {
+      char const *value = getenv(env_var.c_str());
+      if (value != nullptr) {
+        SetParam(env_var, value);
+      }
+    }
+
+    // Command line argument overrides.
+    for (int i = 0; i < argc; ++i) {
+      std::string const key_value = argv[i];
+      auto const delimiter = key_value.find('=');
+      if (delimiter != std::string::npos) {
+        SetParam(key_value.substr(0, delimiter), key_value.substr(delimiter + 1));
+      }
+    }
+  }
+
+  Communicator *Create() {
+    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.";
+    }
+    if (server_cert_.empty()) {
+      LOG(FATAL) << "Federated server cert must be set.";
+    }
+    if (client_key_.empty()) {
+      LOG(FATAL) << "Federated client key must be set.";
+    }
+    if (client_cert_.empty()) {
+      LOG(FATAL) << "Federated client cert must be set.";
+    }
+    return new FederatedCommunicator(world_size_, rank_, server_address_, server_cert_, client_key_,
+                                     client_cert_);
+  }
+
+  std::string const &GetServerAddress() const { return server_address_; }
+  int GetWorldSize() const { return world_size_; }
+  int GetRank() const { return rank_; }
+  std::string const &GetServerCert() const { return server_cert_; }
+  std::string const &GetClientKey() const { return client_key_; }
+  std::string const &GetClientCert() const { return client_cert_; }
+
+ private:
+  void SetParam(std::string const &name, std::string const &val) {
+    if (!strcasecmp(name.c_str(), "FEDERATED_SERVER_ADDRESS")) {
+      server_address_ = val;
+    } else if (!strcasecmp(name.c_str(), "FEDERATED_WORLD_SIZE")) {
+      world_size_ = std::stoi(val);
+    } else if (!strcasecmp(name.c_str(), "FEDERATED_RANK")) {
+      rank_ = std::stoi(val);
+    } else if (!strcasecmp(name.c_str(), "FEDERATED_SERVER_CERT")) {
+      server_cert_ = val;
+    } else if (!strcasecmp(name.c_str(), "FEDERATED_CLIENT_KEY")) {
+      client_key_ = val;
+    } else if (!strcasecmp(name.c_str(), "FEDERATED_CLIENT_CERT")) {
+      client_cert_ = val;
+    }
+  }
+
+  // clang-format off
+  std::vector<std::string> const env_vars_{
+      "FEDERATED_SERVER_ADDRESS",
+      "FEDERATED_WORLD_SIZE",
+      "FEDERATED_RANK",
+      "FEDERATED_SERVER_CERT",
+      "FEDERATED_CLIENT_KEY",
+      "FEDERATED_CLIENT_CERT" };
+  // clang-format on
+
+  std::string server_address_{};
+  int world_size_{0};
+  int rank_{-1};
+  std::string server_cert_{};
+  std::string client_key_{};
+  std::string client_cert_{};
+};
+
+}  // namespace collective
+}  // namespace xgboost

plugin/federated/federated_server.cc

@@ -10,6 +10,8 @@
 #include <fstream>
 #include <sstream>
 
+#include "../../src/common/io.h"
+
 namespace xgboost {
 namespace federated {
 
@@ -201,13 +203,6 @@ grpc::Status FederatedService::Handle(Request const* request, Reply* reply,
   return grpc::Status::OK;
 }
 
-std::string ReadFile(char const* path) {
-  auto stream = std::ifstream(path);
-  std::ostringstream out;
-  out << stream.rdbuf();
-  return out.str();
-}
-
 void RunServer(int port, int world_size, char const* server_key_file, char const* server_cert_file,
                char const* client_cert_file) {
   std::string const server_address = "0.0.0.0:" + std::to_string(port);
@@ -216,10 +211,10 @@ void RunServer(int port, int world_size, char const* server_key_file, char const
   grpc::ServerBuilder builder;
   auto options =
       grpc::SslServerCredentialsOptions(GRPC_SSL_REQUEST_AND_REQUIRE_CLIENT_CERTIFICATE_AND_VERIFY);
-  options.pem_root_certs = ReadFile(client_cert_file);
+  options.pem_root_certs = xgboost::common::ReadAll(client_cert_file);
   auto key = grpc::SslServerCredentialsOptions::PemKeyCertPair();
-  key.private_key = ReadFile(server_key_file);
-  key.cert_chain = ReadFile(server_cert_file);
+  key.private_key = xgboost::common::ReadAll(server_key_file);
+  key.cert_chain = xgboost::common::ReadAll(server_cert_file);
   options.pem_key_cert_pairs.push_back(key);
   builder.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());
   builder.AddListeningPort(server_address, grpc::SslServerCredentials(options));

src/collective/communicator.h

@@ -0,0 +1,96 @@
+/*!
+ * Copyright 2022 XGBoost contributors
+ */
+#pragma once
+#include <xgboost/logging.h>
+
+namespace xgboost {
+namespace collective {
+
+/** @brief Defines the integral and floating data types. */
+enum class DataType { kInt, kFloat, kDouble, kSizeT };
+
+inline std::size_t GetTypeSize(DataType data_type) {
+  std::size_t size{0};
+  switch (data_type) {
+    case DataType::kInt:
+      size = sizeof(int);
+      break;
+    case DataType::kFloat:
+      size = sizeof(float);
+      break;
+    case DataType::kDouble:
+      size = sizeof(double);
+      break;
+    case DataType::kSizeT:
+      size = sizeof(std::size_t);
+      break;
+  }
+  return size;
+}
+
+/** @brief Defines the reduction operation. */
+enum class Operation { kMax, kSum };
+
+/**
+ * @brief A communicator class that handles collective communication.
+ */
+class Communicator {
+ public:
+  /**
+   * @brief Construct a new communicator.
+   *
+   * @param world_size Total number of processes.
+   * @param rank       Rank of the current process.
+   */
+  Communicator(int world_size, int rank) : world_size_(world_size), rank_(rank) {
+    if (world_size < 1) {
+      LOG(FATAL) << "World size " << world_size << " is less than 1.";
+    }
+    if (rank < 0) {
+      LOG(FATAL) << "Rank " << rank << " is less than 0.";
+    }
+    if (rank >= world_size) {
+      LOG(FATAL) << "Rank " << rank << " is greater than world_size - 1: " << world_size - 1 << ".";
+    }
+  }
+
+  virtual ~Communicator() = default;
+
+  /** @brief Get the total number of processes. */
+  int GetWorldSize() const { return world_size_; }
+
+  /** @brief Get the rank of the current processes. */
+  int GetRank() const { return rank_; }
+
+  /** @brief Whether the communicator is running in distributed mode. */
+  bool IsDistributed() const { return world_size_ > 1; };
+
+  /**
+   * @brief Combines values from all processes and distributes the result back to all processes.
+   *
+   * @param send_receive_buffer Buffer storing the data.
+   * @param count               Number of elements in the buffer.
+   * @param data_type           Data type stored in the buffer.
+   * @param op                  The operation to perform.
+   */
+  virtual void AllReduce(void *send_receive_buffer, std::size_t count, DataType data_type,
+                         Operation op) = 0;
+
+  /**
+   * @brief Broadcasts a message from the process with rank `root` to all other processes of the
+   * group.
+   *
+   * @param send_receive_buffer Buffer storing the data.
+   * @param size                Size of the data in bytes.
+   * @param root                Rank of broadcast root.
+   */
+  virtual void Broadcast(void *send_receive_buffer, std::size_t size, int root) = 0;
+
+ private:
+  int const world_size_;
+  int const rank_;
+};
+
+}  // namespace collective
+}  // namespace xgboost

src/collective/communicator_factory.cu

@@ -0,0 +1,68 @@
+/*!
+ * Copyright 2022 XGBoost contributors
+ */
+#include "communicator_factory.h"
+#include "device_communicator_adapter.cuh"
+
+namespace xgboost {
+namespace collective {
+
+thread_local std::unique_ptr<CommunicatorFactory> CommunicatorFactory::instance_{};
+
+void CommunicatorFactory::Init(int argc, char* argv[]) {
+  if (instance_) {
+    LOG(FATAL) << "Communicator factory can only be initialized once.";
+  }
+
+  auto type = GetTypeFromEnv();
+  auto const arg = GetTypeFromArgs(argc, argv);
+  if (arg != CommunicatorType::kUnknown) {
+    type = arg;
+  }
+  switch (type) {
+    case CommunicatorType::kRabit:
+      LOG(FATAL) << "Not implemented yet.";
+      break;
+    case CommunicatorType::kMPI:
+      LOG(FATAL) << "Not implemented yet.";
+      break;
+    case CommunicatorType::kFederated: {
+#if defined(XGBOOST_USE_FEDERATED)
+      FederatedCommunicatorFactory factory{argc, argv};
+      auto* comm = factory.Create();
+      instance_.reset(new CommunicatorFactory(type, comm));
+#else
+      LOG(FATAL) << "XGBoost is not compiled with Federated Learning support.";
+#endif
+      break;
+    }
+    case CommunicatorType::kUnknown:
+      LOG(FATAL) << "Unknown communicator type.";
+      break;
+  }
+}
+
+void CommunicatorFactory::Finalize() { instance_.reset(); }
+
+CommunicatorFactory::CommunicatorFactory(CommunicatorType type, Communicator* communicator)
+    : type_{type}, communicator_{communicator} {}
+
+DeviceCommunicator* CommunicatorFactory::GetDeviceCommunicator(int device_ordinal) {
+  if (!device_communicator_) {
+#ifdef XGBOOST_USE_NCCL
+    if (type_ != CommunicatorType::kFederated) {
+      // Use NCCL communicator.
+      LOG(FATAL) << "Not implemented yet.";
+    } else {
+      device_communicator_.reset(
+          new DeviceCommunicatorAdapter(device_ordinal, communicator_.get()));
+    }
+#else
+    device_communicator_.reset(new DeviceCommunicatorAdapter(device_ordinal, communicator_.get()));
+#endif
+  }
+  return device_communicator_.get();
+}
+
+}  // namespace collective
+}  // namespace xgboost