Initial support for federated learning (#7831)

Federated learning plugin for xgboost:
* A gRPC server to aggregate MPI-style requests (allgather, allreduce, broadcast) from federated workers.
* A Rabit engine for the federated environment.
* Integration test to simulate federated learning.

Additional followups are needed to address GPU support, better security, and privacy, etc.

tests/cpp/CMakeLists.txt

@@ -18,6 +18,14 @@ if (NOT PLUGIN_UPDATER_ONEAPI)
   list(REMOVE_ITEM TEST_SOURCES ${ONEAPI_TEST_SOURCES})
 endif (NOT PLUGIN_UPDATER_ONEAPI)
 
+if (PLUGIN_FEDERATED)
+  target_include_directories(testxgboost PRIVATE ${xgboost_SOURCE_DIR}/plugin/federated)
+  target_link_libraries(testxgboost PRIVATE federated_client)
+else (PLUGIN_FEDERATED)
+  file(GLOB_RECURSE FEDERATED_TEST_SOURCES "plugin/*_federated_*.cc")
+  list(REMOVE_ITEM TEST_SOURCES ${FEDERATED_TEST_SOURCES})
+endif (PLUGIN_FEDERATED)
+
 target_sources(testxgboost PRIVATE ${TEST_SOURCES} ${xgboost_SOURCE_DIR}/plugin/example/custom_obj.cc)
 
 if (USE_CUDA AND PLUGIN_RMM)

tests/cpp/plugin/test_federated_server.cc

@@ -0,0 +1,130 @@
+/*!
+ * Copyright 2017-2020 XGBoost contributors
+ */
+#include <grpcpp/server_builder.h>
+#include <gtest/gtest.h>
+
+#include <thread>
+
+#include "federated_client.h"
+#include "federated_server.h"
+
+namespace xgboost {
+
+class FederatedServerTest : public ::testing::Test {
+ public:
+  static void VerifyAllgather(int rank) {
+    federated::FederatedClient client{kServerAddress, rank};
+    CheckAllgather(client, rank);
+  }
+
+  static void VerifyAllreduce(int rank) {
+    federated::FederatedClient client{kServerAddress, rank};
+    CheckAllreduce(client);
+  }
+
+  static void VerifyBroadcast(int rank) {
+    federated::FederatedClient client{kServerAddress, rank};
+    CheckBroadcast(client, rank);
+  }
+
+  static void VerifyMixture(int rank) {
+    federated::FederatedClient client{kServerAddress, rank};
+    for (auto i = 0; i < 10; i++) {
+      CheckAllgather(client, rank);
+      CheckAllreduce(client);
+      CheckBroadcast(client, rank);
+    }
+  }
+
+ protected:
+  void SetUp() override {
+    server_thread_.reset(new std::thread([this] {
+      grpc::ServerBuilder builder;
+      federated::FederatedService service{kWorldSize};
+      builder.AddListeningPort(kServerAddress, grpc::InsecureServerCredentials());
+      builder.RegisterService(&service);
+      server_ = builder.BuildAndStart();
+      server_->Wait();
+    }));
+  }
+
+  void TearDown() override {
+    server_->Shutdown();
+    server_thread_->join();
+  }
+
+  static void CheckAllgather(federated::FederatedClient& client, int rank) {
+    auto reply = client.Allgather("hello " + std::to_string(rank) + " ");
+    EXPECT_EQ(reply, "hello 0 hello 1 hello 2 ");
+  }
+
+  static void CheckAllreduce(federated::FederatedClient& client) {
+    int data[] = {1, 2, 3, 4, 5};
+    std::string send_buffer(reinterpret_cast<char const*>(data), sizeof(data));
+    auto reply = client.Allreduce(send_buffer, federated::INT, federated::SUM);
+    auto const* result = reinterpret_cast<int const*>(reply.data());
+    int expected[] = {3, 6, 9, 12, 15};
+    for (auto i = 0; i < 5; i++) {
+      EXPECT_EQ(result[i], expected[i]);
+    }
+  }
+
+  static void CheckBroadcast(federated::FederatedClient& client, int rank) {
+    std::string send_buffer{};
+    if (rank == 0) {
+      send_buffer = "hello broadcast";
+    }
+    auto reply = client.Broadcast(send_buffer, 0);
+    EXPECT_EQ(reply, "hello broadcast");
+  }
+
+  static int const kWorldSize{3};
+  static std::string const kServerAddress;
+  std::unique_ptr<std::thread> server_thread_;
+  std::unique_ptr<grpc::Server> server_;
+};
+
+std::string const FederatedServerTest::kServerAddress{"localhost:56789"};  // NOLINT(cert-err58-cpp)
+
+TEST_F(FederatedServerTest, Allgather) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllgather, rank));
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}
+
+TEST_F(FederatedServerTest, Allreduce) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllreduce, rank));
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}
+
+TEST_F(FederatedServerTest, Broadcast) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyBroadcast, rank));
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}
+
+TEST_F(FederatedServerTest, Mixture) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyMixture, rank));
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}
+
+}  // namespace xgboost

