Repository: rwth-i6/returnn

Issue:
Hang in dataset iterator
With PyTorch.

I see this exception in the log:
```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/spawn.py", line 120, in spawn_main
    exitcode = _main(fd, parent_sentinel)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/spawn.py", line 130, in _main
    self = reduction.pickle.load(from_parent)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/u/zeyer/setups/combined/2021-05-31/tools/returnn/returnn/config.py", line 72, in __setstate__
    self.typed_dict = unpickler.load()
                      ^^^^^^^^^^^^^^^^
  File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/pickle.py", line 1213, in load
    dispatch[key[0]](self)
  File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/pickle.py", line 1347, in load_binstring
    self.append(self._decode_string(data))
                ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/pickle.py", line 1329, in _decode_string
    return value.decode(self.encoding, self.errors)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
UnicodeDecodeError: 'ascii' codec can't decode byte 0xe2 in position 0: ordinal not in range(128)
```
But then the trainer hangs.

Ps tree:
```
  90521          \_ python3.11                                                  
  90657          |   \_ python3.11
  90658          |   \_ watch memory
  90777          |   \_ MPD worker 0
  90779          |   \_ MPD worker 1
  90781          |   \_ MPD worker 2
  90782          |   \_ MPD worker 3
  91244          |   \_ python3.11  
  91362          |   \_ MPD worker 0
  91364          |   \_ MPD worker 1
  91365          |   \_ MPD worker 2
  91369          |   \_ MPD worker 3
  91801          |   \_ MPD worker 0
  91802          |   \_ MPD worker 1
  91803          |   \_ MPD worker 2
  91804          |   \_ MPD worker 3
  92265          |   \_ python3.11 <defunct>
```
92265 should be the sub proc which got the exception.

Py-spy on the parent:
```
% py-spy dump -p 90521     
Process 90521: /work/tools/users/zeyer/py-envs/py3.11-torch2.1/bin/python3.11 -u /u/zeyer/setups/combined/2021-05-31/tools/returnn/rnn.py /u/zeyer/setups/combined/2021-05-31/work/i6_core/returnn/training/ReturnnTrainingJob.wlbfTg0IM0nc/output/returnn.config
Python v3.11.2 (/work/tools/users/zeyer/linuxbrew/Cellar/python@3.11/3.11.2_1/bin/python3.11)

Thread 90521 (idle): "MainThread"
    _launch (multiprocessing/popen_spawn_posix.py:62)
    __init__ (multiprocessing/popen_fork.py:19)
    __init__ (multiprocessing/popen_spawn_posix.py:32)
    _Popen (multiprocessing/context.py:288)
    start (multiprocessing/process.py:121)
    start (returnn/returnn/util/multi_proc_non_daemonic_spawn.py:55)
    __init__ (torch/utils/data/dataloader.py:1039)
    _get_iterator (torch/utils/data/dataloader.py:386)
    __iter__ (torch/utils/data/dataloader.py:433)
    train_epoch (returnn/returnn/torch/engine.py:342)
    train (returnn/returnn/torch/engine.py:236)
    execute_main_task (returnn/returnn/__main__.py:543)
    main (returnn/returnn/__main__.py:737)
    <module> (returnn/rnn.py:11)
```

Or pystack:
```
...
    (Python) File "/u/zeyer/setups/combined/2021-05-31/tools/returnn/returnn/util/multi_proc_non_daemonic_spawn.py", line 55, in start
        super().start()
      Arguments:
        self: <NonDaemonicSpawnProcess at 0x7fdda657d890>
      Locals:
        __class__: <cell at 0x7fddda571300>
    (Python) File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/process.py", line 121, in start
        self._popen = self._Popen(self)
      Arguments:
        self: <NonDaemonicSpawnProcess at 0x7fdda657d890>
    (Python) File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/context.py", line 288, in _Popen
        return Popen(process_obj)
      Arguments:
        process_obj: <NonDaemonicSpawnProcess at 0x7fdda657d890>
      Locals:
        Popen: <type at 0x113085f0>
    (Python) File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/popen_spawn_posix.py", line 32, in __init__
        super().__init__(process_obj)
      Arguments:
        process_obj: <NonDaemonicSpawnProcess at 0x7fdda657d890>
        self: <Popen at 0x7fddad9ecbd0>
      Locals:
        __class__: <cell at 0x7fde1b92e620>
    (Python) File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/popen_fork.py", line 19, in __init__
        self._launch(process_obj)
      Arguments:
        process_obj: <NonDaemonicSpawnProcess at 0x7fdda657d890>
        self: <Popen at 0x7fddad9ecbd0>
    (Python) File "/work/tools/users/zeyer/linuxbrew/opt/python@3.11/lib/python3.11/multiprocessing/popen_spawn_posix.py", line 62, in _launch
        f.write(fp.getbuffer())
      Arguments:
        process_obj: <NonDaemonicSpawnProcess at 0x7fdda657d890>
        self: <Popen at 0x7fddad9ecbd0>
      Locals:
        cmd: [b"/work/tools/users/zeyer/py-envs/py3.11-torch2.1/bin/python3.11", "-c", ...]
        f: <_io.BufferedWriter at 0x7fdda65605c0>
        parent_w: 126
        child_r: 125
        parent_r: 123
        fp: <_io.BytesIO at 0x7fdda657ba60>
        child_w: 124
        prep_data: {"log_to_stderr": False, "authkey": <BINARY>, ...}
        tracker_fd: 39
        resource_tracker: <module at 0x7fde1b96f830>
```

