experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Lazy load estimator instead of using virtual pip package. RFC: https://github.com/tensorflow/community/pull/182

Browse Source 
PiperOrigin-RevId: 294847967
Change-Id: I327d075a2065e2ccf8ad5317882ebde14e3dc3d6
This commit is contained in:
Anna R 2020-02-12 23:49:52 -08:00 committed by TensorFlower Gardener
parent d0322efb28
commit 5c00e793c6
7 changed files with 107 additions and 59 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/BUILD
View File

@ -915,7 +915,6 @@ py_library(
"//conditions:default": [":tf_python_api_gen_v1"],
}) + [
":root_init_gen",
":virtual_root_init_gen",
"//tensorflow/python/keras/api:keras_python_api_gen",
"//tensorflow/python/keras/api:keras_python_api_gen_compat_v1",
"//tensorflow/python/keras/api:keras_python_api_gen_compat_v2",

21
tensorflow/api_template.__init__.py
View File

@ -35,9 +35,11 @@ import inspect as _inspect
import logging as _logging
import os as _os
import site as _site
import six as _six
import sys as _sys
from tensorflow.python.tools import module_util as _module_util
from tensorflow.python.util.lazy_loader import LazyLoader as _LazyLoader
# API IMPORTS PLACEHOLDER
@ -69,13 +71,13 @@ except ImportError:
_logging.warning(
"Limited tf.summary API due to missing TensorBoard installation.")
try:
from tensorflow_estimator.python.estimator.api._v2 import estimator
_current_module.__path__ = (
[_module_util.get_parent_dir(estimator)] + _current_module.__path__)
# Lazy-load estimator.
_estimator_module = "tensorflow_estimator.python.estimator.api._v2.estimator"
estimator = _LazyLoader("estimator", globals(), _estimator_module)
_module_dir = _module_util.get_parent_dir_for_name(_estimator_module)
if _module_dir:
_current_module.__path__ = [_module_dir] + _current_module.__path__
setattr(_current_module, "estimator", estimator)
except ImportError:
pass
try:
from .python.keras.api._v2 import keras
@ -85,6 +87,13 @@ try:
except ImportError:
pass
# Explicitly import lazy-loaded modules to support autocompletion.
# pylint: disable=g-import-not-at-top
if not _six.PY2:
import typing as _typing
if _typing.TYPE_CHECKING:
from tensorflow_estimator.python.estimator.api._v2 import estimator
# pylint: enable=g-import-not-at-top
# Enable TF2 behaviors
from tensorflow.python.compat import v2_compat as _compat # pylint: disable=g-import-not-at-top

22
tensorflow/api_template_v1.__init__.py
View File

@ -22,12 +22,14 @@ import distutils as _distutils
import inspect as _inspect
import os as _os
import site as _site
import six as _six
import sys as _sys
# pylint: disable=g-bad-import-order
from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import
from tensorflow.python.tools import module_util as _module_util
from tensorflow.python.platform import tf_logging as _logging
from tensorflow.python.util.lazy_loader import LazyLoader as _LazyLoader
# API IMPORTS PLACEHOLDER
@ -64,13 +66,14 @@ elif _tf_api_dir not in __path__:
# reexport_tf_summary can get compat from sys.modules. Only needed if using
# lazy loading.
_current_module.compat.v2 # pylint: disable=pointless-statement
try:
from tensorflow_estimator.python.estimator.api._v1 import estimator
_current_module.__path__ = (
[_module_util.get_parent_dir(estimator)] + _current_module.__path__)
# Lazy-load estimator.
_estimator_module = "tensorflow_estimator.python.estimator.api._v1.estimator"
estimator = _LazyLoader("estimator", globals(), _estimator_module)
_module_dir = _module_util.get_parent_dir_for_name(_estimator_module)
if _module_dir:
_current_module.__path__ = [_module_dir] + _current_module.__path__
setattr(_current_module, "estimator", estimator)
except ImportError:
pass
try:
from .python.keras.api._v1 import keras
@ -80,6 +83,13 @@ try:
except ImportError:
pass
# Explicitly import lazy-loaded modules to support autocompletion.
# pylint: disable=g-import-not-at-top
if not _six.PY2:
import typing as _typing
if _typing.TYPE_CHECKING:
from tensorflow_estimator.python.estimator.api._v1 import estimator
# pylint: enable=g-import-not-at-top
from tensorflow.python.util.lazy_loader import LazyLoader # pylint: disable=g-import-not-at-top
_CONTRIB_WARNING = """

22
tensorflow/compat_template.__init__.py
View File

@ -20,9 +20,11 @@ from __future__ import print_function as _print_function
import logging as _logging
import os as _os
import six as _six
import sys as _sys
from tensorflow.python.tools import module_util as _module_util
from tensorflow.python.util.lazy_loader import LazyLoader as _LazyLoader
# pylint: disable=g-bad-import-order
@ -36,20 +38,19 @@ try:
from tensorboard.summary._tf import summary
_current_module.__path__ = (
[_module_util.get_parent_dir(summary)] + _current_module.__path__)
# Make sure we get the correct summary module with lazy loading
setattr(_current_module, "summary", summary)
except ImportError:
_logging.warning(
"Limited tf.compat.v2.summary API due to missing TensorBoard "
"installation.")
try:
from tensorflow_estimator.python.estimator.api._v2 import estimator
_current_module.__path__ = (
[_module_util.get_parent_dir(estimator)] + _current_module.__path__)
# Lazy-load estimator.
_estimator_module = "tensorflow_estimator.python.estimator.api._v2.estimator"
estimator = _LazyLoader("estimator", globals(), _estimator_module)
_module_dir = _module_util.get_parent_dir_for_name(_estimator_module)
if _module_dir:
_current_module.__path__ = [_module_dir] + _current_module.__path__
setattr(_current_module, "estimator", estimator)
except ImportError:
pass
try:
from tensorflow.python.keras.api._v2 import keras
@ -59,6 +60,13 @@ try:
except ImportError:
pass
# Explicitly import lazy-loaded modules to support autocompletion.
# pylint: disable=g-import-not-at-top
if not _six.PY2:
import typing as _typing
if _typing.TYPE_CHECKING:
from tensorflow_estimator.python.estimator.api._v2 import estimator
# pylint: enable=g-import-not-at-top
# We would like the following to work for fully enabling 2.0 in a 1.0 install:
#

23
tensorflow/compat_template_v1.__init__.py
View File

@ -20,8 +20,10 @@ from __future__ import print_function as _print_function
import os as _os
import sys as _sys
import six as _six
from tensorflow.python.tools import module_util as _module_util
from tensorflow.python.util.lazy_loader import LazyLoader as _LazyLoader
# pylint: disable=g-bad-import-order
@ -31,13 +33,14 @@ from tensorflow.python.tools import module_util as _module_util
# Hook external TensorFlow modules.
_current_module = _sys.modules[__name__]
try:
from tensorflow_estimator.python.estimator.api._v1 import estimator
_current_module.__path__ = (
[_module_util.get_parent_dir(estimator)] + _current_module.__path__)
# Lazy-load estimator.
_estimator_module = "tensorflow_estimator.python.estimator.api._v1.estimator"
estimator = _LazyLoader("estimator", globals(), _estimator_module)
_module_dir = _module_util.get_parent_dir_for_name(_estimator_module)
if _module_dir:
_current_module.__path__ = [_module_dir] + _current_module.__path__
setattr(_current_module, "estimator", estimator)
except ImportError:
pass
try:
from tensorflow.python.keras.api._v1 import keras
@ -47,6 +50,14 @@ try:
except ImportError:
pass
# Explicitly import lazy-loaded modules to support autocompletion.
# pylint: disable=g-import-not-at-top
if not _six.PY2:
import typing as _typing
if _typing.TYPE_CHECKING:
from tensorflow_estimator.python.estimator.api._v1 import estimator
# pylint: enable=g-import-not-at-top
from tensorflow.python.platform import flags # pylint: disable=g-import-not-at-top
_current_module.app.flags = flags # pylint: disable=undefined-variable

40
tensorflow/python/tools/module_util.py
View File

@ -18,8 +18,48 @@ from __future__ import division
from __future__ import print_function
import os
import six
if six.PY2:
import imp # pylint: disable=g-import-not-at-top
else:
import importlib # pylint: disable=g-import-not-at-top
def get_parent_dir(module):
return os.path.abspath(os.path.join(os.path.dirname(module.__file__), ".."))
def get_parent_dir_for_name(module_name):
"""Get parent directory for module with the given name.
Args:
module_name: Module name for e.g.
tensorflow_estimator.python.estimator.api._v1.estimator.
Returns:
Path to the parent directory if module is found and None otherwise.
Given example above, it should return:
/pathtoestimator/tensorflow_estimator/python/estimator/api/_v1.
"""
name_split = module_name.split(".")
if not name_split:
return None
if six.PY2:
try:
spec = imp.find_module(name_split[0])
except ImportError:
return None
if not spec:
return None
base_path = spec[1]
else:
try:
spec = importlib.util.find_spec(name_split[0])
except ValueError:
return None
if not spec.origin:
return None
base_path = os.path.dirname(spec.origin)
return os.path.join(base_path, *name_split[1:-1])

29
tensorflow/tools/pip_package/build_pip_package.sh
View File

@ -165,35 +165,6 @@ function prepare_src() {
rm -f ${TMPDIR}/tensorflow/libtensorflow_framework.so
rm -f ${TMPDIR}/tensorflow/libtensorflow_framework.so.[0-9].*
# In order to break the circular dependency between tensorflow and
# tensorflow_estimator which forces us to do a multi-step release, we are
# creating a virtual python package called tensorflow and moving all the tf
# code into another python package called tensorflow_core:
#
# * move code from tensorflow to tensorflow_core
# * create the virtual pip package: create folder and __init__.py file with
# needed code for transparent forwarding
#
# This is transparent to internal code or to code not using the pip packages.
mv "${TMPDIR}/tensorflow" "${TMPDIR}/tensorflow_core"
mkdir "${TMPDIR}/tensorflow"
mv "${TMPDIR}/tensorflow_core/virtual_root.__init__.py" "${TMPDIR}/tensorflow/__init__.py"
# In V1 API, we need to remove deprecation warnings from
# ${TMPDIR}/tensorflow_core/__init__.py as those exists in
# ${TMPDIR}/tensorflow/__init__.py which does an import* and if these don't
# get removed users get 6 deprecation warning just on
#
# import tensorflow as tf
#
# which is not ok. We disable deprecation by using sed to toggle the flag
# TODO(mihaimaruseac): When we move the API to root, remove this hack
# Note: Can't do in place sed that works on all OS, so use a temp file instead
sed \
"s/deprecation=True/deprecation=False/g" \
"${TMPDIR}/tensorflow_core/__init__.py" > "${TMPDIR}/tensorflow_core/__init__.out"
mv "${TMPDIR}/tensorflow_core/__init__.out" "${TMPDIR}/tensorflow_core/__init__.py"
}
function build_wheel() {