From d3429f2ff632b042d4da47609b5f3dd78dd19421 Mon Sep 17 00:00:00 2001
From: Rong Ou <rong.ou@gmail.com>
Date: Tue, 31 May 2022 22:21:54 -0700
Subject: [PATCH] Increase gRPC max receive message size for federated learning
 (#7958)

---
 plugin/federated/federated_client.h  | 5 ++++-
 plugin/federated/federated_server.cc | 1 +
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/plugin/federated/federated_client.h b/plugin/federated/federated_client.h
index 5aacef3a621a..b00bf0f0f225 100644
--- a/plugin/federated/federated_client.h
+++ b/plugin/federated/federated_client.h
@@ -8,6 +8,7 @@
 
 #include <cstdio>
 #include <cstdlib>
+#include <limits>
 #include <string>
 
 namespace xgboost {
@@ -25,8 +26,10 @@ class FederatedClient {
           options.pem_root_certs = server_cert;
           options.pem_private_key = client_key;
           options.pem_cert_chain = client_cert;
+          grpc::ChannelArguments args;
+          args.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());
           return Federated::NewStub(
-              grpc::CreateChannel(server_address, grpc::SslCredentials(options)));
+              grpc::CreateCustomChannel(server_address, grpc::SslCredentials(options), args));
         }()},
         rank_{rank} {}
 
diff --git a/plugin/federated/federated_server.cc b/plugin/federated/federated_server.cc
index c9651797a42c..c38bdc4b8171 100644
--- a/plugin/federated/federated_server.cc
+++ b/plugin/federated/federated_server.cc
@@ -221,6 +221,7 @@ void RunServer(int port, int world_size, char const* server_key_file, char const
   key.private_key = ReadFile(server_key_file);
   key.cert_chain = ReadFile(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));
   builder.RegisterService(&service);
   std::unique_ptr<grpc::Server> server(builder.BuildAndStart());