From: notoraptor <notoraptor@users.noreply.github.com>
Date: Wed, 27 May 2020 01:15:18 +0000 (-0400)
Subject: Call previous excepthook in tvm_excepthook. (#5675)
X-Git-Tag: upstream/0.7.0~653
X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=07449f5a579bdde6e98527aa652e6f693ec4c81e;p=platform%2Fupstream%2Ftvm.git

Call previous excepthook in tvm_excepthook. (#5675)

* Call previous excepthook in tvm_excepthook.

* Rename prev_excepthook.

* Create a tvm_wrap_excepthook to wrap a given excepthook with tvm custom excepthook work
and call it on system previous excepthook.

* Add docstring.
---

diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index f781aef..b9fbcb2 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -63,12 +63,17 @@ from . import arith
 # Contrib initializers
 from .contrib import rocm as _rocm, nvcc as _nvcc, sdaccel as _sdaccel
 
-# Clean subprocesses when TVM is interrupted
-def tvm_excepthook(exctype, value, trbk):
-    print('\n'.join(traceback.format_exception(exctype, value, trbk)))
-    if hasattr(multiprocessing, 'active_children'):
-        # pylint: disable=not-callable
-        for p in multiprocessing.active_children():
-            p.terminate()
+def tvm_wrap_excepthook(exception_hook):
+    """Wrap given excepthook with TVM additional work."""
 
-sys.excepthook = tvm_excepthook
+    def wrapper(exctype, value, trbk):
+        """Clean subprocesses when TVM is interrupted."""
+        exception_hook(exctype, value, trbk)
+        if hasattr(multiprocessing, 'active_children'):
+            # pylint: disable=not-callable
+            for p in multiprocessing.active_children():
+                p.terminate()
+
+    return wrapper
+
+sys.excepthook = tvm_wrap_excepthook(sys.excepthook)