From 8f6d831072ee771e8f8b6adffed0f11d6b283b10 Mon Sep 17 00:00:00 2001 From: notoraptor Date: Tue, 26 May 2020 21:15:18 -0400 Subject: [PATCH] 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 f781aef0a8be..b9fbcb29751c 100644 --- a/python/tvm/__init__.py +++ b/python/tvm/__init__.py @@ -63,12 +63,17 @@ # 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)