experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Allow default GRPC options to be seeded via environment variables.

Browse Source 
PiperOrigin-RevId: 272566602
This commit is contained in:
Davide Libenzi 2019-10-02 18:37:39 -07:00 committed by TensorFlower Gardener
parent ff650ad4c4
commit 44c25c8870
2 changed files with 45 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/distributed_runtime/rpc/BUILD
View File

@ -108,6 +108,7 @@ cc_library(
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
],
)

46
tensorflow/core/distributed_runtime/rpc/grpc_channel.cc
View File

@ -15,12 +15,14 @@ limitations under the License.
#include "tensorflow/core/distributed_runtime/rpc/grpc_channel.h"
#include <cstdlib>
#include <limits>
#include <map>
#include <unordered_map>
#include "grpcpp/create_channel.h"
#include "absl/strings/escaping.h"
#include "absl/strings/str_split.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/map_util.h"
@ -58,11 +60,51 @@ Status ValidateHostPortPair(const string& host_port) {
return Status::OK();
}
::grpc::ChannelArguments* CreateDefaultChannelArguments() {
::grpc::ChannelArguments* args = new ::grpc::ChannelArguments();
const char* env = std::getenv("TF_GRPC_DEFAULT_OPTIONS");
if (env != nullptr) {
for (auto& grpc_option : absl::StrSplit(env, ',')) {
std::vector<string> name_value = absl::StrSplit(grpc_option, '=');
if (name_value.size() != 2) {
LOG(ERROR) << "Invalid GRPC options format: " << grpc_option;
continue;
}
VLOG(3) << "Setting GRPC default for '" << name_value[0] << "' to '"
<< name_value[1] << "'";
if (name_value[1].size() >= 2 && name_value[1][0] == '"') {
string ue_value = name_value[1].substr(1, name_value[1].size() - 2);
string value;
string error;
if (!absl::CUnescape(ue_value, &value, &error)) {
LOG(ERROR) << "Failed to parse escaped string for " << grpc_option
<< ": " << error;
} else {
args->SetString(name_value[0], value);
}
} else {
int64 value;
if (strings::safe_strto64(name_value[1], &value)) {
args->SetInt(name_value[0], value);
} else {
LOG(ERROR) << "Invalid integer value: " << grpc_option;
}
}
}
}
return args;
}
const ::grpc::ChannelArguments* GetDefaultChannelArguments() {
static const ::grpc::ChannelArguments* args = CreateDefaultChannelArguments();
return args;
}
} // namespace
::grpc::ChannelArguments GetChannelArguments(const RPCOptions* rpc_options) {
// TODO(mrry): Implement secure channels.
::grpc::ChannelArguments args;
::grpc::ChannelArguments args = *GetDefaultChannelArguments();
args.SetInt(GRPC_ARG_MAX_MESSAGE_LENGTH, std::numeric_limits<int32>::max());
// NOTE(mrry): Some versions of gRPC use a 20-second minimum backoff
// on connection failure, which makes our tests time out.