Looking at the code in `popen_spawn_posix` in `_launch`:
```python
...
            parent_r, child_w = os.pipe()
            child_r, parent_w = os.pipe()
            cmd = spawn.get_command_line(tracker_fd=tracker_fd,
                                         pipe_handle=child_r)
            self._fds.extend([child_r, child_w])
            self.pid = util.spawnv_passfds(spawn.get_executable(),
                                           cmd, self._fds)
            self.sentinel = parent_r
            with open(parent_w, 'wb', closefd=False) as f:
                f.write(fp.getbuffer())   # <---- it hangs here
...
```

So, the sub proc dies because of some pickling error. But then the parent hangs because the fd is not closed yet because the sub proc is now in zombie state, not cleaned up yet (no `os.wait` yet)?



Discussion:
I realize, the zombie subproc is definitely not one of the MPD workers, as you see that all the others are complete.
The first unnamed subproc is `multiprocessing/resource_tracker`.
The second unnamed subproc is `multiprocessing/managers`.
The main proc hangs in starting the subproc for PyTorch `_MultiProcessingDataLoaderIter`. So this is the subproc which raised this exception and then crashed and which hangs in zombie state.

The corresponding code from the main proc in PyTorch dataloader `_MultiProcessingDataLoaderIter` (`torch/utils/data/dataloader.py:1039`) is:
```python
            w = multiprocessing_context.Process(
                target=_utils.worker._worker_loop,
                args=(self._dataset_kind, self._dataset, index_queue,
                      self._worker_result_queue, self._workers_done_event,
                      self._auto_collation, self._collate_fn, self._drop_last,
                      self._base_seed, self._worker_init_fn, i, self._num_workers,
                      self._persistent_workers, self._shared_seed))
            w.daemon = True
            # NB: Process.start() actually take some time as it needs to
            #     start a process and pass the arguments over via a pipe.
            #     Therefore, we only add a worker to self._workers list after
            #     it started, so that we do not call .join() if program dies
            #     before it starts, and __del__ tries to join but will get:
            #     AssertionError: can only join a started process.
            w.start()   # <---- hang here
```

The `w` is a `NonDaemonicSpawnProcess`, and it has an instance of `_DataLoaderWorkerPreInitFunc` for `pre_init_func`. The unpickling of this fails.

As the code happens in our custom `pre_init_func` handling logic (not directly there, but we can catch it there), we will catch the exception there and try to do some proper cleanup such that the parent should properly get the SIGPIPE.


Resolving patch:
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -2672,6 +2672,23 @@ def write_pickled_object(p: typing.BinaryIO, obj: Any):
     p.flush()
 
 
