Merge pull request #125 from jaberg/thunk_hook_refactor

theano/gof/link.py

@@ -8,6 +8,33 @@ from copy import copy
 from theano.gof.python25 import all
 
 __excepthook = sys.excepthook
+
+def log_thunk_trace(value, f=sys.stderr):
+    """Log theano's diagnostic stack trace for an exception
+    raised by raise_with_op.
+    """
+    # in future, consider accepting `write` as arg rather than file
+    # to support writing to a logger
+    def write(msg):
+        print >> f, "log_thunk_trace: %s" % msg.strip()
+
+    if hasattr(value, '__thunk_trace__'):
+        trace2 = value.__thunk_trace__
+        write("There was a problem executing an Op.")
+        if trace2 is None:
+            write("Could not find where this Op was defined.")
+            write(" * You might have instantiated this Op "
+                    "directly instead of using a constructor.")
+            write(" * The Op you constructed might have been"
+                    " optimized. Try turning off optimizations.")
+        elif trace2:
+            write("Definition in: ")
+            for line in traceback.format_list(trace2):
+                write(line)
+            write("For the full definition stack trace set"
+                    " the Theano flags traceback.limit to -1")
+
+
 def thunk_hook(type, value, trace):
     """WRITEME
     This function is meant to replace excepthook and do some
@@ -20,20 +47,7 @@ def thunk_hook(type, value, trace):
 
     :note: This hook replaced by nosetests, so it does not run in nose tests.
     """
-    if hasattr(value, '__thunk_trace__'):
-        trace2 = value.__thunk_trace__
-        if trace2 is None:
-            print>>sys.stderr, "Could not find where this Op was defined."
-            print>>sys.stderr, (" * You might have instantiated this Op "
-                    "directly instead of using a constructor.")
-            print>>sys.stderr, (" * The Op you constructed might have been"
-                    " optimized. Try turning off optimizations.")
-        elif trace2:
-            print>>sys.stderr, "Definition in: "
-            for line in traceback.format_list(trace2):
-                print>>sys.stderr, line,
-            print>>sys.stderr, ("For the full definition stack trace set"
-                    " the Theano flags traceback.limit to -1")
+    log_thunk_trace(value)
     __excepthook(type, value, trace)
 sys.excepthook = thunk_hook
 
@@ -85,8 +99,15 @@ def raise_with_op(op, exc_info=None):
         exc_value.__applynode_index__ = op.env.toposort().index(op)
     else:
         exc_value.__applynode_index__ = None
+
+    # nose and unittest catch the exception and do not run th thunk_hook
+    # so it can be useful to just blurt out errors right here
+    if raise_with_op.print_thunk_trace:
+        log_thunk_trace(exc_value)
+
     raise exc_type, exc_value, exc_trace
 
+raise_with_op.print_thunk_trace = False
 
 class Linker(object):
     """WRITEME"""

theano/gof/vm.py

@@ -20,24 +20,8 @@ AddConfigVar('profile',
         "If VM should collect profile information",
         BoolParam(False))
 
+raise_with_op = link.raise_with_op
 
-def raise_with_op(op, exc_info = None):
-    """WRITEME"""
-    if exc_info is None:
-        exc_info = sys.exc_info()
-    exc_type, exc_value, exc_trace = exc_info
-    if exc_type == KeyboardInterrupt:
-        # print a simple traceback from KeyboardInterrupt
-        raise exc_type, exc_value, exc_trace
-    try:
-        trace = op.tag.trace
-    except AttributeError:
-        trace = ()
-    exc_value.__thunk_trace__ = trace
-    exc_value.args += (op, )
-    if op in op.env.toposort():
-        exc_value.args += ('Sequence id of Apply node='+str(op.env.toposort().index(op)),)
-    raise exc_type, exc_value, exc_trace
 
 class VM(object):
     """
@@ -114,6 +98,7 @@ class VM(object):
             self.call_times[i] = 0.0
             self.call_counts[i] = 0
 
+
 class Loop(VM):
     """
     Unconditional start-to-finish program execution in Python.
@@ -141,6 +126,7 @@ class Loop(VM):
             except:
                 raise_with_op(node)
 
+
 class LoopGC(VM):
     """
     Unconditional start-to-finish program execution in Python.
@@ -179,6 +165,7 @@ class LoopGC(VM):
             except:
                 raise_with_op(node)
 
+
 class Stack(VM):
     """
     Finish-to-start evalution order of thunks.
@@ -385,6 +372,7 @@ class Stack(VM):
                                 if empty_storage_map:
                                     storage_map[i][0] = None
 
+
 try:
     import lazylinker_c
     class CVM(lazylinker_c.CLazyLinker, VM):
@@ -394,6 +382,7 @@ try:
 except ImportError:
     pass
 
+
 class VM_Linker(link.LocalLinker):
     """
     Class that satisfies the Linker interface by acting as a VM factory.