[DeepRec] Support GetModelMetadata request in session_group. (#12)

Author: shanshanpt

tensorflow_serving/model_servers/http_rest_api_handler.cc

@@ -270,9 +270,11 @@ Status HttpRestApiHandler::ProcessModelMetadataRequest(
     request.mutable_model_spec()->mutable_version()->set_value(version);
   }
 
+  ModelMetaOption opt;
+  opt.use_session_group = use_session_group_;
   GetModelMetadataResponse response;
   TF_RETURN_IF_ERROR(
-      GetModelMetadataImpl::GetModelMetadata(core_, request, &response));
+      GetModelMetadataImpl::GetModelMetadata(core_, request, &response, opt));
   JsonPrintOptions opts;
   opts.add_whitespace = true;
   opts.always_print_primitive_fields = true;

tensorflow_serving/model_servers/prediction_service_impl.cc

@@ -61,8 +61,10 @@ ::grpc::Status PredictionServiceImpl::GetModelMetadata(
         errors::InvalidArgument("GetModelMetadata API is only available when "
                                 "use_saved_model is set to true"));
   }
+  ModelMetaOption opt;
+  opt.use_session_group = use_session_group_;
   const ::grpc::Status status = ToGRPCStatus(
-      GetModelMetadataImpl::GetModelMetadata(core_, *request, response));
+      GetModelMetadataImpl::GetModelMetadata(core_, *request, response, opt));
   if (!status.ok()) {
     VLOG(1) << "GetModelMetadata failed: " << status.error_message();
   }

tensorflow_serving/servables/tensorflow/get_model_metadata_impl.cc

@@ -49,17 +49,34 @@ Status ValidateGetModelMetadataRequest(const GetModelMetadataRequest& request) {
 
 Status SavedModelGetSignatureDef(ServerCore* core, const ModelSpec& model_spec,
                                  const GetModelMetadataRequest& request,
-                                 GetModelMetadataResponse* response) {
-  ServableHandle<SavedModelBundle> bundle;
-  TF_RETURN_IF_ERROR(core->GetServableHandle(model_spec, &bundle));
+                                 GetModelMetadataResponse* response,
+                                 ModelMetaOption opt) {
   SignatureDefMap signature_def_map;
-  for (const auto& signature : bundle->meta_graph_def.signature_def()) {
-    (*signature_def_map.mutable_signature_def())[signature.first] =
-        signature.second;
+  std::string model_name;
+  int64 version;
+  if (opt.use_session_group) {
+    ServableHandle<SavedModelBundleV2> bundle;
+    TF_RETURN_IF_ERROR(core->GetServableHandle(model_spec, &bundle));
+    for (const auto& signature : bundle->meta_graph_def.signature_def()) {
+      (*signature_def_map.mutable_signature_def())[signature.first] =
+          signature.second;
+    }
+    model_name = bundle.id().name;
+    version = bundle.id().version;
+  } else {
+    ServableHandle<SavedModelBundle> bundle;
+    TF_RETURN_IF_ERROR(core->GetServableHandle(model_spec, &bundle));
+    for (const auto& signature : bundle->meta_graph_def.signature_def()) {
+      (*signature_def_map.mutable_signature_def())[signature.first] =
+          signature.second;
+    }
+    model_name = bundle.id().name;
+    version = bundle.id().version;
   }
+
   auto response_model_spec = response->mutable_model_spec();
-  response_model_spec->set_name(bundle.id().name);
-  response_model_spec->mutable_version()->set_value(bundle.id().version);
+  response_model_spec->set_name(model_name);
+  response_model_spec->mutable_version()->set_value(version);
 
   (*response->mutable_metadata())[GetModelMetadataImpl::kSignatureDef].PackFrom(
       signature_def_map);
@@ -72,24 +89,26 @@ constexpr const char GetModelMetadataImpl::kSignatureDef[];
 
 Status GetModelMetadataImpl::GetModelMetadata(
     ServerCore* core, const GetModelMetadataRequest& request,
-    GetModelMetadataResponse* response) {
+    GetModelMetadataResponse* response,
+    ModelMetaOption opt) {
   if (!request.has_model_spec()) {
     return tensorflow::Status(tensorflow::error::INVALID_ARGUMENT,
                               "Missing ModelSpec");
   }
   return GetModelMetadataWithModelSpec(core, request.model_spec(), request,
-                                       response);
+                                       response, opt);
 }
 
 Status GetModelMetadataImpl::GetModelMetadataWithModelSpec(
     ServerCore* core, const ModelSpec& model_spec,
     const GetModelMetadataRequest& request,
-    GetModelMetadataResponse* response) {
+    GetModelMetadataResponse* response,
+    ModelMetaOption opt) {
   TF_RETURN_IF_ERROR(ValidateGetModelMetadataRequest(request));
   for (const auto& metadata_field : request.metadata_field()) {
     if (metadata_field == kSignatureDef) {
       TF_RETURN_IF_ERROR(
-          SavedModelGetSignatureDef(core, model_spec, request, response));
+          SavedModelGetSignatureDef(core, model_spec, request, response, opt));
     } else {
       return tensorflow::errors::InvalidArgument(
           "MetadataField %s is not supported", metadata_field);

tensorflow_serving/servables/tensorflow/get_model_metadata_impl.h

@@ -23,20 +23,26 @@ limitations under the License.
 namespace tensorflow {
 namespace serving {
 
+struct ModelMetaOption {
+  bool use_session_group = false;
+};
+
 class GetModelMetadataImpl {
  public:
   static constexpr const char kSignatureDef[] = "signature_def";
 
   static Status GetModelMetadata(ServerCore* core,
                                  const GetModelMetadataRequest& request,
-                                 GetModelMetadataResponse* response);
+                                 GetModelMetadataResponse* response,
+                                 ModelMetaOption opt = ModelMetaOption());
 
   // Like GetModelMetadata(), but uses 'model_spec' instead of the one embedded
   // in 'request'.
   static Status GetModelMetadataWithModelSpec(
       ServerCore* core, const ModelSpec& model_spec,
       const GetModelMetadataRequest& request,
-      GetModelMetadataResponse* response);
+      GetModelMetadataResponse* response,
+      ModelMetaOption opt = ModelMetaOption());
 };
 
 }  // namespace serving