[Debugger Plugin] Fix Python unit tests for TF2 (#2583)

* Motivation for features / changes
  * Re. #1705
  * Make sure all debugger plugin's python tests pass with TF2

* Technical description of changes
  * In session_debug_test.py and interactive_debugger_plugin_test.py: use tf.compat.v1 and disable_v2_behavior(), as tfdbg (V1) is indeed wedded to the Session::Run() paradigm. 
  * In debug_graphs_helper_test.py: make conditional assertions to fit v2-specific behaviors

Author: caisq

tensorboard/plugins/debugger/debug_graphs_helper_test.py

@@ -37,6 +37,8 @@
 
 import tensorflow as tf
 from tensorflow.python import debug as tf_debug
+# See discussion on issue #1996 for private module import justification.
+from tensorflow.python import tf2 as tensorflow_python_tf2
 from tensorflow.python.debug.lib import grpc_debug_test_server
 
 from tensorboard.compat.proto import config_pb2
@@ -85,7 +87,6 @@ def _createTestGraphAndRunOptions(self, sess, gated_grpc=True):
                          debug_urls=self.debug_server_url)
     return z, run_options
 
-  @test_util.run_v1_only('Ops differ. Similar to [1].')
   def testExtractGatedGrpcTensorsFoundGatedGrpcOps(self):
     with tf.compat.v1.Session() as sess:
       z, run_options = self._createTestGraphAndRunOptions(sess, gated_grpc=True)
@@ -108,15 +109,11 @@ def testExtractGatedGrpcTensorsFoundGatedGrpcOps(self):
       gated_debug_ops = [
           (item[0], item[2], item[3]) for item in gated_debug_ops]
 
-      # TODO(#1705): TF 2.0 breaks below.
       self.assertIn(('a', 0, 'DebugIdentity'), gated_debug_ops)
-      self.assertIn(('a/read', 0, 'DebugIdentity'), gated_debug_ops)
       self.assertIn(('b', 0, 'DebugIdentity'), gated_debug_ops)
-      self.assertIn(('b/read', 0, 'DebugIdentity'), gated_debug_ops)
       self.assertIn(('c', 0, 'DebugIdentity'), gated_debug_ops)
-      self.assertIn(('c/read', 0, 'DebugIdentity'), gated_debug_ops)
       self.assertIn(('d', 0, 'DebugIdentity'), gated_debug_ops)
-      self.assertIn(('d/read', 0, 'DebugIdentity'), gated_debug_ops)
+
       self.assertIn(('x', 0, 'DebugIdentity'), gated_debug_ops)
       self.assertIn(('y', 0, 'DebugIdentity'), gated_debug_ops)
       self.assertIn(('z', 0, 'DebugIdentity'), gated_debug_ops)
@@ -151,9 +148,6 @@ def testExtractGatedGrpcTensorsFoundNoGatedGrpcOps(self):
       self.assertEqual([], gated_debug_ops)
 
 
-@test_util.run_v1_only((
-    'Graph creates different op structure in v2. See '
-    'debug_graphs_helper_test.py[1].'))
 class BaseExpandedNodeNameTest(tf.test.TestCase):
 
   def testMaybeBaseExpandedNodeName(self):
@@ -174,9 +168,14 @@ def testMaybeBaseExpandedNodeName(self):
       self.assertEqual(
           'bar/b/read',
           graph_wrapper.maybe_base_expanded_node_name('bar/b/read'))
-      # TODO(#1705): TF 2.0 tf.add creates nested nodes.
-      self.assertEqual(
-          'baz/c', graph_wrapper.maybe_base_expanded_node_name('baz/c'))
+
+      if tensorflow_python_tf2.enabled():
+        # NOTE(#1705): TF 2.0 tf.add creates nested nodes.
+        self.assertEqual(
+            'baz/c/(c)', graph_wrapper.maybe_base_expanded_node_name('baz/c'))
+      else:
+        self.assertEqual(
+            'baz/c', graph_wrapper.maybe_base_expanded_node_name('baz/c'))
 
 
 if __name__ == '__main__':

