feat(mge): add mge._exit

GitOrigin-RevId: e7f348a677

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

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: