Fix support for sub-process (process pool)

trains/binding/environ_bind.py

@@ -1,5 +1,7 @@
 import os
 
+import six
+
 from ..config import TASK_LOG_ENVIRONMENT, running_remotely
 
 
@@ -34,3 +36,43 @@ class EnvironmentBind(object):
         if running_remotely():
             # put back into os:
             os.environ.update(env_param)
+
+
+class PatchOsFork(object):
+    _original_fork = None
+
+    @classmethod
+    def patch_fork(cls):
+        # only once
+        if cls._original_fork:
+            return
+        if six.PY2:
+            cls._original_fork = staticmethod(os.fork)
+        else:
+            cls._original_fork = os.fork
+        os.fork = cls._patched_fork
+
+    @staticmethod
+    def _patched_fork(*args, **kwargs):
+        ret = PatchOsFork._original_fork(*args, **kwargs)
+        # Make sure the new process stdout is logged
+        if not ret:
+            from ..task import Task
+            if Task.current_task() is not None:
+                # bind sub-process logger
+                task = Task.init()
+                task.get_logger().flush()
+
+                # 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