Set max message size in insecure gRPC (#8203)

dmlc · Aug 26, 2022 · d6e2013 · d6e2013
1 parent 651f0a8
commit d6e2013
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/plugin/federated/federated_client.h b/plugin/federated/federated_client.h
@@ -35,8 +35,12 @@ class FederatedClient {
 
   /** @brief Insecure client for connecting to localhost only. */
   FederatedClient(std::string const &server_address, int rank)
-      : stub_{Federated::NewStub(
-            grpc::CreateChannel(server_address, grpc::InsecureChannelCredentials()))},
+      : stub_{[&] {
+          grpc::ChannelArguments args;
+          args.SetMaxReceiveMessageSize(std::numeric_limits<int>::max());
+          return Federated::NewStub(
+              grpc::CreateCustomChannel(server_address, grpc::InsecureChannelCredentials(), args));
+        }()},
         rank_{rank} {}
 
   std::string Allgather(std::string const &send_buffer) {