tensorboard/plugins/debugger/interactive_debugger_plugin_test.py

@@ -33,7 +33,7 @@
 import numpy as np
 import portpicker  # pylint: disable=import-error
 from six.moves import urllib  # pylint: disable=wrong-import-order
-import tensorflow as tf  # pylint: disable=wrong-import-order
+import tensorflow.compat.v1 as tf  # pylint: disable=wrong-import-order
 from tensorflow.python import debug as tf_debug  # pylint: disable=wrong-import-order
 from werkzeug import test as werkzeug_test  # pylint: disable=wrong-import-order
 from werkzeug import wrappers  # pylint: disable=wrong-import-order
@@ -44,11 +44,14 @@
 from tensorboard.plugins.debugger import interactive_debugger_plugin
 from tensorboard.util import test_util
 
+# These unit tests for Debugger Plugin V1 are tied to TF1.x behavior
+# (`tf.Session`s).
+tf.disable_v2_behavior()
+
 
 _SERVER_URL_PREFIX = '/data/plugin/debugger/'
 
 
-@test_util.run_v1_only('Test fails to run and clean up properly; they time out.')
 class InteractiveDebuggerPluginTest(tf.test.TestCase):
 
   def setUp(self):
@@ -111,14 +114,14 @@ def _deserializeResponse(self, response):
   def _runSimpleAddMultiplyGraph(self, variable_size=1):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         a = tf.Variable([10.0] * variable_size, name='a')
         b = tf.Variable([20.0] * variable_size, name='b')
         c = tf.Variable([30.0] * variable_size, name='c')
         x = tf.multiply(a, b, name="x")
         y = tf.add(x, c, name="y")
 
-        sess.run(tf.compat.v1.global_variables_initializer())
+        sess.run(tf.global_variables_initializer())
 
         sess = tf_debug.TensorBoardDebugWrapperSession(sess, self._debugger_url)
         session_run_results.append(sess.run(y))
@@ -129,12 +132,12 @@ def session_run_job():
   def _runMultiStepAssignAddGraph(self, steps):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         a = tf.Variable(10, dtype=tf.int32, name='a')
         b = tf.Variable(1, dtype=tf.int32, name='b')
-        inc_a = tf.compat.v1.assign_add(a, b, name='inc_a')
+        inc_a = tf.assign_add(a, b, name='inc_a')
 
-        sess.run(tf.compat.v1.global_variables_initializer())
+        sess.run(tf.global_variables_initializer())
 
         sess = tf_debug.TensorBoardDebugWrapperSession(sess, self._debugger_url)
         for _ in range(steps):
@@ -146,15 +149,15 @@ def session_run_job():
   def _runTfGroupGraph(self):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         a = tf.Variable(10, dtype=tf.int32, name='a')
         b = tf.Variable(20, dtype=tf.int32, name='b')
         d = tf.constant(1, dtype=tf.int32, name='d')
-        inc_a = tf.compat.v1.assign_add(a, d, name='inc_a')
-        inc_b = tf.compat.v1.assign_add(b, d, name='inc_b')
+        inc_a = tf.assign_add(a, d, name='inc_a')
+        inc_b = tf.assign_add(b, d, name='inc_b')
         inc_ab = tf.group([inc_a, inc_b], name="inc_ab")
 
-        sess.run(tf.compat.v1.global_variables_initializer())
+        sess.run(tf.global_variables_initializer())
 
         sess = tf_debug.TensorBoardDebugWrapperSession(sess, self._debugger_url)
         session_run_results.append(sess.run(inc_ab))
@@ -703,7 +706,7 @@ def testGetSourceOpTraceback(self):
   def _runInitializer(self):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         a = tf.Variable([10.0] * 10, name='a')
         sess = tf_debug.TensorBoardDebugWrapperSession(sess, self._debugger_url)
         # Run the initializer with a debugger-wrapped tf.Session.
