Re-add dynamic loading flags and logic in Python.

PiperOrigin-RevId: 305712320
Change-Id: Iaa25ab0397e91ffa83cba805e81a6c9e315549d8

tensorflow/python/BUILD

@@ -5828,7 +5828,11 @@ py_library(
     name = "pywrap_tensorflow",
     srcs = [
         "pywrap_tensorflow.py",
-    ],
+    ] + if_static(
+        ["pywrap_dlopen_global_flags.py"],
+        # Import will fail, indicating no global dlopen flags
+        otherwise = [],
+    ),  # b/153585257
     srcs_version = "PY2AND3",
     deps = [":pywrap_tensorflow_internal"],
 )

tensorflow/python/pywrap_tensorflow.py

@@ -18,6 +18,41 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import ctypes
+import sys
+import traceback
+
+from tensorflow.python.platform import self_check
+
+# Perform pre-load sanity checks in order to produce a more actionable error.
+self_check.preload_check()
+
+# pylint: disable=wildcard-import,g-import-not-at-top,unused-import,line-too-long
+
+try:
+  # This import is expected to fail if there is an explicit shared object
+  # dependency (with_framework_lib=true), since we do not need RTLD_GLOBAL.
+  from tensorflow.python import pywrap_dlopen_global_flags
+  _use_dlopen_global_flags = True
+except ImportError:
+  _use_dlopen_global_flags = False
+
+# On UNIX-based platforms, pywrap_tensorflow is a python library that
+# dynamically loads _pywrap_tensorflow.so.
+_can_set_rtld_local = (
+    hasattr(sys, 'getdlopenflags') and hasattr(sys, 'setdlopenflags'))
+if _can_set_rtld_local:
+  _default_dlopen_flags = sys.getdlopenflags()
+
+try:
+  if _use_dlopen_global_flags:
+    pywrap_dlopen_global_flags.set_dlopen_flags()
+  elif _can_set_rtld_local:
+    # Ensure RTLD_LOCAL behavior for platforms where it isn't the default
+    # (macOS). On Linux RTLD_LOCAL is 0, so this does nothing (and would not
+    # override an RTLD_GLOBAL in _default_dlopen_flags).
+    sys.setdlopenflags(_default_dlopen_flags | ctypes.RTLD_LOCAL)
+
   # Python2.7 does not have a ModuleNotFoundError.
   try:
     ModuleNotFoundError
@@ -35,3 +70,16 @@ try:
   # This logic is used in other internal projects using py_extension.
   except ModuleNotFoundError:
     pass
+
+  if _use_dlopen_global_flags:
+    pywrap_dlopen_global_flags.reset_dlopen_flags()
+  elif _can_set_rtld_local:
+    sys.setdlopenflags(_default_dlopen_flags)
+except ImportError:
+  msg = """%s\n\nFailed to load the native TensorFlow runtime.\n
+See https://www.tensorflow.org/install/errors\n
+for some common reasons and solutions.  Include the entire stack trace
+above this error message when asking for help.""" % traceback.format_exc()
+  raise ImportError(msg)
+
+# pylint: enable=wildcard-import,g-import-not-at-top,unused-import,line-too-long