From 2f395cc76b97c75c5600be84e856eaaee824a203 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 9 Apr 2020 12:50:04 +0300
Subject: [PATCH] Use sub-process fork signal hooks instead of os._exit patch

---
 trains/binding/environ_bind.py | 22 ++++++++++++----------
 trains/task.py                 |  6 ++++--
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/trains/binding/environ_bind.py b/trains/binding/environ_bind.py
index 35678771..39714506 100644
--- a/trains/binding/environ_bind.py
+++ b/trains/binding/environ_bind.py
@@ -68,16 +68,18 @@ class PatchOsFork(object):
 
                 # Hack: now make sure we setup the reporter thread
                 task._setup_reporter()
-                # if we got here patch the os._exit of our instance to call us
-                def _at_exit_callback(*args, **kwargs):
-                    # call at exit manually
-                    # noinspection PyProtectedMember
-                    task._at_exit()
-                    # noinspection PyProtectedMember
-                    return os._org_exit(*args, **kwargs)
 
-                if not hasattr(os, '_org_exit'):
-                    os._org_exit = os._exit
-                os._exit = _at_exit_callback
+                # TODO: Check if the signal handler method is enough, for the time being, disable the _exit hook
+                # # if we got here patch the os._exit of our instance to call us
+                # def _at_exit_callback(*args, **kwargs):
+                #     # call at exit manually
+                #     # noinspection PyProtectedMember
+                #     task._at_exit()
+                #     # noinspection PyProtectedMember
+                #     return os._org_exit(*args, **kwargs)
+                #
+                # if not hasattr(os, '_org_exit'):
+                #     os._org_exit = os._exit
+                # os._exit = _at_exit_callback
 
         return ret
diff --git a/trains/task.py b/trains/task.py
index 72968ca1..ccf17986 100644
--- a/trains/task.py
+++ b/trains/task.py
@@ -218,7 +218,8 @@ class Task(_Task):
                 cls.__main_task._artifacts_manager = Artifacts(cls.__main_task)
                 # unregister signal hooks, they cause subprocess to hang
                 cls.__main_task.__register_at_exit(cls.__main_task._at_exit)
-                cls.__main_task.__register_at_exit(None, only_remove_signal_and_exception_hooks=True)
+                # TODO: Check if the signal handler method is safe enough, for the time being, do not unhook
+                # cls.__main_task.__register_at_exit(None, only_remove_signal_and_exception_hooks=True)
 
             if not running_remotely():
                 verify_defaults_match()
@@ -1633,7 +1634,8 @@ class Task(_Task):
                 if self._exit_callback:
                     atexit.register(self._exit_callback)
 
-                if not self._org_handlers and not Task._Task__is_subprocess():
+                # TODO: check if sub-process hooks are safe enough, for the time being allow it
+                if not self._org_handlers:  # ## and not Task._Task__is_subprocess():
                     if sys.platform == 'win32':
                         catch_signals = [signal.SIGINT, signal.SIGTERM, signal.SIGSEGV, signal.SIGABRT,
                                          signal.SIGILL, signal.SIGFPE]