+def serialize_object(obj: Any) -> bytes:
+    """
+    Uses :func:`write_pickled_object`.
+    """
+    stream = BytesIO()
+    write_pickled_object(stream, obj)
+    return stream.getvalue()
+
+
+def deserialize_object(data: bytes) -> Any:
+    """
+    Uses :func:`read_pickled_object`.
+    """
+    stream = BytesIO(data)
+    return read_pickled_object(stream)
+
+
 def load_txt_vector(filename):
     """
     Expect line-based text encoding in file.
diff --git a/returnn/util/multi_proc_non_daemonic_spawn.py b/returnn/util/multi_proc_non_daemonic_spawn.py
--- a/returnn/util/multi_proc_non_daemonic_spawn.py
+++ b/returnn/util/multi_proc_non_daemonic_spawn.py
@@ -18,15 +18,22 @@
 
 from __future__ import annotations
 
+import io
+import pickle
 import time
-from typing import Optional, Any, Callable, Tuple
+from typing import Optional, Callable
 import os
 import signal
 import atexit
+from .basic import serialize_object, deserialize_object
+from . import better_exchook
 
 # noinspection PyProtectedMember
 from multiprocessing.context import BaseContext, SpawnProcess
 
+# noinspection PyUnresolvedReferences
+from multiprocessing import reduction
+
 # noinspection PyUnresolvedReferences
 from multiprocessing.util import is_exiting
 
@@ -111,10 +118,14 @@ def _close_cleanup(self):
         self._at_exit_cleanup_handler = None
 
     def __reduce__(self):
-        res_tuple = super().__reduce__()
         if not self.pre_init_func:
-            return res_tuple
-        reconstruct_func, reconstruct_args, *other = res_tuple
+            return super().__reduce__()
+        reconstruct_func, reconstruct_args, reconstruct_state = super().__reduce__()
+        reconstruct_state = reconstruct_state.copy()
+        reconstruct_state.pop("pre_init_func")
+        # For pickling the process object, we need to use the multiprocessing pickler.
+        buffer = io.BytesIO()
+        reduction.dump((reconstruct_func, reconstruct_args, reconstruct_state), buffer)
         # Use our own reconstruct function to call the pre_init_func.
         # This is unpickled and executed *before* the other state is unpickled.
         # This is important: This allows to potentially prepare some global state,
@@ -124,16 +135,39 @@ def __reduce__(self):
         # this is the way to do it.
         # Note that internally, multiprocessing SpawnProcess does sth similar,
         # see multiprocessing.spawn._main, spawn.prepare.
-        return (self._reconstruct_with_pre_init_func, (reconstruct_func, reconstruct_args, self.pre_init_func)) + tuple(
-            other
+        return (
+            self._reconstruct_with_pre_init_func,
+            (
+                serialize_object(self.pre_init_func),
+                buffer.getvalue(),
+            ),
         )
 
     @staticmethod
     def _reconstruct_with_pre_init_func(
-        reconstruct_func: Callable, reconstruct_args: Tuple[Any, ...], pre_init_func: Callable[[], None]
+        pre_init_func_serialized: bytes,
+        reconstruct_func_and_args_and_state_serialized: bytes,
     ):
-        pre_init_func()
-        return reconstruct_func(*reconstruct_args)
+        better_exchook.install()
+        try:
+            pre_init_func: Callable[[], None] = deserialize_object(pre_init_func_serialized)
+            pre_init_func()
+            buffer = io.BytesIO(reconstruct_func_and_args_and_state_serialized)
+            reconstruct_func, reconstruct_args, reconstruct_state = pickle.load(buffer)
+            obj = reconstruct_func(*reconstruct_args)
+            for k, v in reconstruct_state.items():
+                setattr(obj, k, v)
+            return obj
+        except Exception as exc:
+            print(
+                f"[PID {os.getpid()}, PPID {os.getppid()}]"
+                f" Error in NonDaemonicSpawnProcess._reconstruct_with_pre_init_func: {exc}"
+            )
+            # Note: All relevant data should have been read already from the pipe in the child proc,
+            # so we should not hang anymore in the parent while writing to the pipe.
+            # Thus, it should be safe to just reraise here.
+            # https://github.com/rwth-i6/returnn/issues/1514
+            raise
 
 
 class NonDaemonicSpawnContext(BaseContext):


Test patch:
diff --git a/tests/test_Util.py b/tests/test_Util.py
--- a/tests/test_Util.py
+++ b/tests/test_Util.py
@@ -10,6 +10,7 @@
 import numpy
 import unittest
 import textwrap
+import signal
 
 from returnn.util import better_exchook
 from returnn.util.py_ext_mod_compiler import PyExtModCompiler
@@ -19,6 +20,25 @@
 my_dir = os.path.dirname(os.path.realpath(__file__))
 
 
+def _sig_alarm_handler(signum, frame):
+    raise Exception(f"Alarm (timeout) signal handler")
+
+
+signal.signal(signal.SIGALRM, _sig_alarm_handler)
+
+
+@contextlib.contextmanager
+def timeout(seconds=10):
+    """
+    :param seconds: when the context is not closed within this time, an exception will be raised
+    """
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+
+
 def test_cmd_true():
     r = sys_cmd_out_lines("true")
     assert_equal(r, [])
@@ -537,6 +557,46 @@ def test_native_signal_handler():
     install_native_signal_handler(reraise_exceptions=True)
 
 
+class _PreInitUnpicklingRaiseException:
+    def __getstate__(self):
+        return 42
+
+    def __setstate__(self, state):
+        raise Exception("_PreInitUnpicklingRaiseException")
+
+
+class _PreInitDummy:
+    def __init__(self, *, payload):
+        self.payload = payload
+
+    def __call__(self):
+        pass  # nothing
+
+
+def _noop_func():
+    pass
+
+
+def test_NonDaemonicSpawnProcess_hang_1514():
+    # https://github.com/rwth-i6/returnn/issues/1514
+    from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnProcess
+
+    with timeout():
+        proc = NonDaemonicSpawnProcess(target=_noop_func)
+        proc.start()
+        proc.join()
+        assert proc.exitcode == 0
+
+        proc = NonDaemonicSpawnProcess(target=_noop_func)
+        # Give it some payload large enough such that it will the pipe buffer to trigger the potential hang.
+        proc.pre_init_func = _PreInitDummy(payload=["A" * 100_000, _PreInitUnpicklingRaiseException(), "B" * 100_000])
+        # After fixing this, we expect that the unpickling of pre_init_func will raise an exception,
+        # however, the parent then will not hang.
+        proc.start()
+        proc.join()
+        assert proc.exitcode != 0
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1: