From cb7050c5356ec0012620bc5d665cd2b0cb34353b Mon Sep 17 00:00:00 2001 From: ghostplant Date: Tue, 6 Aug 2019 00:31:55 +0800 Subject: [PATCH] Quit and clean when TVM is interrupted (#3640) --- python/tvm/__init__.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py index 56b8b3d..9e3eb0f 100644 --- a/python/tvm/__init__.py +++ b/python/tvm/__init__.py @@ -18,6 +18,10 @@ """TVM: Low level DSL/IR stack for tensor computation.""" from __future__ import absolute_import as _abs +import multiprocessing +import sys +import traceback + from . import _pyversion from . import tensor @@ -59,3 +63,13 @@ from .tag import tag_scope # 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() + +sys.excepthook = tvm_excepthook -- 2.7.4