Repository: pfnet/pytorch-pfn-extras

Issue:
The new ppe profiler introduces an unexpected ill-effect where users cannot set multiprocessing process start mode after importing ppe
I found out that importing the latest ppe will cause 

```python
>>> import multiprocessing as mp
>>> import pytorch_pfn_extras as ppe
>>> print(ppe.__version__)
>>> mp.set_start_method('forkserver')   # 'fork' or whatever

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/xxxx/.local/lib/python3.8/site-packages/pytorch_pfn_extras/__init__.py", line 10, in <module>
    from pytorch_pfn_extras import profiler  # NOQA
  File "/home/xxxx/.local/lib/python3.8/site-packages/pytorch_pfn_extras/profiler/__init__.py", line 1, in <module>
    from pytorch_pfn_extras.profiler._record import record  # NOQA
  File "/home/xxxx/.local/lib/python3.8/site-packages/pytorch_pfn_extras/profiler/_record.py", line 4, in <module>
    from pytorch_pfn_extras.profiler._time_summary import time_summary
  File "/home/xxxx/.local/lib/python3.8/site-packages/pytorch_pfn_extras/profiler/_time_summary.py", line 157, in <module>
    mp.set_start_method('forkserver')
  File "/usr/local/lib/python3.8/multiprocessing/context.py", line 243, in set_start_method
    raise RuntimeError('context has already been set')
RuntimeError: context has already been set
```

What is actually happening inside is very similar to this. 
https://github.com/numba/numba/pull/5724
Only the difference between ppe is that ppe creates `multiprocessing.JoinableQueue` instead of `get_context()` as in numba. In either case the process start method is implicitly fixed to "fork" right then and we can no longer change it anymore during the process lifetime, no matter which mode actually we want.

This happens at **import time** of ppe, specifically in instantiation of `TimeSummary`, not when the first time actually use it.
https://github.com/pfnet/pytorch-pfn-extras/blob/94a0d5092731ca94cc2deefaf5a55d037ec30dcc/pytorch_pfn_extras/profiler/_time_summary.py#L155
https://github.com/pfnet/pytorch-pfn-extras/blob/94a0d5092731ca94cc2deefaf5a55d037ec30dcc/pytorch_pfn_extras/profiler/_time_summary.py#L18

### Simple workaround

As written in https://github.com/numba/numba/pull/5724, if we set start method before importing ppe, it works fine.

```python
>>> import multiprocessing as mp
>>> mp.set_start_method('forkserver')
>>> import pytorch_pfn_extras as ppe
>>> print(ppe.__version__)
```

### Fundamental fix

I think we need to delay instantiation of _CPUWorker (or at least creation of `JoinableQueue`) to when it actually becomes necessary.


Discussion:
UGh, this is an ugly one, sorry for the problem and thanks for the debugging
I initially thought it's gonna be a little tough to solve, but soon have came up with an idea to just simply delay instantiation of `JoinableQueue`, which I attempted in #240 .
But I'm actually unaware that when the lazy open will happen, since I couldn't find how to use the profiling mechanism.

