[DeepRec] Support session_group to use multiple GPUs. (#15)

Author: shanshanpt

WORKSPACE

@@ -11,8 +11,8 @@ load("//tensorflow_serving:repo.bzl", "tensorflow_http_archive")
 
 tensorflow_http_archive(
     name = "org_tensorflow",
-    sha256 = "f11c490ff7cb90fba7e1fa31822de425e3243bddc767bc5ebf7ae0d373bfdbc1",
-    git_commit = "dcdc32e2e565b7a815a5906a7e1e1df202322d37",
+    sha256 = "836a1a5425a4f853abfe1dd03e2222bb9004540163fc2c56d57979b390cd5a31",
+    git_commit = "c3a656197cae9647fad47ec7d84ccd722a502815",
 )
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")

tensorflow_serving/model_servers/main.cc

@@ -160,6 +160,10 @@ int main(int argc, char** argv) {
                        &options.session_num_per_group,
                        "Session num for a session group, "
                        "default 0 means we not use session group."),
+      tensorflow::Flag("gpu_ids_list",
+                       &options.gpu_ids_list,
+                       "GPU id list for a session group, "
+                       "default '', the legal format is '0,1,2...'."),
       tensorflow::Flag(
           "ssl_config_file", &options.ssl_config_file,
           "If non-empty, read an ascii SSLConfig protobuf from "

tensorflow_serving/model_servers/server.cc

@@ -264,6 +264,15 @@ Status CreatePlatformConfigMapV2(const Server::Options& server_options,
   model_session_config->set_session_num(
       server_options.session_num_per_group);
 
+  // gpu_ids_list for a session group
+  if (!server_options.gpu_ids_list.empty()) {
+    std::vector<string> ids =
+        str_util::Split(server_options.gpu_ids_list, ',');
+    for (auto id : ids) {
+      model_session_config->add_gpu_ids(std::stoi(id));
+    }
+  }
+
   // use_per_session_threads
   model_session_config->mutable_session_config()
       ->set_use_per_session_threads(
@@ -366,7 +375,7 @@ Status Server::BuildAndStart(const Options& server_options) {
   if (server_options.model_config_file.empty()) {
     options.model_server_config = BuildSingleModelConfig(
         server_options.model_name, server_options.model_base_path);
-    use_session_group = server_options.session_num_per_group > 1;
+    use_session_group = server_options.session_num_per_group > 0;
   } else {
     TF_RETURN_IF_ERROR(ParseProtoTextFile<ModelServerConfig>(
         server_options.model_config_file, &options.model_server_config));

tensorflow_serving/model_servers/server.h

@@ -88,6 +88,7 @@ class Server {
     bool use_per_session_threads = false;
     bool use_session_group = false;
     tensorflow::int32 session_num_per_group = 0;
+    tensorflow::string gpu_ids_list = "";
     bool use_multi_stream = false;
 
     Options();

tensorflow_serving/servables/tensorflow/bundle_factory_util.cc

@@ -84,16 +84,18 @@ SessionGroupOptions GetSessionOptions(const SessionGroupBundleConfig& config, in
                << config.model_session_config_size();
   }
   options.config = config.model_session_config()[model_id].session_config();
-  options.metadata.session_num = config.model_session_config()[model_id].session_num();
-  if (options.metadata.session_num == 0) {
+  options.metadata.session_count = config.model_session_config()[model_id].session_num();
+  if (options.metadata.session_count == 0) {
     LOG(WARNING) << "User set use_session_group=true, but the #" << model_id
                  << " model config don't contain session_num field, "
                  << "please check platform_config_file config file. "
                  << "Now use default value 1.";
-    options.metadata.session_num = 1;
+    options.metadata.session_count = 1;
   }
-  for (auto& conf : config.model_session_config()) {
-    options.metadata.streams_vec.emplace_back(conf.session_num());
+  if (!config.model_session_config()[model_id].gpu_ids().empty()) {
+    for (auto id : config.model_session_config()[model_id].gpu_ids()) {
+      options.metadata.gpu_ids.emplace_back(id);
+    }
   }
   options.metadata.model_id = model_id;
   options.metadata.cpusets = config.model_session_config()[model_id].cpusets();

tensorflow_serving/servables/tensorflow/session_bundle_config.proto

@@ -170,13 +170,13 @@ message SessionGroupBundleConfig {
     // See details at tensorflow/core/protobuf/config.proto.
     ConfigProto session_config = 1;
     // session_num of the session group.
-    int32 session_num = 2;
-    // gpu num
-    //int32 gpu_num = 3;
+    int32 session_num = 2;// [default = 1];
+    // gpu list for current session_group
+    repeated int32 gpu_ids = 3;
     // model id in multi-models
-    int32 model_id = 3;
+    int32 model_id = 4;
     // cpusets
-    string cpusets = 4;
+    string cpusets = 5;
   }
 
   repeated ModelSessionConfig model_session_config = 784;