Repository: pfnet/pytorch-pfn-extras

Issue:
EOF Error in multiprocessing
After the upgrading of ppe==0.5.0, the following test code shows an EOF error.

Also under ppe==0.4.5, there is an issue that the last reported metric is not recorded on tensorbord (I mean the third `train/loss` and `val/loss` in the following code),
But, there is a workaround that is just adding time.sleep(1) before the end of the subprocess like the following.

However, under the new version of ppe, the following code fails completely, and I haven't found any workaround or fixing solution.

Does anyone have solutions?

```python
import multiprocessing as mp
import random
import time

import pytorch_pfn_extras as ppe
from torch.utils.tensorboard import SummaryWriter


def subprocess(num: int):
    ppe_extensions = [
        ppe.training.extensions.LogReport(),
        ppe.training.extensions.PlotReport(y_keys=["train/loss", "val/loss"], x_key="epoch",
                                           file_name="loss{}.png".format(num)),
    ]
    writer = ppe.writing.SimpleWriter(out_dir="save/logs-{}".format(num))
    manager = ppe.training.ExtensionsManager(models={}, optimizers={}, max_epochs=3,
                                             extensions=ppe_extensions, writer=writer, iters_per_epoch=1)
    tb_writer = SummaryWriter(log_dir="save/tensorboard-{}".format(num))

    @ppe.training.make_extension(trigger=(1, 'epoch'))
    def tensorboard_writer(manager):
        m = manager.observation
        tb_writer.add_scalars("loss", {'train/loss': m['train/loss'], 'val/loss': m['val/loss']}, manager.epoch)

    manager.extend(tensorboard_writer)

    for _ in range(3):
        with manager.run_iteration():
            ppe.reporting.report({"train/loss": random.random()})
            ppe.reporting.report({"val/loss": random.random()})
    time.sleep(1)
    return

def main():
    process = []
    for i in range(2):
        p = mp.Process(target=subprocess, args=(i,))
        process.append(p)
    for p in process:
        p.start()
    for p in process:
        p.join()
        p.close()


if __name__ == '__main__':
    main()
```

Result:
```
Exception in thread Thread-2:
Traceback (most recent call last):
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/threading.py", line 910, in run
    self._target(*self._args, **self._kwargs)
  File "/*************/ppe-debug/.venv/lib/python3.9/site-packages/pytorch_pfn_extras/profiler/_time_summary.py", line 80, in _worker
    v = self._queue.get()
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/queues.py", line 103, in get
    res = self._recv_bytes()
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/connection.py", line 221, in recv_bytes
    buf = self._recv_bytes(maxlength)
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/connection.py", line 419, in _recv_bytes
    buf = self._recv(4)
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/connection.py", line 388, in _recv
    raise EOFError
EOFError
Exception in thread Thread-2:
Traceback (most recent call last):
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/threading.py", line 910, in run
    self._target(*self._args, **self._kwargs)
  File "/*************/ppe-debug/.venv/lib/python3.9/site-packages/pytorch_pfn_extras/profiler/_time_summary.py", line 80, in _worker
    v = self._queue.get()
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/queues.py", line 103, in get
    res = self._recv_bytes()
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/connection.py", line 221, in recv_bytes
    buf = self._recv_bytes(maxlength)
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/connection.py", line 419, in _recv_bytes
    buf = self._recv(4)
  File "/usr/local/Cellar/python@3.9/3.9.6/Frameworks/Python.framework/Versions/3.9/lib/python3.9/multiprocessing/connection.py", line 388, in _recv
    raise EOFError
EOFError
```


Discussion:
Thanks for reporting and sorry for the delayed response :bow:
It seems this is caused by `time_summary`, which is a global singleton, is shared between processes. Let us fix it.
We've identified that the problem occurs because that initialization of TimeSummary (introduced in v0.5.0) is thread-unsafe: https://github.com/pfnet/pytorch-pfn-extras/blob/9073555b17bb22ce45fefcf601dada94076906c2/pytorch_pfn_extras/profiler/_time_summary.py#L106-L113

