Allow localized disabling of tests for asan/msan/tsan/ubsan.

Currently we can disable entire test targets via the bazel tags `noasan`,
`nomsan`, `notsan`, `noubsan`.  However, in some instances, we may expect
one particular edge-case to fail a sanitizer test (for known reason), but
still want the rest of the test suite to be analyzed. By disabling individual
test methods or parts thereof locally, we allow for this without the need to
separate the offending case into a different target.

There already was some functionality in tflite, which seems to currently
be unused (though has been previously to locally skip tests).
Refactored and moved to tf so can be more widely used.

Also added corresponding defines to the `.bazelrc` file.

PiperOrigin-RevId: 374231232
Change-Id: I62ca36ebe9edec18e0d6d358496b6d07efb43574

.bazelrc

@@ -599,7 +599,7 @@ build:asan --linkopt -fsanitize=address
 # CC=clang bazel build --config msan
 build:msan --strip=never
 build:msan --copt -fsanitize=memory
-build:msan --copt -DADDRESS_SANITIZER
+build:msan --copt -DMEMORY_SANITIZER
 build:msan --copt -g
 build:msan --copt -O3
 build:msan --copt -fno-omit-frame-pointer
@@ -609,6 +609,7 @@ build:msan --linkopt -fsanitize=memory
 # CC=clang bazel build --config ubsan
 build:ubsan --strip=never
 build:ubsan --copt -fsanitize=undefined
+build:msan --copt -DUNDEFINED_BEHAVIOR_SANITIZER
 build:ubsan --copt -g
 build:ubsan --copt -O3
 build:ubsan --copt -fno-omit-frame-pointer

tensorflow/lite/tools/sanitizers/BUILD

@@ -1,21 +0,0 @@
-load("//tensorflow:tensorflow.bzl", "pybind_extension")
-
-package(
-    default_visibility = [
-        "//visibility:public",
-    ],
-    licenses = ["notice"],
-)
-
-pybind_extension(
-    name = "_pywrap_tensorflow_lite_sanitizers",
-    srcs = [
-        "sanitizers_pybind11.cc",
-    ],
-    link_in_framework = True,
-    module_name = "_pywrap_tensorflow_lite_sanitizers",
-    deps = [
-        "//tensorflow/python:pybind11_lib",
-        "@pybind11",
-    ],
-)

tensorflow/python/BUILD

@@ -4742,6 +4742,26 @@ tf_python_pybind_extension(
     ],
 )
 
