Add new flag --flagfile to support load flags from file

tensorflow_serving/model_servers/BUILD

@@ -416,6 +416,19 @@ cc_library(
  ] + SUPPORTED_TENSORFLOW_OPS,
 )
 
+cc_library(
+ name = "flagfile",
+ hdrs = [
+ "flagfile.h",
+ ],
+ srcs = [
+ "flagfile.cc",
+ ],
+ deps = [
+ "@com_google_absl//absl/strings",
+ ],
+)
+
 cc_library(
  name = "tensorflow_model_server_main_lib",
  srcs = [
@@ -431,6 +444,7 @@ cc_library(
  ],
  deps = [
  ":server_lib",
+ ":flagfile",
  "@org_tensorflow//tensorflow/c:c_api",
  "@org_tensorflow//tensorflow/compiler/jit:xla_cpu_jit",
  "@org_tensorflow//tensorflow/core:lib",

tensorflow_serving/model_servers/flagfile.cc

@@ -0,0 +1,52 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_serving/model_servers/flagfile.h"
+
+#include <iostream>
+#include <fstream>
+#include <sstream>
+#include <vector>
+
+#include "absl/strings/str_split.h"
+
+namespace tensorflow {
+namespace serving {
+namespace main {
+
+static std::string ReadFileIntoString(const std::string& file) {
+ std::ifstream inf(file);
+ if (!inf.is_open()) {
+ return "";
+ }
+
+ std::stringstream ss;
+ ss << inf.rdbuf();
+
+ return ss.str();
+}
+
+std::vector<std::string> LoadFlagsFromFile(const std::string& file) {
+ std::string content = ReadFileIntoString(file);
+
+ std::vector<std::string> flags = absl::StrSplit(
+ content, '\n', absl::SkipEmpty());
+
+ return flags;
+}
+
+} // namespace main
+} // namespace serving
+} // namespace tensorflow

tensorflow_serving/model_servers/flagfile.h

@@ -0,0 +1,33 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_SERVING_MODEL_SERVERS_FLAGFILE_H_
+#define TENSORFLOW_SERVING_MODEL_SERVERS_FLAGFILE_H_
+
+#include <string>
+#include <vector>
+
+namespace tensorflow {
+namespace serving {
+namespace main {
+
+std::vector<std::string> LoadFlagsFromFile(const std::string& file);
+
+} // namespace main
+} // namespace serving
+} // namespace tensorflow
+
+
+#endif // TENSORFLOW_SERVING_MODEL_SERVERS_FLAGFILE_H_ 

tensorflow_serving/model_servers/main.cc

@@ -54,13 +54,15 @@ limitations under the License.
 #include "tensorflow/core/tpu/tpu_global_init.h"
 #endif
 #include "tensorflow/core/util/command_line_flags.h"
+#include "tensorflow_serving/model_servers/flagfile.h"
 #include "tensorflow_serving/model_servers/server.h"
 #include "tensorflow_serving/model_servers/version.h"
 
 int main(int argc, char** argv) {
  tensorflow::serving::main::Server::Options options;
  bool display_version = false;
  bool xla_cpu_compilation_enabled = false;
+ std::string flagfile;
  std::vector<tensorflow::Flag> flag_list = {
  tensorflow::Flag("port", &options.grpc_port,
  "TCP port to listen on for gRPC/HTTP API. Disabled if "
@@ -69,6 +71,8 @@ int main(int argc, char** argv) {
  "If non-empty, listen to a UNIX socket for gRPC API "
  "on the given path. Can be either relative or absolute "
  "path."),
+ tensorflow::Flag("flagfile", &flagfile,
+ "If non-empty, load flags from file"),
  tensorflow::Flag("rest_api_port", &options.http_port,
  "Port to listen on for HTTP/REST API. If set to zero "
  "HTTP/REST API will not be exported. This port must be "
@@ -269,6 +273,26 @@ int main(int argc, char** argv) {
  return -1;
  }
 
+ if (!flagfile.empty()) {
+ auto flags = tensorflow::serving::main::LoadFlagsFromFile(flagfile);
+ // With one default value to skip argv[0]
+ std::vector<char *> char_flags = {nullptr};
+ for (auto& flag : flags) {
+ if (flag.empty() || flag[0] == '#') {
+ continue;
+ }
+
+ char_flags.push_back(flag.data());
+ }
+
+ int flag_cnt = char_flags.size();
+ if (!char_flags.empty() &&
+ !tensorflow::Flags::Parse(&flag_cnt, char_flags.data(), flag_list)) {
+ std::cout << usage;
+ return -1;
+ }
+ }
+
 #if defined(LIBTPU_ON_GCE)
  std::cout << "Initializing TPU system.";
  TF_QCHECK_OK(tensorflow::InitializeTPUSystemGlobally())