[service] Set the same options consistently for gRPC in Python and C++.

facebookresearch · Feb 17, 2022 · f287e36 · f287e36
1 parent f30e0e0
commit f287e36
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 9 deletions.
diff --git a/compiler_gym/service/connection.py b/compiler_gym/service/connection.py
@@ -37,6 +37,7 @@
 GRPC_CHANNEL_OPTIONS = [
     # Disable the inbound message length filter to allow for large messages such
     # as observations.
+    ("grpc.max_send_message_length", -1),
     ("grpc.max_receive_message_length", -1),
     # Fix for "received initial metadata size exceeds limit"
     ("grpc.max_metadata_size", 512 * 1024),

diff --git a/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.cc b/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.cc
@@ -20,4 +20,15 @@ void shutdown_handler(int signum) {
   shutdownSignal.set_value();
 }
 
+// Configure the gRPC server, using the same options as the python client. See
+// GRPC_CHANNEL_OPTIONS in compiler_gym/service/connection.py for the python
+// equivalents and the rationale for each.
+void setGrpcChannelOptions(grpc::ServerBuilder& builder) {
+  builder.SetMaxReceiveMessageSize(-1);
+  builder.SetMaxSendMessageSize(-1);
+  builder.AddChannelArgument(GRPC_ARG_MAX_METADATA_SIZE, 512 * 1024);
+  builder.AddChannelArgument(GRPC_ARG_ENABLE_HTTP_PROXY, 0);
+  builder.AddChannelArgument(GRPC_ARG_ALLOW_REUSEPORT, 0);
+}
+
 }  // namespace compiler_gym::runtime
diff --git a/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h b/compiler_gym/service/runtime/CreateAndRunCompilerGymServiceImpl.h
@@ -29,12 +29,10 @@ namespace compiler_gym::runtime {
 
 extern std::promise<void> shutdownSignal;
 
-// Increase maximum message size beyond the 4MB default as inbound message
-// may be larger (e.g., in the case of IR strings).
-constexpr size_t kMaxMessageSizeInBytes = 512 * 1024 * 1024;
-
 void shutdown_handler(int signum);
 
+void setGrpcChannelOptions(grpc::ServerBuilder& builder);
+
 // Create a service, configured using --port and --working_dir flags, and run
 // it. This function never returns.
 //
@@ -87,7 +85,7 @@ template <typename CompilationSessionType>
   grpc::ServerBuilder builder;
   builder.RegisterService(&service);
 
-  builder.SetMaxMessageSize(kMaxMessageSizeInBytes);
+  setGrpcChannelOptions(builder);
 
   // Start a channel on the port.
   int port;

diff --git a/compiler_gym/service/runtime/create_and_run_compiler_gym_service.py b/compiler_gym/service/runtime/create_and_run_compiler_gym_service.py
@@ -18,6 +18,7 @@
 import grpc
 from absl import app, flags, logging
 
+from compiler_gym.service import connections
 from compiler_gym.service.compilation_session import CompilationSession
 from compiler_gym.service.proto import compiler_gym_service_pb2_grpc
 from compiler_gym.service.runtime.compiler_gym_service import CompilerGymService
@@ -94,10 +95,7 @@ def main(argv):
         # Create the service.
         server = grpc.server(
             futures.ThreadPoolExecutor(max_workers=FLAGS.rpc_service_threads),
-            options=[
-                ("grpc.max_send_message_length", MAX_MESSAGE_SIZE_IN_BYTES),
-                ("grpc.max_receive_message_length", MAX_MESSAGE_SIZE_IN_BYTES),
-            ],
+            options=connections.CLIENT_CONNECTION_OPTIONS,
         )
         service = CompilerGymService(
             working_directory=working_dir,