[DeepRec] Support multi-model deployment in SessionGroup. (#13)

Author: shanshanpt

WORKSPACE

@@ -11,8 +11,8 @@ load("//tensorflow_serving:repo.bzl", "tensorflow_http_archive")
 
 tensorflow_http_archive(
     name = "org_tensorflow",
-    sha256 = "46ad5154cec11995d5feba1401b6d0d72be457e48a9776bfce04b13c009ac412",
-    git_commit = "0fe26688a57eee31bda56a57d8f05e7071c78c9b",
+    sha256 = "95c2e401a57024a57fcf757498f0e962519e6967381b630d6afc3a4a80c7ee37",
+    git_commit = "4b8c11fa0c72e1b3483ef8b1960c4790dda0e437",
 )
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")

tensorflow_serving/config/model_server_config.proto

@@ -65,6 +65,9 @@ message ModelConfig {
   //
   // (This can be changed once a model is in serving.)
   LoggingConfig logging_config = 6;
+
+  // model_id in multi-models.
+  int32 model_id = 9;
 }
 
 // Static list of models to be loaded for serving.

tensorflow_serving/core/servable_data.h

@@ -65,12 +65,16 @@ class ServableData {
   // !this->ok().
   T ConsumeDataOrDie();
 
+  void SetModelId(int id) { model_id_ = id; }
+  int GetModelId() const { return model_id_; }
+
  private:
   ServableData() = delete;
 
   const ServableId id_;
   const Status status_;
   T data_;
+  int model_id_ = 0;
 };
 
 // Helper static method to create a ServableData object. Caller may skip

tensorflow_serving/core/source_adapter.h

@@ -135,6 +135,10 @@ class UnarySourceAdapter : public SourceAdapter<InputType, OutputType> {
   // Converts a single InputType instance into a corresponding OutputType
   // instance.
   virtual Status Convert(const InputType& data, OutputType* converted_data) = 0;
+  virtual Status Convert(const InputType& data, int model_id,
+      OutputType* converted_data) {
+    return Convert(data, converted_data);
+  }
 };
 
 // A source adapter that converts every incoming ServableData<InputType> item
@@ -210,7 +214,8 @@ UnarySourceAdapter<InputType, OutputType>::Adapt(
   for (const ServableData<InputType>& version : versions) {
     if (version.status().ok()) {
       OutputType adapted_data;
-      Status adapt_status = Convert(version.DataOrDie(), &adapted_data);
+      Status adapt_status = Convert(version.DataOrDie(),
+          version.GetModelId(), &adapted_data);
       if (adapt_status.ok()) {
         adapted_versions.emplace_back(
             ServableData<OutputType>{version.id(), std::move(adapted_data)});

tensorflow_serving/core/test_util/session_test_util.cc

@@ -62,7 +62,7 @@ class DelegatingSessionFactory : public SessionFactory {
 
   Status NewSessionGroup(const SessionOptions& options,
                          SessionGroup** out_session_group,
-                         int session_num = 1) {
+                         const SessionGroupMetadata& metadata) {
     return errors::Internal(
         "NewSessionGroup method not implemented in DelegatingSessionFactory.");
   }

tensorflow_serving/example/config_files/README

@@ -0,0 +1,3 @@
+Usage:
+
+CUDA_VISIBLE_DEVICES=0,1 tensorflow_model_server --use_session_group=true --model_config_file=session_group_multi_models_config --platform_config_file=session_group_multi_models_platform_config

tensorflow_serving/example/config_files/models.config

@@ -0,0 +1,12 @@
+model_config_list:{
+    config:{
+      name:"pb1",
+      base_path:"/data/workspace/serving-model/multi_wdl_model/pb1",
+      model_platform:"tensorflow"
+    },
+    config:{
+      name:"pb2",
+      base_path:"/data/workspace/serving-model/multi_wdl_model/pb2",
+      model_platform:"tensorflow"
+    },
+}

tensorflow_serving/example/config_files/platform_config

@@ -0,0 +1,17 @@
+platform_configs {
+  key: "tensorflow"
+  value {
+    source_adapter_config {
+      [type.googleapis.com/tensorflow.serving.SavedModelBundleSourceAdapterConfig] {
+        legacy_config {
+          session_config {
+            gpu_options {
+              allow_growth: true
+            }
+          }
+        }
+      }
+    }
+  }
+}
+

tensorflow_serving/example/config_files/session_group_multi_models_config

@@ -0,0 +1,14 @@
+model_config_list:{
+    config:{
+      name:"pb1",
+      base_path:"/data/workspace/serving-model/multi_wdl_model/pb1",
+      model_platform:"tensorflow",
+      model_id: 0
+    },
+    config:{
+      name:"pb2",
+      base_path:"/data/workspace/serving-model/multi_wdl_model/pb2",
+      model_platform:"tensorflow",
+      model_id: 1
+    },
+}

tensorflow_serving/example/config_files/session_group_multi_models_platform_config

@@ -0,0 +1,36 @@
+platform_configs {
+  key: "tensorflow"
+  value {
+    source_adapter_config {
+      [type.googleapis.com/tensorflow.serving.SavedModelBundleV2SourceAdapterConfig] {
+        legacy_config {
+          model_session_config {
+            session_config {
+              gpu_options {
+                allow_growth: true
+              }
+              intra_op_parallelism_threads: 8
+              inter_op_parallelism_threads: 8
+              use_per_session_threads: true
+              use_per_session_stream: true
+            }
+            session_num: 2
+          }
+          model_session_config {
+            session_config {
+              gpu_options {
+                allow_growth: true
+              }
+              intra_op_parallelism_threads: 16
+              inter_op_parallelism_threads: 16
+              use_per_session_threads: true
+              use_per_session_stream: true
+            }
+            session_num: 4
+          }
+        }
+      }
+    }
+  }
+}
+