Increase gRPC max receive message size for federated learning (#7958)

dmlc · Jun 1, 2022 · d3429f2 · d3429f2
1 parent 545fd45
commit d3429f2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git 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
@@ -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());