+py_library(
+    name = "pywrap_sanitizers",
+    srcs = ["pywrap_sanitizers.py"],
+    srcs_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":_pywrap_sanitizers",
+        ":pywrap_tensorflow",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_sanitizers",
+    srcs = ["sanitizers_wrapper.cc"],
+    module_name = "_pywrap_sanitizers",
+    deps = [
+        "@pybind11",
+    ],
+)
+
 cc_library(
     name = "unified_api_pywrap_required_headers",
     textual_hdrs = [

tensorflow/python/framework/BUILD

@@ -1234,6 +1234,7 @@ py_library(
         "//tensorflow/python:client",
         "//tensorflow/python:platform",
         "//tensorflow/python:platform_test",
+        "//tensorflow/python:pywrap_sanitizers",
         "//tensorflow/python:pywrap_tf_session",
         "//tensorflow/python:resource_variable_ops",
         "//tensorflow/python:session",

tensorflow/python/framework/test_util.py

@@ -43,6 +43,7 @@ from google.protobuf import text_format
 
 from tensorflow.core.framework import graph_pb2
 from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python import pywrap_sanitizers
 from tensorflow.python import tf2
 from tensorflow.python.client import device_lib
 from tensorflow.python.client import pywrap_tf_session
@@ -120,6 +121,26 @@ except ImportError:
     pass
 
 
+def is_asan_enabled():
+  """Check if ASAN is enabled."""
+  return pywrap_sanitizers.is_asan_enabled()
+
+
+def is_msan_enabled():
+  """Check if MSAN is enabled."""
+  return pywrap_sanitizers.is_msan_enabled()
+
+
+def is_tsan_enabled():
+  """Check if TSAN is enabled."""
+  return pywrap_sanitizers.is_tsan_enabled()
+
+
+def is_ubsan_enabled():
+  """Check if UBSAN is enabled."""
+  return pywrap_sanitizers.is_ubsan_enabled()
+
+
 def _get_object_count_by_type(exclude=()):
   return (
       collections.Counter([type(obj).__name__ for obj in gc.get_objects()]) -
@@ -1855,6 +1876,34 @@ def disable_mlir_bridge(description):  # pylint: disable=unused-argument
   return _disable_test(execute_func)
 
 
+# The description is just for documentation purposes.
+def disable_asan(description):  # pylint: disable=unused-argument
+  """Execute the test method only if ASAN is not enabled."""
+  execute_func = not is_asan_enabled()
+  return _disable_test(execute_func)
+
+
+# The description is just for documentation purposes.
+def disable_msan(description):  # pylint: disable=unused-argument
+  """Execute the test method only if MSAN is not enabled."""
+  execute_func = not is_msan_enabled()
+  return _disable_test(execute_func)
+
+
+# The description is just for documentation purposes.
+def disable_tsan(description):  # pylint: disable=unused-argument
+  """Execute the test method only if TSAN is not enabled."""
+  execute_func = not is_tsan_enabled()
+  return _disable_test(execute_func)
+
+
+# The description is just for documentation purposes.
+def disable_ubsan(description):  # pylint: disable=unused-argument
+  """Execute the test method only if UBSAN is not enabled."""
+  execute_func = not is_ubsan_enabled()
+  return _disable_test(execute_func)
+
+
 # The description is just for documentation purposes.
 def disable_tfrt(unused_description):
 

tensorflow/python/framework/test_util_test.py

@@ -32,6 +32,7 @@ from google.protobuf import text_format
 
 from tensorflow.core.framework import graph_pb2
 from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python import pywrap_sanitizers
 from tensorflow.python.compat import compat
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
@@ -147,6 +148,22 @@ class TestUtilTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     else:
       print("MKL is disabled")
 
+  @test_util.disable_asan("Skip test if ASAN is enabled.")
+  def testDisableAsan(self):
+    self.assertFalse(pywrap_sanitizers.is_asan_enabled())
+
+  @test_util.disable_msan("Skip test if MSAN is enabled.")
+  def testDisableMsan(self):
+    self.assertFalse(pywrap_sanitizers.is_msan_enabled())
+
+  @test_util.disable_tsan("Skip test if TSAN is enabled.")
+  def testDisableTsan(self):
+    self.assertFalse(pywrap_sanitizers.is_tsan_enabled())
+
+  @test_util.disable_ubsan("Skip test if UBSAN is enabled.")
+  def testDisableUbsan(self):
+    self.assertFalse(pywrap_sanitizers.is_ubsan_enabled())
+
   @test_util.run_in_graph_and_eager_modes
   def testAssertProtoEqualsStr(self):
 

tensorflow/python/pywrap_sanitizers.py

@@ -0,0 +1,23 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Python module for sanitizer detection functions exported by pybind11."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# pylint: disable=invalid-import-order,g-bad-import-order, wildcard-import, unused-import
+from tensorflow.python import pywrap_tensorflow
+from tensorflow.python._pywrap_sanitizers import *

tensorflow/python/sanitizers_wrapper.cc

@@ -14,25 +14,35 @@ limitations under the License.
 ==============================================================================*/
 
 #include "pybind11/pybind11.h"
-#include "tensorflow/python/lib/core/pybind11_lib.h"
 
-PYBIND11_MODULE(_pywrap_tensorflow_lite_sanitizers, m) {
-  m.def("TSan_Enabled", []() -> bool {
-#ifdef THREAD_SANITIZER
+// Check if specific santizers are enabled.
+PYBIND11_MODULE(_pywrap_sanitizers, m) {
+  m.def("is_asan_enabled", []() -> bool {
+#if defined(ADDRESS_SANITIZER)
     return true;
 #else
     return false;
 #endif
   });
-  m.def("MSan_Enabled", []() -> bool {
-#ifdef MEMORY_SANITIZER
+
+  m.def("is_msan_enabled", []() -> bool {
+#if defined(MEMORY_SANITIZER)
     return true;
 #else
     return false;
 #endif
   });
-  m.def("ASan_Enabled", []() -> bool {
-#ifdef ADDRESS_SANITIZER
+
+  m.def("is_tsan_enabled", []() -> bool {
+#if defined(THREAD_SANITIZER)
+    return true;
+#else
+    return false;
+#endif
+  });
+
+  m.def("is_ubsan_enabled", []() -> bool {
+#if defined(UNDEFINED_BEHAVIOR_SANITIZER)
     return true;
 #else
     return false;