From f14b6e477ba27b57805268741c3515524951fbf9 Mon Sep 17 00:00:00 2001
From: Eric Liang <ekhliang@gmail.com>
Date: Mon, 24 Feb 2020 23:22:49 -0800
Subject: [PATCH] Raise gRPC message size limit to 100MB (#7269)

---
 python/ray/tests/test_basic.py  | 16 ++++++++++++++++
 src/ray/common/ray_config_def.h |  4 ++++
 src/ray/rpc/grpc_client.h       |  5 +++++
 src/ray/rpc/grpc_server.cc      |  6 ++++++
 4 files changed, 31 insertions(+)

diff --git a/python/ray/tests/test_basic.py b/python/ray/tests/test_basic.py
index d1f84358e..9e504cf70 100644
--- a/python/ray/tests/test_basic.py
+++ b/python/ray/tests/test_basic.py
@@ -38,6 +38,22 @@ def test_ignore_http_proxy(shutdown_only):
     assert ray.get(f.remote()) == 1
 
 
+# https://github.com/ray-project/ray/issues/7263
+def test_grpc_message_size(shutdown_only):
+    ray.init(num_cpus=1)
+
+    @ray.remote
+    def bar(*a):
+        return
+
+    # 50KiB, not enough to spill to plasma, but will be inlined.
+    def f():
+        return np.zeros(50000, dtype=np.uint8)
+
+    # Executes a 10MiB task spec
+    ray.get(bar.remote(*[f() for _ in range(200)]))
+
+
 # https://github.com/ray-project/ray/issues/7287
 def test_omp_threads_set(shutdown_only):
     ray.init(num_cpus=1)
diff --git a/src/ray/common/ray_config_def.h b/src/ray/common/ray_config_def.h
index 47f61d3d0..c68b178ef 100644
--- a/src/ray/common/ray_config_def.h
+++ b/src/ray/common/ray_config_def.h
@@ -91,6 +91,10 @@ RAY_CONFIG(bool, new_scheduler_enabled, false)
 // Objects larger than this size will be spilled/promoted to plasma.
 RAY_CONFIG(int64_t, max_direct_call_object_size, 100 * 1024)
 
+// The max gRPC message size (the gRPC internal default is 4MB). We use a higher
+// limit in Ray to avoid crashing with many small inlined task arguments.
+RAY_CONFIG(int64_t, max_grpc_message_size, 100 * 1024 * 1024)
+
 // The min number of retries for direct actor creation tasks. The actual number
 // of creation retries will be MAX(actor_creation_min_retries, max_reconstructions).
 RAY_CONFIG(uint64_t, actor_creation_min_retries, 3)
diff --git a/src/ray/rpc/grpc_client.h b/src/ray/rpc/grpc_client.h
index 1999ef7e3..93550dbb6 100644
--- a/src/ray/rpc/grpc_client.h
+++ b/src/ray/rpc/grpc_client.h
@@ -5,6 +5,7 @@
 #include <boost/asio.hpp>
 
 #include "ray/common/grpc_util.h"
+#include "ray/common/ray_config.h"
 #include "ray/common/status.h"
 #include "ray/rpc/client_call.h"
 
@@ -42,6 +43,8 @@ class GrpcClient {
     // Disable http proxy since it disrupts local connections. TODO(ekl) we should make
     // this configurable, or selectively set it for known local connections only.
     argument.SetInt(GRPC_ARG_ENABLE_HTTP_PROXY, 0);
+    argument.SetMaxSendMessageSize(RayConfig::instance().max_grpc_message_size());
+    argument.SetMaxReceiveMessageSize(RayConfig::instance().max_grpc_message_size());
     std::shared_ptr<grpc::Channel> channel =
         grpc::CreateCustomChannel(address + ":" + std::to_string(port),
                                   grpc::InsecureChannelCredentials(), argument);
@@ -56,6 +59,8 @@ class GrpcClient {
     grpc::ChannelArguments argument;
     argument.SetResourceQuota(quota);
     argument.SetInt(GRPC_ARG_ENABLE_HTTP_PROXY, 0);
+    argument.SetMaxSendMessageSize(RayConfig::instance().max_grpc_message_size());
+    argument.SetMaxReceiveMessageSize(RayConfig::instance().max_grpc_message_size());
     std::shared_ptr<grpc::Channel> channel =
         grpc::CreateCustomChannel(address + ":" + std::to_string(port),
                                   grpc::InsecureChannelCredentials(), argument);
diff --git a/src/ray/rpc/grpc_server.cc b/src/ray/rpc/grpc_server.cc
index 556ec2f19..45606e9d4 100644
--- a/src/ray/rpc/grpc_server.cc
+++ b/src/ray/rpc/grpc_server.cc
@@ -3,6 +3,8 @@
 #include <grpcpp/impl/service_type.h>
 #include <boost/asio/detail/socket_holder.hpp>
 
+#include "ray/common/ray_config.h"
+
 namespace ray {
 namespace rpc {
 
@@ -20,6 +22,10 @@ void GrpcServer::Run() {
   // (default behavior in grpc), we may see multiple workers listen on the same port and
   // the requests sent to this port may be handled by any of the workers.
   builder.AddChannelArgument(GRPC_ARG_ALLOW_REUSEPORT, 0);
+  builder.AddChannelArgument(GRPC_ARG_MAX_SEND_MESSAGE_LENGTH,
+                             RayConfig::instance().max_grpc_message_size());
+  builder.AddChannelArgument(GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH,
+                             RayConfig::instance().max_grpc_message_size());
   // TODO(hchen): Add options for authentication.
   builder.AddListeningPort(server_address, grpc::InsecureServerCredentials(), &port_);
   // Register all the services to this server.