Repository: scalableminds/webknossos-libs

Issue:
Always use `spawn` in multiprocessing
I just debugged an issue where my multiprocessing worker would hang on Linux, but not on macOS. After 4 hours of debugging, I found out that `s3fs` seems to have issues with the `fork` start method in multiprocessing. See https://github.com/fsspec/s3fs/issues/464.

I know there have been other occasions where switching the start method was the solution. I wonder if there are any downsides to always use `spawn` (or `forkserver` where available) within the ProcessPoolExecutors from `cluster_tools`?


Discussion:

Resolving patch:
diff --git a/cluster_tools/cluster_tools/__init__.py b/cluster_tools/cluster_tools/__init__.py
--- a/cluster_tools/cluster_tools/__init__.py
+++ b/cluster_tools/cluster_tools/__init__.py
@@ -25,7 +25,7 @@ def get_existent_kwargs_subset(whitelist, kwargs):
     return new_kwargs
 
 
-PROCESS_POOL_KWARGS_WHITELIST = ["max_workers", "mp_context", "initializer", "initargs"]
+PROCESS_POOL_KWARGS_WHITELIST = ["max_workers", "initializer", "initargs"]
 
 
 class WrappedProcessPoolExecutor(ProcessPoolExecutor):
@@ -37,31 +37,28 @@ class WrappedProcessPoolExecutor(ProcessPoolExecutor):
     """
 
     def __init__(self, **kwargs):
-        new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)
+        assert (not "start_method" in kwargs or kwargs["start_method"] is None) or (
+            not "mp_context" in kwargs
+        ), "Cannot use both `start_method` and `mp_context` kwargs."
 
-        self.did_overwrite_start_method = False
-        if kwargs.get("start_method", None) is not None:
-            self.did_overwrite_start_method = True
-            self.old_start_method = multiprocessing.get_start_method()
-            start_method = kwargs["start_method"]
-            logging.info(
-                f"Overwriting start_method to {start_method}. Previous value: {self.old_start_method}"
-            )
-            multiprocessing.set_start_method(start_method, force=True)
+        new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)
 
-        ProcessPoolExecutor.__init__(self, **new_kwargs)
+        mp_context = None
 
-    def shutdown(self, *args, **kwargs):
+        if "mp_context" in kwargs:
+            mp_context = kwargs["mp_context"]
+        elif "start_method" in kwargs and kwargs["start_method"] is not None:
+            mp_context = multiprocessing.get_context(kwargs["start_method"])
+        elif "MULTIPROCESSING_DEFAULT_START_METHOD" in os.environ:
+            mp_context = multiprocessing.get_context(
+                os.environ["MULTIPROCESSING_DEFAULT_START_METHOD"]
+            )
+        else:
+            mp_context = multiprocessing.get_context("spawn")
 
-        super().shutdown(*args, **kwargs)
+        new_kwargs["mp_context"] = mp_context
 
-        if self.did_overwrite_start_method:
-            logging.info(
-                f"Restoring start_method to original value: {self.old_start_method}."
-            )
-            multiprocessing.set_start_method(self.old_start_method, force=True)
-            self.old_start_method = None
-            self.did_overwrite_start_method = False
+        ProcessPoolExecutor.__init__(self, **new_kwargs)
 
     def submit(self, *args, **kwargs):
 
@@ -88,7 +85,7 @@ def submit(self, *args, **kwargs):
         # where wrapper_fn_1 is called, which eventually calls wrapper_fn_2, which eventually calls actual_fn.
         call_stack = []
 
-        if multiprocessing.get_start_method() != "fork":
+        if self._mp_context.get_start_method() != "fork":
             # If a start_method other than the default "fork" is used, logging needs to be re-setup,
             # because the programming context is not inherited in those cases.
             multiprocessing_logging_setup_fn = get_multiprocessing_logging_setup_fn()


Test patch:
diff --git a/cluster_tools/tests/test_multiprocessing.py b/cluster_tools/tests/test_multiprocessing.py
--- a/cluster_tools/tests/test_multiprocessing.py
+++ b/cluster_tools/tests/test_multiprocessing.py
@@ -14,6 +14,11 @@ def expect_fork():
     return True
 
 
+def expect_forkserver():
+    assert mp.get_start_method() == "forkserver"
+    return True
+
+
 def expect_spawn():
     assert mp.get_start_method() == "spawn"
     return True
@@ -22,22 +27,29 @@ def expect_spawn():
 def test_map_with_spawn():
     with cluster_tools.get_executor("multiprocessing", max_workers=5) as executor:
         assert executor.submit(
-            expect_fork
-        ).result(), "Multiprocessing should use fork by default"
+            expect_spawn
+        ).result(), "Multiprocessing should use `spawn` by default"
 
     with cluster_tools.get_executor(
         "multiprocessing", max_workers=5, start_method=None
     ) as executor:
         assert executor.submit(
-            expect_fork
-        ).result(), "Multiprocessing should use fork if start_method is None"
+            expect_spawn
+        ).result(), "Multiprocessing should use `spawn` if start_method is None"
 
     with cluster_tools.get_executor(
-        "multiprocessing", max_workers=5, start_method="spawn"
+        "multiprocessing", max_workers=5, start_method="forkserver"
     ) as executor:
         assert executor.submit(
-            expect_spawn
-        ).result(), "Multiprocessing should use spawn if requested"
+            expect_forkserver
+        ).result(), "Multiprocessing should use `forkserver` if requested"
+
+    with cluster_tools.get_executor(
+        "multiprocessing", max_workers=5, start_method="fork"
+    ) as executor:
+        assert executor.submit(
+            expect_fork
+        ).result(), "Multiprocessing should use `fork` if requested"
 
 
 def accept_high_mem(data):
diff --git a/webknossos/tests/conftest.py b/webknossos/tests/conftest.py
--- a/webknossos/tests/conftest.py
+++ b/webknossos/tests/conftest.py
@@ -1,5 +1,4 @@
 import json
-import multiprocessing
 import re
 import warnings
 from io import BytesIO
@@ -318,11 +317,11 @@ def pytest_collection_modifyitems(items: List[pytest.Item]) -> None:
         if item.get_closest_marker("vcr") is None:
             item.add_marker("vcr")
 
-        if (
-            item.get_closest_marker("block_network") is None
-            and multiprocessing.get_start_method() != "fork"
-        ):
-            # To allow for UNIX socket communication necessary for spawn multiprocessing
-            # addresses starting with `/` are allowed
-            marker = pytest.mark.block_network(allowed_hosts=["/.*"])
-            item.add_marker(marker)
+        # To allow for UNIX socket communication necessary for spawn multiprocessing
+        # addresses starting with `/` are allowed
+        marker = item.get_closest_marker("block_network")
+        if marker is None:
+            new_marker = pytest.mark.block_network(allowed_hosts=["/.*"])
+            item.add_marker(new_marker)
+        else:
+            marker.kwargs["allowed_hosts"].append("/.*")