Update TensorBoard to import pywrap_tensorflow from tensorflow.python (#2096)

* Update TensorBoard to import pywrap_tensorflow from tensorflow.python instead of tensorflow.compat.v1. We are removing pywrap_tensorflow from TF API including tf.compat.v1. Importing from tensorflow.python is a temporary solution until a better API for reading records is available for TensorBoard to use.

* Change call to NewCheckpointReader to tf.train.load_checkpoint instead. Also, added TODO

Author: annarev

tensorboard/backend/event_processing/event_file_loader.py

@@ -21,6 +21,7 @@
 import inspect
 
 from tensorboard.compat import tf
+from tensorboard.compat import _pywrap_tensorflow
 from tensorboard.compat.proto import event_pb2
 from tensorboard.util import platform_util
 from tensorboard.util import tb_logging
@@ -38,7 +39,7 @@ def __init__(self, file_path):
     file_path = platform_util.readahead_file_path(file_path)
     logger.debug('Opening a record reader pointing at %s', file_path)
     with tf.errors.raise_exception_on_not_ok_status() as status:
-      self._reader = tf.compat.v1.pywrap_tensorflow.PyRecordReader_New(
+      self._reader = _pywrap_tensorflow.PyRecordReader_New(
           tf.compat.as_bytes(file_path), 0, tf.compat.as_bytes(''), status)
     # Store it for logging purposes.
     self._file_path = file_path

tensorboard/compat/__init__.py

@@ -72,3 +72,34 @@ def tf2():
     # As a fallback, try `tensorflow.compat.v2` if it's defined.
     return tf.compat.v2
   raise ImportError('cannot import tensorflow 2.0 API')
+
+
+# TODO(https://github.com/tensorflow/tensorboard/issues/1711): remove this
+@_lazy.lazy_load('tensorboard.compat._pywrap_tensorflow')
+def _pywrap_tensorflow():
+  """Provide pywrap_tensorflow access in TensorBoard.
+
+  pywrap_tensorflow cannot be accessed from tf.python.pywrap_tensorflow
+  and needs to be imported using
+  `from tensorflow.python import pywrap_tensorflow`. Therefore, we provide
+  a separate accessor function for it here.
+
+  NOTE: pywrap_tensorflow is not part of TensorFlow API and this
+  dependency will go away soon.
+
+  Returns:
+    pywrap_tensorflow import, if available.
+
+  Raises:
+    ImportError: if we couldn't import pywrap_tensorflow.
+  """
+  try:
+    from tensorboard.compat import notf  # pylint: disable=g-import-not-at-top
+  except ImportError:
+    try:
+      from tensorflow.python import pywrap_tensorflow  # pylint: disable=g-import-not-at-top
+      return pywrap_tensorflow
+    except ImportError:
+      pass
+  from tensorboard.compat.tensorflow_stub import pywrap_tensorflow  # pylint: disable=g-import-not-at-top
+  return pywrap_tensorflow

tensorboard/compat/tensorflow_stub/__init__.py

@@ -35,8 +35,5 @@
 from . import pywrap_tensorflow  # noqa
 from . import tensor_shape  # noqa
 
-# Set pywrap_tensorflow on v1 and avoid cycles on some imports
-compat.v1.pywrap_tensorflow = pywrap_tensorflow
-
 # Set a fake __version__ to help distinguish this as our own stub API.
 __version__ = 'stub'

tensorboard/plugins/debugger/debugger_plugin_testlib.py

@@ -25,6 +25,7 @@
 
 import portpicker  # pylint: disable=import-error
 import tensorflow as tf
+from tensorflow.python import pywrap_tensorflow
 from werkzeug import wrappers
 from werkzeug import test as werkzeug_test
 
@@ -60,7 +61,7 @@ def setUp(self):
     self.log_dir = self.get_temp_dir()
     file_prefix = tf.compat.as_bytes(
         os.path.join(self.log_dir, 'events.debugger'))
-    writer = tf.compat.v1.pywrap_tensorflow.EventsWriter(file_prefix)
+    writer = pywrap_tensorflow.EventsWriter(file_prefix)
     device_name = '/job:localhost/replica:0/task:0/cpu:0'
     writer.WriteEvent(
         self._CreateEventWithDebugNumericSummary(
@@ -106,7 +107,7 @@ def setUp(self):
     os.mkdir(run_foo_directory)
     file_prefix = tf.compat.as_bytes(
         os.path.join(run_foo_directory, 'events.debugger'))
-    writer = tf.compat.v1.pywrap_tensorflow.EventsWriter(file_prefix)
+    writer = pywrap_tensorflow.EventsWriter(file_prefix)
     writer.WriteEvent(
         self._CreateEventWithDebugNumericSummary(
             device_name=device_name,

tensorboard/plugins/debugger/events_writer_manager.py

@@ -27,6 +27,7 @@
 import time
 
 import tensorflow as tf
+from tensorflow.python import pywrap_tensorflow
 from tensorboard.util import tb_logging
 
 logger = tb_logging.get_logger()
@@ -208,7 +209,7 @@ def _create_events_writer(self, directory):
         os.path.join(directory, DEBUGGER_EVENTS_FILE_STARTING_TEXT),
         time.time(), self._events_file_count)
     logger.info("Creating events file %s", file_path)
-    return tf.compat.v1.pywrap_tensorflow.EventsWriter(tf.compat.as_bytes(file_path))
+    return pywrap_tensorflow.EventsWriter(tf.compat.as_bytes(file_path))
 
   def _fetch_events_files_on_disk(self):
     """Obtains the names of debugger-related events files within the directory.

tensorboard/plugins/projector/projector_plugin.py

@@ -32,6 +32,7 @@
 
 from tensorboard.backend.http_util import Respond
 from tensorboard.compat import tf
+from tensorboard.compat import _pywrap_tensorflow
 from tensorboard.plugins import base_plugin
 from tensorboard.plugins.projector.projector_config_pb2 import ProjectorConfig
 from tensorboard.util import tb_logging
@@ -423,8 +424,7 @@ def _get_reader_for_run(self, run):
     reader = None
     if config.model_checkpoint_path and _using_tf():
       try:
-        reader = tf.compat.v1.pywrap_tensorflow.NewCheckpointReader(
-            config.model_checkpoint_path)
+        reader = tf.train.load_checkpoint(config.model_checkpoint_path)
       except Exception:  # pylint: disable=broad-except
         logger.warn('Failed reading "%s"', config.model_checkpoint_path)
     self.readers[run] = reader