Maybe.. I'm afraid `_CPUWorker.__init__` will immediately start `self._worker` in a background thread so this fix won't solve the issue, will it?
Right, we need to delay creation of both workers too :(

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
@@ -1,9 +1,9 @@
-import atexit
 from contextlib import contextmanager
+import os
 import time
 from typing import Tuple
-from threading import Lock, Thread
-from queue import Queue
+import threading
+import queue
 import multiprocessing as mp
 import torch
 from pytorch_pfn_extras.reporting import DictSummary
@@ -12,24 +12,43 @@
 Events = Tuple[torch.cuda.Event, torch.cuda.Event]
 
 
-class _CPUWorker(object):
+class _CPUWorker:
     def __init__(self, add, max_queue_size: int):
         self._add = add
-        self._queue = mp.JoinableQueue(max_queue_size)
-        self._thread = Thread(target=self._worker)
-        self._thread.setDaemon(True)
+        self._max_queue_size = max_queue_size
+        self._initialized = False
+        self._queue = None
+        self._thread = None
+
+    def __del__(self):
+        self.finalize()
+
+    def initialize(self):
+        if self._initialized:
+            return
+        self._queue = mp.JoinableQueue(self._max_queue_size)
+        self._thread = threading.Thread(target=self._worker, daemon=True)
         self._thread.start()
+        self._initialized = True
 
-    def close(self):
+    def finalize(self):
+        if not self._initialized:
+            return
         self._queue.put(None)
-        self.wait()
-        self._thread.join()
+        self._queue.join()
         self._queue.close()
         self._queue.join_thread()
+        self._thread.join()
+        self._initialized = False
 
-    def wait(self):
+    def synchronize(self):
+        assert self._initialized
         self._queue.join()
 
+    def put(self, name, value):
+        assert self._initialized
+        self._queue.put((name, value))
+
     def _worker(self):
         while True:
             v = self._queue.get()
@@ -41,32 +60,51 @@ def _worker(self):
             self._queue.task_done()
 
 
-class _CUDAWorker(object):
+class _CUDAWorker:
     def __init__(self, add, max_queue_size: int):
         self._add = add
-        self._queue = Queue(max_queue_size)
-        self._thread = Thread(target=self._worker)
-        self._thread.setDaemon(True)
+        self._max_queue_size = max_queue_size
+        self._initialized = False
+        self._thread = None
+        self._queue = None
+        self._event_lock = threading.Lock()
+        self._events = None
+
+    def __del__(self):
+        self.finalize()
+
+    def initialize(self):
+        if self._initialized:
+            return
+        self._queue = queue.Queue(self._max_queue_size)
+        self._events = queue.Queue(self._max_queue_size * 2)
+        self._thread = threading.Thread(target=self._worker, daemon=True)
         self._thread.start()
-        self._event_lock = Lock()
-        self._events = Queue(max_queue_size)
+        self._initialized = True
 
-    def close(self):
+    def finalize(self):
+        if not self._initialized:
+            return
         self._queue.put(None)
-        self.wait()
+        self._queue.join()
         self._thread.join()
+        self._initialized = False
 
-    def wait(self):
+    def synchronize(self):
+        assert self._initialized
         self._queue.join()
 
+    def put(self, name, events):
+        assert self._initialized
+        self._queue.put((name, events))
+
     def _worker(self):
         while True:
             v = self._queue.get()
             if v is None:
                 self._queue.task_done()
                 break
-            name, value = v
-            begin, end = value
+            name, (begin, end) = v
             end.synchronize()
             t_ms = begin.elapsed_time(end)
             self._add(name, t_ms / 1000)
@@ -75,39 +113,79 @@ def _worker(self):
                 self._events.put(end)
             self._queue.task_done()
 
-    def _get_cuda_event(self) -> torch.cuda.Event:
+    def get_cuda_event(self) -> torch.cuda.Event:
+        assert self._initialized
         with self._event_lock:
             if self._events.empty():
                 self._events.put(torch.cuda.Event(enable_timing=True))
             return self._events.get()
 
 
-class TimeSummary(object):
+class TimeSummary:
     """Online summarization of execution times.
 
     `TimeSummary` computes the average and standard deviation of exeuction
     times in both cpu and gpu devices.
 
+    Args:
+        max_queue_size (int): Length limit of the internal queues that keep
+            reported time info until they are summarized.
+        auto_init (bool): Whether to automatically call `initialize()`
+            when the instance is created.
     """
 
-    def __init__(self, max_queue_size: int = 1000):
-        self._summary_lock = Lock()
+    def __init__(self, *, max_queue_size: int = 1000, auto_init: bool = True):
+        self._summary_lock = threading.Lock()
         self._summary = DictSummary()
         self._additional_stats = {}
 
-        self._cpu_worker = _CPUWorker(self._add, max_queue_size)
+        self._cpu_worker = _CPUWorker(self._add_from_worker, max_queue_size)
         if torch.cuda.is_available():
-            self._cuda_worker = _CUDAWorker(self._add, max_queue_size)
+            self._cuda_worker = _CUDAWorker(self._add_from_worker, max_queue_size)
         else:
             self._cuda_worker = None
 
-    def close(self):
-        self.wait()
-        self._cpu_worker.close()
+        self._initialized = False
+        self._master_pid = os.getpid()
+        if auto_init:
+            self.initialize()
+
+    def initialize(self) -> None:
+        """Initializes the worker threads for TimeSummary.
+
+        Usually you do not have to call it for yourself.
+        However in case you directly use ``ppe.time_summary`` outside of
+        :class:`pytorch_pfn_extras.training.extensions.ProfileReport`,
+        you have to explicitly call ``initialize()`` in advance.
+        """
+        if self._initialized:
+            return
+        if os.getpid() != self._master_pid:
+            raise RuntimeError(
+                "TimeSummary must be initialized in the same process as the "
+                "one created the instance. Please call initialize() in the "
+                "main process.")
+        self._cpu_worker.initialize()
+        if self._cuda_worker is not None:
+            self._cuda_worker.initialize()
+        self._initialized = True
+
+    def finalize(self):
+        if not self._initialized:
+            return
+        self._cpu_worker.finalize()
         if self._cuda_worker is not None:
-            self._cuda_worker.close()
+            self._cuda_worker.finalize()
+        self._initialized = False
 
-    def _add(self, name, value):
+    def synchronize(self):
+        self.initialize()
+        self._cpu_worker.synchronize()
+        if self._cuda_worker is not None:
+            self._cuda_worker.synchronize()
+
+    def _add_from_worker(self, name, value):
+        assert self._initialized
         with self._summary_lock:
             self._summary.add({name: value})
             min_value = self._additional_stats.get(f"{name}.min", value)
@@ -115,13 +193,12 @@ def _add(self, name, value):
             max_value = self._additional_stats.get(f"{name}.max", value)
             self._additional_stats[f"{name}.max"] = max(value, max_value)
 
-    def wait(self):
-        self._cpu_worker.wait()
-        if self._cuda_worker is not None:
-            self._cuda_worker.wait()
+    def add(self, name, value):
+        self._add_from_worker(name, value)
 
     @contextmanager
     def summary(self, clear: bool = False):
+        self.initialize()
         try:
             with self._summary_lock:
                 yield self._summary, self._additional_stats
@@ -142,21 +219,20 @@ def report(self, tag: str, use_cuda: bool = False) -> None:
             tag (str): A name to identify the section of code being profiled.
             use_cuda (bool): Indicates if GPU time should also be profiled.
         """
+        self.initialize()
         if use_cuda:
-            begin_event = self._cuda_worker._get_cuda_event()
+            begin_event = self._cuda_worker.get_cuda_event()
             begin_event.record()
         try:
             begin = time.time()
             yield
         finally:
             end = time.time()
-            self._cpu_worker._queue.put((tag, end - begin))
+            self._cpu_worker.put(tag, end - begin)
             if use_cuda:
-                end_event = self._cuda_worker._get_cuda_event()
+                end_event = self._cuda_worker.get_cuda_event()
                 end_event.record()
-                self._cuda_worker._queue.put(
-                    (f"{tag}.cuda", (begin_event, end_event)))
+                self._cuda_worker.put(f"{tag}.cuda", (begin_event, end_event))
 
 
-time_summary = TimeSummary()
-atexit.register(time_summary.close)
+time_summary = TimeSummary(auto_init=False)
diff --git a/pytorch_pfn_extras/training/extensions/profile_report.py b/pytorch_pfn_extras/training/extensions/profile_report.py
--- a/pytorch_pfn_extras/training/extensions/profile_report.py
+++ b/pytorch_pfn_extras/training/extensions/profile_report.py
@@ -52,6 +52,9 @@ def __init__(
         format=None,
         **kwargs,
     ):
+        # Initializes global TimeSummary.
+        time_summary.initialize()
+
         if store_keys is None:
             self._store_keys = store_keys
         else:


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/profiler_tests/test_time_summary.py b/tests/pytorch_pfn_extras_tests/profiler_tests/test_time_summary.py
--- a/tests/pytorch_pfn_extras_tests/profiler_tests/test_time_summary.py
+++ b/tests/pytorch_pfn_extras_tests/profiler_tests/test_time_summary.py
@@ -1,4 +1,5 @@
 import multiprocessing as mp
+import subprocess
 import sys
 
 import pytest
@@ -10,12 +11,12 @@ def test_report():
     summary = TimeSummary()
     with summary.report("foo"):
         pass
-    summary.wait()
+    summary.synchronize()
     with summary.summary() as s:
         assert "foo" in s[0].compute_mean()
         assert "foo.min" in s[1]
         assert "foo.max" in s[1]
-    summary.close()
+    summary.finalize()
 
 
 def worker(summary):
@@ -31,12 +32,12 @@ def test_report_from_other_process():
     p = mp.Process(target=worker, args=(summary,))
     p.start()
     p.join()
-    summary.wait()
+    summary.synchronize()
     with summary.summary() as s:
         assert "foo" in s[0].compute_mean()
         assert "foo.min" in s[1]
         assert "foo.max" in s[1]
-    summary.close()
+    summary.finalize()
 
 
 def worker1():
@@ -48,10 +49,11 @@ def worker1():
     sys.platform == 'win32',
     reason='Multiprocessing not fully supported on Windows')
 def test_global_summary():
+    time_summary.initialize()
     p = mp.Process(target=worker1)
     p.start()
     p.join()
-    time_summary.wait()
+    time_summary.synchronize()
     with time_summary.summary() as s:
         assert "foo" in s[0].compute_mean()
         assert "foo.min" in s[1]
@@ -60,14 +62,23 @@ def test_global_summary():
 
 def test_clear():
     summary = TimeSummary()
-    summary._add("foo", 10)
-    summary._add("foo", 5)
-    summary._add("foo", 15)
-    summary.wait()
+    summary.add("foo", 10)
+    summary.add("foo", 5)
+    summary.add("foo", 15)
+    summary.synchronize()
     with summary.summary(clear=True) as s:
         assert s[0].compute_mean() == {"foo": 10}
         assert s[1] == {"foo.min": 5, "foo.max": 15}
-    with summary.summary(clear=True) as s:
-        assert s[0].compute_mean() == {}
-        assert s[1] == {}
-    summary.close()
+    summary.finalize()
+
+
+def test_multiprocessing_start_method():
+    # Ensure that importing PPE does not initialize multiprocessing context.
+    # See #238 for the context.
+    subprocess.check_call([
+        sys.executable,
+        '-c',
+        ('import multiprocessing as mp; '
+         + 'import pytorch_pfn_extras; '
+         + 'mp.set_start_method("spawn"); ')
+    ])