Fixes tensorflow_serving continuous build.

PiperOrigin-RevId: 615451402
tensorflow · Mar 13, 2024 · a99fb9c · a99fb9c
1 parent f82600a
commit a99fb9c
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/tensorflow_serving/model_servers/BUILD b/tensorflow_serving/model_servers/BUILD
@@ -219,7 +219,8 @@ cc_library(
         "//tensorflow_serving/apis:model_service_cc_proto",
         "//tensorflow_serving/util:status_util",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@org_tensorflow//tensorflow/core:lib",
+        "@local_tsl//tsl/lib/monitoring:collected_metrics",
+        "@local_tsl//tsl/lib/monitoring:collection_registry",
     ],
 )
 

diff --git a/tensorflow_serving/model_servers/model_service_impl.cc b/tensorflow_serving/model_servers/model_service_impl.cc
@@ -92,9 +92,10 @@ absl::flat_hash_map<std::string, int64_t> ModelServiceImpl::GetMetrics(
 
   for (const std::string &metric_name : request->metric_names()) {
     int64_t metric_value = 0;
-    if (collected_metrics->point_set_map.contains(metric_name)) {
+    auto it = collected_metrics->point_set_map.find(metric_name);
+    if (it != collected_metrics->point_set_map.end()) {
       std::vector<std::unique_ptr<tsl::monitoring::Point>> *points =
-          &collected_metrics->point_set_map[metric_name]->points;
+          &it->second->points;
       if (!points->empty()) {
         metric_value = (*points)[0]->int64_value;
       }

diff --git a/tensorflow_serving/model_servers/model_service_impl.h b/tensorflow_serving/model_servers/model_service_impl.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "grpcpp/server_context.h"
 #include "grpcpp/support/status.h"
+#include "absl/container/flat_hash_map.h"
 #include "tensorflow_serving/apis/model_management.pb.h"
 #include "tensorflow_serving/apis/model_service.grpc.pb.h"
 #include "tensorflow_serving/apis/model_service.pb.h"