diff --git a/plugin/federated/federated_client.h b/plugin/federated/federated_client.h index 5aacef3a6..b00bf0f0f 100644 --- a/plugin/federated/federated_client.h +++ b/plugin/federated/federated_client.h @@ -8,6 +8,7 @@ #include #include +#include #include 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::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 c9651797a..c38bdc4b8 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::max()); builder.AddListeningPort(server_address, grpc::SslServerCredentials(options)); builder.RegisterService(&service); std::unique_ptr server(builder.BuildAndStart());