Use OSS-compatible logging lib directly

PiperOrigin-RevId: 553315782
tensorflow · Aug 3, 2023 · 3e8dd82 · 3e8dd82
1 parent 9c7bfd1
commit 3e8dd82
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/tensorflow_serving/servables/tensorflow/util.cc b/tensorflow_serving/servables/tensorflow/util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow_serving/servables/tensorflow/util.h"
 
 #include <atomic>
+#include <cstdlib>
 
 #include "google/protobuf/wrappers.pb.h"
 #include "tensorflow/cc/saved_model/signature_constants.h"
@@ -358,5 +359,10 @@ std::set<string> SetDifference(std::set<string> set_a, std::set<string> set_b) {
   return result;
 }
 
+bool IsTfrtErrorLoggingEnabled() {
+  const char* env = getenv("ENABLE_TFRT_SERVING_ERROR_LOGGING");
+  return env != nullptr && env[0] == '1' && env[1] == '\0';
+}
+
 }  // namespace serving
 }  // namespace tensorflow
diff --git a/tensorflow_serving/servables/tensorflow/util.h b/tensorflow_serving/servables/tensorflow/util.h
@@ -135,6 +135,9 @@ std::set<string> GetMapKeys(const T& map) {
 // Returns a \ b, i.e. all items that are in `set_a` but not in `set_b`.
 std::set<string> SetDifference(std::set<string> set_a, std::set<string> set_b);
 
+// Returns true if errors should be exported to logging service.
+bool IsTfrtErrorLoggingEnabled();
+
 }  // namespace serving
 }  // namespace tensorflow