From b0a148ab025bafe26961ee167ed3ff54c3bd12fb Mon Sep 17 00:00:00 2001
From: Rong Ou <rong.ou@gmail.com>
Date: Thu, 25 Aug 2022 15:14:03 -0700
Subject: [PATCH] Set max message size in insecure gRPC

---
 plugin/federated/federated_client.h | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/plugin/federated/federated_client.h b/plugin/federated/federated_client.h
index ab9fc895e342..592ab75edce8 100644
--- 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) {