This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git
The following commit(s) were added to refs/heads/master by this push:
new 07449f5 Call previous excepthook in tvm_excepthook. (#5675)
07449f5 is described below
commit 07449f5a579bdde6e98527aa652e6f693ec4c81e
Author: notoraptor <[email protected]>
AuthorDate: Tue May 26 21:15:18 2020 -0400
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.
---
python/tvm/__init__.py | 21 +++++++++++++--------
1 file changed, 13 insertions(+), 8 deletions(-)
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)
