Fix 2 bugs where errors reported directly to the coordinator were not propagated

out of a "with tf.MonitoredSession()" block.

1: If the _CoordinatedSession detected a request for stop during a call to run() the
   request was not honored.  (Missing a finally: clause)
2: If a thread requested a stop but the with.MonitoredSession block terminated without
   calling run() the coordinated threads where not joined and the optional exception
   was not raised.

Added 2 tests, one for each bug.

Modernize the code to take benefit of the fact that the threads now register themselves
with the Coordinator directly.
Change: 130890075

Author: tensorflower-gardener

tensorflow/contrib/learn/python/learn/monitored_session.py

@@ -267,6 +267,7 @@ def __init__(self,
     self._config = config
     self._hooks = hooks or []
     self._scaffold = scaffold or Scaffold()
+    self._coord = None
     for h in self._hooks:
       h.begin()
     # Create the session.
@@ -299,12 +300,10 @@ def _create_session(self):
     # Keep the tf_sess for quick runs of global step when needed.
     self._tf_sess = tf_sess
     # We don't want coordinator to suppress any exception.
-    coord = coordinator.Coordinator(clean_stop_exception_types=[])
-    coordinated_threads_to_join = queue_runner.start_queue_runners(
-        sess=tf_sess, coord=coord)
-    return _CoordinatedSession(
-        _HookedSession(tf_sess, self._hooks), coord,
-        coordinated_threads_to_join)
+    self._coord = coordinator.Coordinator(clean_stop_exception_types=[])
+    queue_runner.start_queue_runners(sess=tf_sess, coord=self._coord)
+    return _CoordinatedSession(_HookedSession(tf_sess, self._hooks),
+                               self._coord)
 
   @property
   def scaffold(self):
@@ -346,10 +345,20 @@ def _close_internal(self, exception_type=None):
       if not exception_type:
         for h in self._hooks:
           h.end(self._tf_sess)
+      if not self._coord.joined:
+        # We exited cleanly without stopping.  Some things now.  This will also
+        # re-raise exceptions from the coordinated threads, as needed.
+        self._coord.request_stop()
+        self._coord.join()
     finally:
       self._sess.close()
       self._sess = None
       self._tf_sess = None
+      self._coord = None
+
+  @property
+  def coord(self):
+    return self._coord
 
   def _is_closed(self):
     """Return True if the supervised session is closed.  For tests only.
@@ -491,24 +500,22 @@ class _CoordinatedSession(_WrappedSession):
   raises an exception, the exception is reported to the coordinator.
 
   In addition, after each call to `run()` this session ask the coordinator if
-  the session should stop.  In that case it will will join all the coordinated
-  threads passed to the constructor before returning.
+  the session should stop.  In that case it will will join all the threads
+  registered with the coordinator before returning.
 
   If the coordinator was requested to stop with an exception, that exception
   will be re-raised from the call to `run()`.
   """
 
-  def __init__(self, sess, coord, coordinated_threads_to_join):
+  def __init__(self, sess, coord):
     """Create a new `_CoordinatedSession`.
 
     Args:
       sess: A `tf.Session` object.  The wrapped session.
       coord: A `tf.train.Coordinator` object.
-      coordinated_threads_to_join: A list of threads.
     """
     _WrappedSession.__init__(self, sess)
     self._coord = coord
-    self._coordinated_threads_to_join = coordinated_threads_to_join
 
   def _check_stop(self):
     # Check with the coordinator if we should stop.
@@ -518,7 +525,7 @@ def close(self):
     try:
       if not self._coord.should_stop():
         self._coord.request_stop()
-        self._coord.join(self._coordinated_threads_to_join)
+        self._coord.join()
     except Exception:  # pylint: disable=broad-except
       # Don't raise exception at close
       pass
@@ -530,8 +537,9 @@ def run(self, *args, **kwargs):
       return self._sess.run(*args, **kwargs)
     except Exception as e:  # pylint: disable=broad-except
       self._coord.request_stop(e)
-    if self._coord.should_stop():
-      self._coord.join(self._coordinated_threads_to_join)
+    finally:
+      if self._coord.should_stop():
+        self._coord.join()
 
 
 class _HookedSession(_WrappedSession):

tensorflow/contrib/learn/python/learn/tests/monitored_session_test.py

@@ -187,7 +187,7 @@ def test_properties(self):
     with self.test_session() as sess:
       tf.constant(0.0)
       coord = tf.train.Coordinator()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, [])
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       self.assertEquals(sess.graph, coord_sess.graph)
       self.assertEquals(sess.sess_str, coord_sess.sess_str)
 
@@ -196,21 +196,21 @@ def test_run(self):
       c = tf.constant(0)
       v = tf.identity(c)
       coord = tf.train.Coordinator()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, [])
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       self.assertEqual(42, coord_sess.run(v, feed_dict={c: 42}))
 
   def test_should_stop_on_close(self):
     with self.test_session() as sess:
       coord = tf.train.Coordinator()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, [])
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       self.assertFalse(coord_sess.should_stop())
       coord_sess.close()
       self.assertTrue(coord_sess.should_stop())
 
   def test_should_stop_on_coord_stop(self):
     with self.test_session() as sess:
       coord = tf.train.Coordinator()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, [])
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       self.assertFalse(coord_sess.should_stop())
       coord.request_stop()
       self.assertTrue(coord_sess.should_stop())
@@ -220,7 +220,7 @@ def test_request_stop_on_exception(self):
       c = tf.constant(0)
       v = tf.identity(c)
       coord = tf.train.Coordinator()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, [])
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       self.assertFalse(coord_sess.should_stop())
       self.assertEqual(0, coord_sess.run(c))
       self.assertEqual(1, coord_sess.run(v, feed_dict={c: 1}))
@@ -237,8 +237,9 @@ def test_stop_threads_on_exception(self):
       threads = [threading.Thread(
           target=busy_wait_for_coord_stop, args=(coord,)) for _ in range(3)]
       for t in threads:
+        coord.register_thread(t)
         t.start()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, threads)
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       self.assertFalse(coord_sess.should_stop())
       for t in threads:
         self.assertTrue(t.is_alive())
@@ -261,8 +262,9 @@ def test_stop_threads_on_close(self):
       threads = [threading.Thread(
           target=busy_wait_for_coord_stop, args=(coord,)) for _ in range(3)]
       for t in threads:
+        coord.register_thread(t)
         t.start()
-      coord_sess = monitored_session._CoordinatedSession(sess, coord, threads)
+      coord_sess = monitored_session._CoordinatedSession(sess, coord)
       coord_sess.close()
       for t in threads:
         self.assertFalse(t.is_alive())
@@ -766,6 +768,46 @@ def test_regular_exception_pass_through_run(self):
       self.assertTrue(hook.raised)
       self.assertTrue(session.should_stop())
 
+  def test_regular_exception_reported_to_coord_pass_through_run(self):
+    # Tests that regular exceptions reported to the coordinator from a thread
+    # passes through a "run()" call within a "with MonitoredSession" block and
+    # set the session in stop mode.
+    with tf.Graph().as_default():
+      gstep = tf.contrib.framework.get_or_create_global_step()
+      scaffold = monitored_session.Scaffold()
+      session = monitored_session.MonitoredSession('', scaffold=scaffold)
+      with self.assertRaisesRegexp(RuntimeError, 'a thread wants to stop'):
+        with session:
+          self.assertEqual(0, session.run(gstep))
+          # Report an exception through the coordinator.
+          try:
+            raise RuntimeError('a thread wants to stop')
+          except RuntimeError as e:
+            session.coord.request_stop(e)
+          # Call run() which should raise the reported exception.
+          self.assertEqual(0, session.run(gstep))
+          # We should not hit this
+          self.assertFalse(True)
+
+  def test_regular_exception_reported_to_coord_pass_through_return(self):
+    # Tests that regular exceptions reported to the coordinator from a thread
+    # passes through returning from a "with MonitoredSession" block and
+    # set the session in stop mode.
+    with tf.Graph().as_default():
+      gstep = tf.contrib.framework.get_or_create_global_step()
+      scaffold = monitored_session.Scaffold()
+      session = monitored_session.MonitoredSession('', scaffold=scaffold)
+      with self.assertRaisesRegexp(RuntimeError, 'a thread wants to stop'):
+        with session:
+          self.assertEqual(0, session.run(gstep))
+          # Report an exception through the coordinator.
+          try:
+            raise RuntimeError('a thread wants to stop')
+          except RuntimeError as e:
+            session.coord.request_stop(e)
+          # Do not call run, just terminate the with.session block cleanly.
+          self.assertTrue(session.should_stop())
+
   # This set of tests, verifies the session behavior when exceptions are raised
   # from code inside a "with MonitoredSession:" context.