From f4733e451c8c99e89637a9d2045615259b854671 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 56b8b3d9d298..9e3eb0faefb8 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 @@ # 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