non-fork start_method as default (#662)

* non-fork start_method as default

* tests

* fixes

* cleanup

* fixes

* new slurm image

* env var

* debugging

* remove prints

* fixes?

* tests

* changelog

* revert

* unix-paths for allowed_hosts

* lint

* topfschlagen

* fix block-network

* pr feedback

* ignore types

* fix types

Author: normanrz

cluster_tools/Changelog.md

@@ -10,6 +10,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v0.9.14...HEAD)
 
 ### Breaking Changes
+- The `multiprocessing` executor now uses `spawn` as default start method. `fork` and `forkserver` can be used by supplying a `start_method` argument (e.g. `cluster_tools.get_executor("multiprocessing", start_method="forkserver")`) or by setting the `MULTIPROCESSING_DEFAULT_START_METHOD` environment variable. [#662](https://github.com/scalableminds/webknossos-libs/pull/662)
 
 ### Added
 

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()

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):

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("/.*")