Remove RTLD_GLOBAL when loading pywrap_tensorflow

Splits out a shared object (//tensorflow/libtensorflow_framework.so) with core TensorFlow functionality but neither ops nor kernels. This object does include registries for ops, kernels, filesystems, etc. The expectation is that shared objects containing custom ops will have a runtime dependency on this framework shared object: TensorFlow will load the custom op shared object, and the custom op shared object will use the symbols from the framework shared object to register its ops/kernels/etc. rather than (as before this change) relying on those symbols being in the global symbol table.

In this mode, TensorFlow artifacts (_pywrap_tensorflow.so for Python, libtensorflow.so for the C API; currently excluding Android artifacts) will depend on the framework shared object, which will be packaged with the Python pip package and other language distributions. This means that custom ops targeting the framework shared object will work in any language (C++, Java, Go; previously custom ops in these languages required custom Bazel builds).

Adds a config option which reproduces the old behavior (--config=monolithic), which for Python means building a monolithic pywrap_tensorflow shared object and loading its symbols into the global symbol table (with RTLD_GLOBAL). As before, there will be no extra-Bazel custom op support for other languages when compiling in this mode.

Does not change behavior on Windows; the cmake build is still monolithic.

Requires using tf_cc_binary, tf_cc_test, and (rarely) tf_cc_shared_object rules to link in the framework shared object when adding new TensorFlow build rules.

PiperOrigin-RevId: 169572746

1 files changed, 27 insertions, 14 deletions

diff --git a/tensorflow/python/pywrap_tensorflow.py b/tensorflow/python/pywrap_tensorflow.py
index 48ae17ca94..000ed8df8b 100644
--- a/tensorflow/python/pywrap_tensorflow.py
+++ b/tensorflow/python/pywrap_tensorflow.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =============================================================================
-"""pywrap_tensorflow wrapper that exports all symbols with RTLD_GLOBAL."""
+"""A wrapper for TensorFlow SWIG-generated bindings."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -31,25 +31,38 @@ 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 SWIG-generated
-# python library that dynamically loads _pywrap_tensorflow.so. The
-# default mode for loading keeps all the symbol private and not
-# visible to other libraries that may be loaded. Setting the mode to
-# RTLD_GLOBAL to make the symbols visible, so that custom op libraries
-# imported using `tf.load_op_library()` can access symbols defined in
-# _pywrap_tensorflow.so.
+# 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:
-  # TODO(keveman,mrry): Support dynamic op loading on platforms that do not
-  # use `dlopen()` for dynamic loading.
-  _use_rtld_global = hasattr(sys, 'getdlopenflags') and hasattr(sys, 'setdlopenflags')
-  if _use_rtld_global:
-    _default_dlopen_flags = sys.getdlopenflags()
-    sys.setdlopenflags(_default_dlopen_flags | ctypes.RTLD_GLOBAL)
+  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)
+
   from tensorflow.python.pywrap_tensorflow_internal import *
   from tensorflow.python.pywrap_tensorflow_internal import __version__
   from tensorflow.python.pywrap_tensorflow_internal import __git_version__
   from tensorflow.python.pywrap_tensorflow_internal import __compiler_version__
-  if _use_rtld_global:
+
+  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