When two processes (or threads) call `initialize()` almost at the same time, initialization runs two times, causing worker threads to be launched two times. `_CUDAWorker` gets lost track of the first worker thread thus not taken care of by `finalize()`, so it continues to consume the queue, causing EOFError.

I think TensorBoard issue in v0.4.5 is a different issue. Let me investigate further.

Resolving patch:
diff --git a/pytorch_pfn_extras/profiler/_time_summary.py b/pytorch_pfn_extras/profiler/_time_summary.py
--- a/pytorch_pfn_extras/profiler/_time_summary.py
+++ b/pytorch_pfn_extras/profiler/_time_summary.py
@@ -50,6 +50,7 @@ def __init__(
         self._initialized = False
         self._queue: Optional[mp.JoinableQueue[Optional[Tuple[str, float]]]] = None
         self._thread: Optional[threading.Thread] = None
+        self._thread_exited = False
 
     def initialize(self) -> None:
         if self._initialized:
@@ -58,17 +59,21 @@ def initialize(self) -> None:
         self._thread = threading.Thread(target=self._worker, daemon=True)
         self._thread.start()
         self._initialized = True
+        self._thread_exited = False
 
     def finalize(self) -> None:
         if not self._initialized:
             return
         assert self._queue is not None
         assert self._thread is not None
-        self._queue.put(None)
+        # In some situations, (when this runs in a subprocess), the queue might have
+        # been cut in the worker thread before this function is called
+        # due to the non-deterministic shutdown process.
+        if not self._thread_exited:
+            self._queue.put(None)
         self._queue.join()
         self._queue.close()
         self._queue.join_thread()
-        self._thread.join()
         self._initialized = False
 
     def synchronize(self) -> None:
@@ -77,12 +82,19 @@ def synchronize(self) -> None:
 
     def put(self, name: str, value: float) -> None:
         assert self._queue is not None
+        assert not self._thread_exited
         self._queue.put((name, value))
 
     def _worker(self) -> None:
         assert self._queue is not None
         while True:
-            v = self._queue.get()
+            try:
+                v = self._queue.get()
+            # If this runs in a subprocess, the cleanup may throw an EOF here
+            # before the queue cleanup code is executed
+            except EOFError:
+                self._thread_exited = True
+                break
             if v is None:
                 self._queue.task_done()
                 break
@@ -107,6 +119,7 @@ def __init__(
         self._queue: Optional['queue.Queue[Optional[_QueueElem]]'] = None
         self._event_lock = threading.Lock()
         self._events: Optional['queue.Queue[torch.cuda.Event]'] = None
+        self._thread_exited = False
 
     def initialize(self) -> None:
         if self._initialized:
@@ -116,15 +129,16 @@ def initialize(self) -> None:
         self._thread = threading.Thread(target=self._worker, daemon=True)
         self._thread.start()
         self._initialized = True
+        self._thread_exited = False
 
     def finalize(self) -> None:
         if not self._initialized:
             return
         assert self._queue is not None
         assert self._thread is not None
-        self._queue.put(None)
+        if not self._thread_exited:
+            self._queue.put(None)
         self._queue.join()
-        self._thread.join()
         self._initialized = False
 
     def synchronize(self) -> None:
@@ -137,13 +151,18 @@ def put(
             events: Tuple[torch.cuda.Event, torch.cuda.Event],
     ) -> None:
         assert self._queue is not None
+        assert not self._thread_exited
         self._queue.put((name, events))
 
     def _worker(self) -> None:
         assert self._queue is not None
         assert self._events is not None
         while True:
-            v = self._queue.get()
+            try:
+                v = self._queue.get()
+            except EOFError:
+                self._thread_exited = True
+                break
             if v is None:
                 self._queue.task_done()
                 break


Test patch: