Forwarding out of range errors rather than capturing them within the base TFDBG

tensorflower-gardener · tensorflower-gardener · commit d627ca4e4d57 · 2017-10-02T16:40:25.000-07:00
wrapper session. Several use cases expect the error to be raised to cancel their iterations (namely the NMT Tutorial [1]). [1]: https://research.googleblog.com/2017/07/building-your-own-neural-machine.html PiperOrigin-RevId: 170773133
diff --git a/tensorflow/python/debug/wrappers/dumping_wrapper.py b/tensorflow/python/debug/wrappers/dumping_wrapper.py
@@ -36,6 +36,7 @@ def __init__(self,
                session_root,
                watch_fn=None,
                thread_name_filter=None,
+               pass_through_operrors=None,
                log_usage=True):
     """Constructor of DumpingDebugWrapperSession.
 
@@ -56,6 +57,8 @@ def __init__(self,
       thread_name_filter: Regular-expression white list for threads on which the
         wrapper session will be active. See doc of `BaseDebugWrapperSession` for
         more details.
+      pass_through_operrors: If true, all captured OpErrors will be
+        propagated. By default this captures all OpErrors.
       log_usage: (`bool`) whether the usage of this class is to be logged.
 
     Raises:
@@ -67,7 +70,8 @@ def __init__(self,
       pass  # No logging for open-source.
 
     framework.NonInteractiveDebugWrapperSession.__init__(
-        self, sess, watch_fn=watch_fn, thread_name_filter=thread_name_filter)
+        self, sess, watch_fn=watch_fn, thread_name_filter=thread_name_filter,
+        pass_through_operrors=pass_through_operrors)
 
     if gfile.Exists(session_root):
       if not gfile.IsDirectory(session_root):
diff --git a/tensorflow/python/debug/wrappers/framework.py b/tensorflow/python/debug/wrappers/framework.py
@@ -337,7 +337,8 @@ class BaseDebugWrapperSession(session.SessionInterface):
   # TODO(cais): Add on_cont_start and on_cont_end callbacks once the stepper is
   # is available.
 
-  def __init__(self, sess, thread_name_filter=None):
+  def __init__(self, sess, thread_name_filter=None,
+               pass_through_operrors=False):
     """Constructor of `BaseDebugWrapperSession`.
 
     Args:
@@ -349,6 +350,8 @@ def __init__(self, sess, thread_name_filter=None):
         by applying the `match` method of the compiled pattern. The default
         `None` means that the wrapper session will be active on all threads.
         E.g., r"MainThread$", r"QueueRunnerThread.*".
+      pass_through_operrors: If True, all captured OpErrors will be
+        propagated.  By default this captures all OpErrors.
 
     Raises:
       ValueError: On invalid `OnSessionInitAction` value.
@@ -361,6 +364,8 @@ def __init__(self, sess, thread_name_filter=None):
     self._sess = sess
     self._thread_name_filter_pattern = (re.compile(thread_name_filter)
                                         if thread_name_filter else None)
+    # TODO(cais/kstevens): Unittest this pass through feature.
+    self._pass_through_operrors = pass_through_operrors
 
     # Keeps track of number of run calls that have been performed on this
     # debug-wrapper session. The count can be used for purposes such as
@@ -480,6 +485,8 @@ def run(self,
                                    options=decorated_run_options,
                                    run_metadata=run_metadata)
       except errors.OpError as op_error:
+        if self._pass_through_operrors:
+          raise op_error
         tf_error = op_error
         retvals = op_error
 
@@ -783,7 +790,8 @@ def __repr__(self):
 class NonInteractiveDebugWrapperSession(BaseDebugWrapperSession):
   """Base class for non-interactive (i.e., non-CLI) debug wrapper sessions."""
 
-  def __init__(self, sess, watch_fn=None, thread_name_filter=None):
+  def __init__(self, sess, watch_fn=None, thread_name_filter=None,
+               pass_through_operrors=False):
     """Constructor of DumpingDebugWrapperSession.
 
     Args:
@@ -802,12 +810,15 @@ def __init__(self, sess, watch_fn=None, thread_name_filter=None):
       thread_name_filter: Regular-expression white list for threads on which the
         wrapper session will be active. See doc of `BaseDebugWrapperSession` for
         more details.
+      pass_through_operrors: If true, all captured OpErrors will be
+        propagated.  By default this captures all OpErrors.
     Raises:
        TypeError: If a non-None `watch_fn` is specified and it is not callable.
     """
 
     BaseDebugWrapperSession.__init__(
-        self, sess, thread_name_filter=thread_name_filter)
+        self, sess, thread_name_filter=thread_name_filter,
+        pass_through_operrors=pass_through_operrors)
 
     self._watch_fn = None
     if watch_fn is not None:
