From ce70229ba6b67a9ed3d4a5a315e88a9c1e26389d Mon Sep 17 00:00:00 2001
From: gongweibao <weibao.gong@gmail.com>
Date: Fri, 28 Dec 2018 15:45:05 +0800
Subject: [PATCH] Add max_body_size flags to brpc (#15084)

---
 .../distributed/collective_server_test.cc     |  5 +-
 paddle/fluid/pybind/pybind.cc                 | 10 ++--
 paddle/testing/paddle_gtest_main.cc           | 53 +++++++++++++++----
 python/paddle/fluid/__init__.py               |  4 ++
 4 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/paddle/fluid/operators/distributed/collective_server_test.cc b/paddle/fluid/operators/distributed/collective_server_test.cc
index c5d18f7c6..46c761000 100644
--- a/paddle/fluid/operators/distributed/collective_server_test.cc
+++ b/paddle/fluid/operators/distributed/collective_server_test.cc
@@ -52,12 +52,12 @@ std::unique_ptr<framework::Scope> GenerateVars(platform::Place place) {
   framework::Scope* scope = new framework::Scope();
   framework::Variable* var = scope->Var("var1");
   auto* slr = var->GetMutable<framework::SelectedRows>();
-  slr->set_height(1000);
+  slr->set_height(20000);
 
   auto* tensor = slr->mutable_value();
   auto* rows = slr->mutable_rows();
 
-  tensor->Resize(framework::make_ddim({3, 5}));
+  tensor->Resize(framework::make_ddim({20000, 1024}));
   tensor->mutable_data<float>(place);
 
   paddle::operators::math::set_constant(ctx, tensor, 32.7);
@@ -83,6 +83,7 @@ void Gather(const std::vector<distributed::RemoteVar>& vars,
 }
 
 TEST(PREFETCH, GPU) {
+  setenv("FLAGS_max_body_size", "2147483647", 1);
   platform::CUDAPlace place;
   platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
   auto& ctx = *pool.Get(place);
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
index 2ffdc90d8..d664107d5 100644
--- a/paddle/fluid/pybind/pybind.cc
+++ b/paddle/fluid/pybind/pybind.cc
@@ -84,11 +84,15 @@ bool IsCompiledWithCUDA() {
 }
 
 bool IsCompiledWithBrpc() {
-#if defined(PADDLE_WITH_BRPC) || defined(PADDLE_WITH_BRPC_RDMA)
-  return true;
-#else
+#ifndef PADDLE_WITH_DISTRIBUTE
   return false;
 #endif
+
+#ifdef PADDLE_WITH_GRPC
+  return false;
+#endif
+
+  return true;
 }
 
 bool IsCompiledWithDIST() {
diff --git a/paddle/testing/paddle_gtest_main.cc b/paddle/testing/paddle_gtest_main.cc
index ef43d13e1..47c5248b5 100644
--- a/paddle/testing/paddle_gtest_main.cc
+++ b/paddle/testing/paddle_gtest_main.cc
@@ -28,20 +28,53 @@ int main(int argc, char** argv) {
   for (int i = 0; i < argc; ++i) {
     new_argv.push_back(argv[i]);
   }
+
+  std::vector<std::string> envs;
+  std::vector<std::string> undefok;
+#if defined(PADDLE_WITH_DISTRIBUTE) && !defined(PADDLE_WITH_GRPC)
+  envs.push_back("max_body_size");
+#endif
+
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-  new_argv.push_back(
-      strdup("--tryfromenv=fraction_of_gpu_memory_to_use,allocator_strategy"));
+  envs.push_back("fraction_of_gpu_memory_to_use");
+  envs.push_back("allocator_strategy");
 #elif __clang__
-  new_argv.push_back(
-      strdup("--tryfromenv=use_mkldnn,initial_cpu_memory_in_"
-             "mb,allocator_strategy"));
-  new_argv.push_back(strdup("--undefok=use_mkldnn,initial_cpu_memory_in_mb"));
+  envs.push_back("use_mkldnn");
+  envs.push_back("initial_cpu_memory_in_mb");
+  envs.push_back("allocator_strategy");
+
+  undefok.push_back("use_mkldnn");
+  undefok.push_back("initial_cpu_memory_in_mb");
 #else
-  new_argv.push_back(
-      strdup("--tryfromenv=use_pinned_memory,use_mkldnn,initial_cpu_memory_in_"
-             "mb,allocator_strategy"));
-  new_argv.push_back(strdup("--undefok=use_mkldnn,initial_cpu_memory_in_mb"));
+  envs.push_back("use_pinned_memory");
+  envs.push_back("use_mkldnn");
+  envs.push_back("initial_cpu_memory_in_mb");
+  envs.push_back("allocator_strategy");
+
+  undefok.push_back("use_mkldnn");
+  undefok.push_back("initial_cpu_memory_in_mb");
 #endif
+
+  if (envs.size() > 0) {
+    std::string env_string = "--tryfromenv=";
+    for (auto t : envs) {
+      env_string += t + ",";
+    }
+    env_string = env_string.substr(0, env_string.length() - 1);
+    new_argv.push_back(strdup(env_string.c_str()));
+    VLOG(1) << "gtest env_string:" << env_string;
+  }
+
+  if (undefok.size() > 0) {
+    std::string undefok_string = "--undefok=";
+    for (auto t : undefok) {
+      undefok_string += t + ",";
+    }
+    undefok_string = undefok_string.substr(0, undefok_string.length() - 1);
+    new_argv.push_back(strdup(undefok_string.c_str()));
+    VLOG(1) << "gtest undefok_string:" << undefok_string;
+  }
+
   int new_argc = static_cast<int>(new_argv.size());
   char** new_argv_address = new_argv.data();
   google::ParseCommandLineFlags(&new_argc, &new_argv_address, false);
diff --git a/python/paddle/fluid/__init__.py b/python/paddle/fluid/__init__.py
index 7433c2cbb..7a7267093 100644
--- a/python/paddle/fluid/__init__.py
+++ b/python/paddle/fluid/__init__.py
@@ -151,6 +151,10 @@ def __bootstrap__():
         read_env_flags.append('rpc_get_thread_num')
         read_env_flags.append('rpc_prefetch_thread_num')
         read_env_flags.append('rpc_disable_reuse_port')
+        if core.is_compiled_with_brpc():
+            read_env_flags.append('max_body_size')
+            #set brpc max body size
+            os.environ['FLAGS_max_body_size'] = "2147483647"
 
     if core.is_compiled_with_cuda():
         read_env_flags += [
-- 
GitLab