@@ -781,7 +784,7 @@ def testCommDataForUninitializedTensorIsHandledCorrectly(self):
   def _runHealthPillNetwork(self):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         a = tf.Variable(
             [np.nan, np.inf, np.inf, -np.inf, -np.inf, -np.inf, 10, 20, 30],
             dtype=tf.float32, name='a')
@@ -828,11 +831,11 @@ def testHealthPill(self):
   def _runAsciiStringNetwork(self):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         str1 = tf.Variable('abc', name='str1')
         str2 = tf.Variable('def', name='str2')
         str_concat = tf.add(str1, str2, name='str_concat')
-        sess.run(tf.compat.v1.global_variables_initializer())
+        sess.run(tf.global_variables_initializer())
         sess = tf_debug.TensorBoardDebugWrapperSession(sess, self._debugger_url)
         session_run_results.append(sess.run(str_concat))
     session_run_thread = threading.Thread(target=session_run_job)
@@ -885,11 +888,11 @@ def testAsciiStringTensorIsHandledCorrectly(self):
   def _runBinaryStringNetwork(self):
     session_run_results = []
     def session_run_job():
-      with tf.compat.v1.Session() as sess:
+      with tf.Session() as sess:
         str1 = tf.Variable([b'\x01' * 3, b'\x02' * 3], name='str1')
         str2 = tf.Variable([b'\x03' * 3, b'\x04' * 3], name='str2')
         str_concat = tf.add(str1, str2, name='str_concat')
-        sess.run(tf.compat.v1.global_variables_initializer())
+        sess.run(tf.global_variables_initializer())
         sess = tf_debug.TensorBoardDebugWrapperSession(sess, self._debugger_url)
         session_run_results.append(sess.run(str_concat))
     session_run_thread = threading.Thread(target=session_run_job)

tensorboard/plugins/debugger/session_debug_test.py

@@ -35,15 +35,18 @@
 
 import numpy as np
 import portpicker  # pylint: disable=import-error
-import tensorflow as tf  # pylint: disable=wrong-import-order
+import tensorflow.compat.v1 as tf  # pylint: disable=wrong-import-order
 from tensorflow.python import debug as tf_debug  # pylint: disable=wrong-import-order
 
 from tensorboard.plugins.debugger import constants
 from tensorboard.plugins.debugger import debugger_server_lib
 from tensorboard.util import test_util
 
+# These unit tests for Debugger Plugin V1 are tied to TF1.x behavior
+# (`tf.Session`s).
+tf.disable_v2_behavior()
+
 
-@test_util.run_v1_only("Server fails to come up. Requires study.")
 class SessionDebugTestBase(tf.test.TestCase):
 
   def setUp(self):
@@ -67,17 +70,17 @@ def tearDown(self):
     if os.path.isdir(self._logdir):
       shutil.rmtree(self._logdir)
 
-    tf.compat.v1.reset_default_graph()
+    tf.reset_default_graph()
 
   def _poll_server_till_success(self, max_tries, poll_interval_seconds):
     for _ in range(max_tries):
       try:
-        with tf.compat.v1.Session() as sess:
+        with tf.Session() as sess:
           a_init_val = np.array([42.0])
           a_init = tf.constant(a_init_val, shape=[1], name="a_init")
           a = tf.Variable(a_init, name="a")
 
