From 90d39057993ca5a7b51ff7328eea0ae63ec37678 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 26 Apr 2021 13:14:53 +0800
Subject: [PATCH] feat(mge): add mge._exit

GitOrigin-RevId: e7f348a677483b6cd3a457426cc6a90b67b9f276
---
 imperative/python/megengine/__init__.py       | 21 +++++++++++++++++++
 .../python/megengine/distributed/launcher.py  |  2 ++
 2 files changed, 23 insertions(+)

diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py
index fc352f11a..cf3c98838 100644
--- a/imperative/python/megengine/__init__.py
+++ b/imperative/python/megengine/__init__.py
@@ -95,6 +95,27 @@ atexit.register(_close)
 
 del _set_fork_exec_path_for_timed_func
 
+_exit_handlers = []
+
+
+def _run_exit_handlers():
+    for handler in _exit_handlers:
+        handler()
+    _exit_handlers.clear()
+
+
+atexit.register(_run_exit_handlers)
+
+
+def _exit(code):
+    _run_exit_handlers()
+    sys.exit(code)
+
+
+def _atexit(handler):
+    _exit_handlers.append(handler)
+
+
 # subpackages
 import megengine.autodiff
 import megengine.data
diff --git a/imperative/python/megengine/distributed/launcher.py b/imperative/python/megengine/distributed/launcher.py
index 1e2a3dff9..f88f9ab45 100644
--- a/imperative/python/megengine/distributed/launcher.py
+++ b/imperative/python/megengine/distributed/launcher.py
@@ -11,6 +11,7 @@ import multiprocessing as mp
 import os
 import queue
 
+from .. import _exit
 from ..core._imperative_rt.core2 import sync
 from ..logger import get_logger
 from .group import group_barrier, init_process_group
@@ -53,6 +54,7 @@ def _run_wrapped(
     sync()
     if is_multimachine:
         group_barrier()
+    _exit(0)
 
 
 class launcher:
-- 
GitLab