tests/distributed/runtests-federated.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+set -e
+
+rm -f ./*.model* ./agaricus* ./*.pem
+
+world_size=3
+
+# Generate server and client certificates.
+openssl req -x509 -newkey rsa:2048 -days 7 -nodes -keyout server-key.pem -out server-cert.pem -subj "/C=US/CN=localhost"
+openssl req -x509 -newkey rsa:2048 -days 7 -nodes -keyout client-key.pem -out client-cert.pem -subj "/C=US/CN=localhost"
+
+# Split train and test files manually to simulate a federated environment.
+split -n l/${world_size} -d ../../demo/data/agaricus.txt.train agaricus.txt.train-
+split -n l/${world_size} -d ../../demo/data/agaricus.txt.test agaricus.txt.test-
+
+python test_federated.py ${world_size}

tests/distributed/test_federated.py

@@ -0,0 +1,78 @@
+#!/usr/bin/python
+import multiprocessing
+import sys
+import time
+
+import xgboost as xgb
+import xgboost.federated
+
+SERVER_KEY = 'server-key.pem'
+SERVER_CERT = 'server-cert.pem'
+CLIENT_KEY = 'client-key.pem'
+CLIENT_CERT = 'client-cert.pem'
+
+
+def run_server(port: int, world_size: int) -> None:
+    xgboost.federated.run_federated_server(port, world_size, SERVER_KEY, SERVER_CERT,
+                                           CLIENT_CERT)
+
+
+def run_worker(port: int, world_size: int, rank: int) -> None:
+    # Always call this before using distributed module
+    rabit_env = [
+        f'federated_server_address=localhost:{port}',
+        f'federated_world_size={world_size}',
+        f'federated_rank={rank}',
+        f'federated_server_cert={SERVER_CERT}',
+        f'federated_client_key={CLIENT_KEY}',
+        f'federated_client_cert={CLIENT_CERT}'
+    ]
+    xgb.rabit.init([e.encode() for e in rabit_env])
+
+    # Load file, file will not be sharded in federated mode.
+    dtrain = xgb.DMatrix('agaricus.txt.train-%02d' % rank)
+    dtest = xgb.DMatrix('agaricus.txt.test-%02d' % rank)
+
+    # Specify parameters via map, definition are same as c++ version
+    param = {'max_depth': 2, 'eta': 1, 'objective': 'binary:logistic'}
+
+    # Specify validations set to watch performance
+    watchlist = [(dtest, 'eval'), (dtrain, 'train')]
+    num_round = 20
+
+    # Run training, all the features in training API is available.
+    # Currently, this script only support calling train once for fault recovery purpose.
+    bst = xgb.train(param, dtrain, num_round, evals=watchlist, early_stopping_rounds=2)
+
+    # Save the model, only ask process 0 to save the model.
+    if xgb.rabit.get_rank() == 0:
+        bst.save_model("test.model.json")
+        xgb.rabit.tracker_print("Finished training\n")
+
+    # Notify the tracker all training has been successful
+    # This is only needed in distributed training.
+    xgb.rabit.finalize()
+
+
+def run_test() -> None:
+    port = 9091
+    world_size = int(sys.argv[1])
+
+    server = multiprocessing.Process(target=run_server, args=(port, world_size))
+    server.start()
+    time.sleep(1)
+    if not server.is_alive():
+        raise Exception("Error starting Federated Learning server")
+
+    workers = []
+    for rank in range(world_size):
+        worker = multiprocessing.Process(target=run_worker, args=(port, world_size, rank))
+        workers.append(worker)
+        worker.start()
+    for worker in workers:
+        worker.join()
+    server.terminate()
+
+
+if __name__ == '__main__':
+    run_test()