-          run_options = tf.compat.v1.RunOptions(output_partition_graphs=True)
+          run_options = tf.RunOptions(output_partition_graphs=True)
           tf_debug.watch_graph(run_options,
                                sess.graph,
                                debug_ops=["DebugNumericSummary"],
@@ -125,7 +128,7 @@ def _compute_health_pill(self, x):
   def _check_health_pills_in_events_file(self,
                                          events_file_path,
                                          debug_key_to_tensors):
-    reader = tf.compat.v1.python_io.tf_record_iterator(events_file_path)
+    reader = tf.python_io.tf_record_iterator(events_file_path)
     event_read = tf.Event()
 
     # The first event in the file should contain the events version, which is
@@ -159,7 +162,7 @@ def _check_health_pills_in_events_file(self,
             health_pills[debug_key][i])
 
   def testRunSimpleNetworkoWithInfAndNaNWorks(self):
-    with tf.compat.v1.Session() as sess:
+    with tf.Session() as sess:
       x_init_val = np.array([[2.0], [-1.0]])
       y_init_val = np.array([[0.0], [-0.25]])
       z_init_val = np.array([[0.0, 3.0], [-1.0, 0.0]])
@@ -171,14 +174,14 @@ def testRunSimpleNetworkoWithInfAndNaNWorks(self):
       z_init = tf.constant(z_init_val, shape=[2, 2])
       z = tf.Variable(z_init, name="z")
 
-      u = tf.compat.v1.div(x, y, name="u")  # Produces an Inf.
+      u = tf.div(x, y, name="u")  # Produces an Inf.
       v = tf.matmul(z, u, name="v")  # Produces NaN and Inf.
 
       sess.run(x.initializer)
       sess.run(y.initializer)
       sess.run(z.initializer)
 
-      run_options = tf.compat.v1.RunOptions(output_partition_graphs=True)
+      run_options = tf.RunOptions(output_partition_graphs=True)
       tf_debug.watch_graph(run_options,
                            sess.graph,
                            debug_ops=["DebugNumericSummary"],
@@ -221,18 +224,18 @@ def testRunSimpleNetworkoWithInfAndNaNWorks(self):
     self.assertEqual(0, report[1].pos_inf_event_count)
 
   def testMultipleInt32ValuesOverMultipleRunsAreRecorded(self):
-    with tf.compat.v1.Session() as sess:
+    with tf.Session() as sess:
       x_init_val = np.array([10], dtype=np.int32)
       x_init = tf.constant(x_init_val, shape=[1], name="x_init")
       x = tf.Variable(x_init, name="x")
 
       x_inc_val = np.array([2], dtype=np.int32)
       x_inc = tf.constant(x_inc_val, name="x_inc")
-      inc_x = tf.compat.v1.assign_add(x, x_inc, name="inc_x")
+      inc_x = tf.assign_add(x, x_inc, name="inc_x")
 
       sess.run(x.initializer)
 
-      run_options = tf.compat.v1.RunOptions(output_partition_graphs=True)
+      run_options = tf.RunOptions(output_partition_graphs=True)
       tf_debug.watch_graph(run_options,
                            sess.graph,
                            debug_ops=["DebugNumericSummary"],
@@ -266,7 +269,7 @@ def testConcurrentNumericsAlertsAreRegisteredCorrectly(self):
     # Before any Session runs, the report ought to be empty.
     self.assertEqual([], self._debug_data_server.numerics_alert_report())
 
-    with tf.compat.v1.Session() as sess:
+    with tf.Session() as sess:
       x_init_val = np.array([[2.0], [-1.0]])
       y_init_val = np.array([[0.0], [-0.25]])
       z_init_val = np.array([[0.0, 3.0], [-1.0, 0.0]])
@@ -278,7 +281,7 @@ def testConcurrentNumericsAlertsAreRegisteredCorrectly(self):
       z_init = tf.constant(z_init_val, shape=[2, 2])
       z = tf.Variable(z_init, name="z")
 
-      u = tf.compat.v1.div(x, y, name="u")  # Produces an Inf.
+      u = tf.div(x, y, name="u")  # Produces an Inf.
       v = tf.matmul(z, u, name="v")  # Produces NaN and Inf.
 
       sess.run(x.initializer)
@@ -287,7 +290,7 @@ def testConcurrentNumericsAlertsAreRegisteredCorrectly(self):
 
       run_options_list = []
       for i in range(num_threads):
-        run_options = tf.compat.v1.RunOptions(output_partition_graphs=True)
+        run_options = tf.RunOptions(output_partition_graphs=True)
         # Use different grpc:// URL paths so that each thread opens a separate
         # gRPC stream to the debug data server, simulating multi-worker setting.
         tf_debug.watch_graph(run_options,