We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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`?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `cluster_tools/cluster_tools/__init__.py`
Content:
```
1 import logging
2 import multiprocessing
3 import os
4 import shutil
5 import tempfile
6 from concurrent import futures
7 from concurrent.futures import ProcessPoolExecutor
8 from functools import partial
9 from pathlib import Path
10 
11 from . import pickling
12 from .multiprocessing_logging_handler import get_multiprocessing_logging_setup_fn
13 from .schedulers.kube import KubernetesExecutor
14 from .schedulers.pbs import PBSExecutor
15 from .schedulers.slurm import SlurmExecutor
16 from .util import enrich_future_with_uncaught_warning
17 
18 
19 def get_existent_kwargs_subset(whitelist, kwargs):
20     new_kwargs = {}
21     for arg_name in whitelist:
22         if arg_name in kwargs:
23             new_kwargs[arg_name] = kwargs[arg_name]
24 
25     return new_kwargs
26 
27 
28 PROCESS_POOL_KWARGS_WHITELIST = ["max_workers", "mp_context", "initializer", "initargs"]
29 
30 
31 class WrappedProcessPoolExecutor(ProcessPoolExecutor):
32     """
33     Wraps the ProcessPoolExecutor to add various features:
34     - map_to_futures and map_unordered method
35     - pickling of job's output (see output_pickle_path_getter and output_pickle_path)
36     - job submission via pickling to circumvent bug in python < 3.8 (see MULTIPROCESSING_VIA_IO_TMP_DIR)
37     """
38 
39     def __init__(self, **kwargs):
40         new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)
41 
42         self.did_overwrite_start_method = False
43         if kwargs.get("start_method", None) is not None:
44             self.did_overwrite_start_method = True
45             self.old_start_method = multiprocessing.get_start_method()
46             start_method = kwargs["start_method"]
47             logging.info(
48                 f"Overwriting start_method to {start_method}. Previous value: {self.old_start_method}"
49             )
50             multiprocessing.set_start_method(start_method, force=True)
51 
52         ProcessPoolExecutor.__init__(self, **new_kwargs)
53 
54     def shutdown(self, *args, **kwargs):
55 
56         super().shutdown(*args, **kwargs)
57 
58         if self.did_overwrite_start_method:
59             logging.info(
60                 f"Restoring start_method to original value: {self.old_start_method}."
61             )
62             multiprocessing.set_start_method(self.old_start_method, force=True)
63             self.old_start_method = None
64             self.did_overwrite_start_method = False
65 
66     def submit(self, *args, **kwargs):
67 
68         output_pickle_path = None
69         if "__cfut_options" in kwargs:
70             output_pickle_path = kwargs["__cfut_options"]["output_pickle_path"]
71             del kwargs["__cfut_options"]
72 
73         if os.environ.get("MULTIPROCESSING_VIA_IO"):
74             # If MULTIPROCESSING_VIA_IO is set, _submit_via_io is used to
75             # workaround size constraints in pythons multiprocessing
76             # implementation. Also see https://github.com/python/cpython/pull/10305/files
77             # This should be fixed in python 3.8
78             submit_fn = self._submit_via_io
79         else:
80             submit_fn = super().submit
81 
82         # Depending on the start_method and output_pickle_path, wrapper functions may need to be
83         # executed in the new process context, before the actual code is ran.
84         # These wrapper functions consume their arguments from *args, **kwargs and assume
85         # that the next argument will be another function that is then called.
86         # The call_stack holds all of these wrapper functions and their arguments in the correct order.
87         # For example, call_stack = [wrapper_fn_1, wrapper_fn_1_arg_1, wrapper_fn_2, actual_fn, actual_fn_arg_1]
88         # where wrapper_fn_1 is called, which eventually calls wrapper_fn_2, which eventually calls actual_fn.
89         call_stack = []
90 
91         if multiprocessing.get_start_method() != "fork":
92             # If a start_method other than the default "fork" is used, logging needs to be re-setup,
93             # because the programming context is not inherited in those cases.
94             multiprocessing_logging_setup_fn = get_multiprocessing_logging_setup_fn()
95             call_stack.extend(
96                 [
97                     WrappedProcessPoolExecutor._setup_logging_and_execute,
98                     multiprocessing_logging_setup_fn,
99                 ]
100             )
101 
102         if output_pickle_path is not None:
103             call_stack.extend(
104                 [
105                     WrappedProcessPoolExecutor._execute_and_persist_function,
106                     output_pickle_path,
107                 ]
108             )
109 
110         fut = submit_fn(*call_stack, *args, **kwargs)
111 
112         enrich_future_with_uncaught_warning(fut)
113         return fut
114 
115     def _submit_via_io(self, *args, **kwargs):
116 
117         func = args[0]
118         args = args[1:]
119 
120         opt_tmp_dir = os.environ.get("MULTIPROCESSING_VIA_IO_TMP_DIR")
121         if opt_tmp_dir is not None:
122             dirpath = tempfile.mkdtemp(dir=opt_tmp_dir)
123         else:
124             dirpath = tempfile.mkdtemp()
125 
126         output_pickle_path = Path(dirpath) / "jobdescription.pickle"
127 
128         with open(output_pickle_path, "wb") as file:
129             pickling.dump((func, args, kwargs), file)
130 
131         future = super().submit(
132             WrappedProcessPoolExecutor._execute_via_io, output_pickle_path
133         )
134 
135         future.add_done_callback(
136             partial(WrappedProcessPoolExecutor._remove_tmp_file, dirpath)
137         )
138 
139         return future
140 
141     @staticmethod
142     def _remove_tmp_file(path, _future):
143 
144         shutil.rmtree(path)
145 
146     @staticmethod
147     def _setup_logging_and_execute(multiprocessing_logging_setup_fn, *args, **kwargs):
148 
149         func = args[0]
150         args = args[1:]
151 
152         multiprocessing_logging_setup_fn()
153 
154         return func(*args, **kwargs)
155 
156     @staticmethod
157     def _execute_via_io(serialized_function_info_path):
158 
159         with open(serialized_function_info_path, "rb") as file:
160             (func, args, kwargs) = pickling.load(file)
161         return func(*args, **kwargs)
162 
163     @staticmethod
164     def _execute_and_persist_function(output_pickle_path, *args, **kwargs):
165 
166         func = args[0]
167         args = args[1:]
168 
169         result = func(*args, **kwargs)
170 
171         with open(output_pickle_path, "wb") as file:
172             pickling.dump(result, file)
173 
174         return result
175 
176     def map_unordered(self, func, args):
177 
178         futs = self.map_to_futures(func, args)
179 
180         # Return a separate generator to avoid that map_unordered
181         # is executed lazily (otherwise, jobs would be submitted
182         # lazily, as well).
183         def result_generator():
184             for fut in futures.as_completed(futs):
185                 yield fut.result()
186 
187         return result_generator()
188 
189     def map_to_futures(self, func, args, output_pickle_path_getter=None):
190 
191         if output_pickle_path_getter is not None:
192             futs = [
193                 self.submit(
194                     func,
195                     arg,
196                     __cfut_options={
197                         "output_pickle_path": output_pickle_path_getter(arg)
198                     },
199                 )
200                 for arg in args
201             ]
202         else:
203             futs = [self.submit(func, arg) for arg in args]
204 
205         return futs
206 
207     def forward_log(self, fut):
208         """
209         Similar to the cluster executor, this method Takes a future from which the log file is forwarded to the active
210         process. This method blocks as long as the future is not done.
211         """
212 
213         # Since the default behavior of process pool executors is to show the log in the main process
214         # we don't need to do anything except for blocking until the future is done.
215         return fut.result()
216 
217 
218 class SequentialExecutor(WrappedProcessPoolExecutor):
219     """
220     The same as WrappedProcessPoolExecutor, but always uses only one core. In essence,
221     this is a sequential executor approach, but it still makes use of the standard pool approach.
222     That way, switching between different executors should always work without any problems.
223     """
224 
225     def __init__(self, **kwargs):
226         kwargs["max_workers"] = 1
227         WrappedProcessPoolExecutor.__init__(self, **kwargs)
228 
229 
230 class DebugSequentialExecutor(SequentialExecutor):
231     """
232     Only use for debugging purposes. This executor does not spawn new processes for its jobs. Therefore,
233     setting breakpoint()'s should be possible without context-related problems.
234     """
235 
236     def submit(self, *args, **kwargs):
237 
238         output_pickle_path = None
239         if "__cfut_options" in kwargs:
240             output_pickle_path = kwargs["__cfut_options"]["output_pickle_path"]
241             del kwargs["__cfut_options"]
242 
243         if output_pickle_path is not None:
244             fut = self._blocking_submit(
245                 WrappedProcessPoolExecutor._execute_and_persist_function,
246                 output_pickle_path,
247                 *args,
248                 **kwargs,
249             )
250         else:
251             fut = self._blocking_submit(*args, **kwargs)
252 
253         enrich_future_with_uncaught_warning(fut)
254         return fut
255 
256     def _blocking_submit(self, *args, **kwargs):
257 
258         func = args[0]
259         args = args[1:]
260 
261         fut = futures.Future()
262         result = func(*args, **kwargs)
263         fut.set_result(result)
264 
265         return fut
266 
267 
268 def pickle_identity(obj):
269     return pickling.loads(pickling.dumps(obj))
270 
271 
272 def pickle_identity_executor(func, *args, **kwargs):
273     result = func(*args, **kwargs)
274     return pickle_identity(result)
275 
276 
277 class PickleExecutor(WrappedProcessPoolExecutor):
278     """
279     The same as WrappedProcessPoolExecutor, but always pickles input and output of the jobs.
280     When using this executor for automated tests, it is ensured that using cluster executors in production
281     won't provoke pickling-related problems.
282     """
283 
284     def submit(self, _func, *_args, **_kwargs):
285 
286         (func, args, kwargs) = pickle_identity((_func, _args, _kwargs))
287         return super().submit(pickle_identity_executor, func, *args, **kwargs)
288 
289 
290 def noop():
291     return True
292 
293 
294 did_start_test_multiprocessing = False
295 
296 
297 def test_valid_multiprocessing():
298 
299     msg = """
300     ###############################################################
301     An attempt has been made to start a new process before the
302     current process has finished its bootstrapping phase.
303 
304     This probably means that you are not using fork to start your
305     child processes and you have forgotten to use the proper idiom
306     in the main module:
307 
308         if __name__ == '__main__':
309             main()
310             ...
311     ###############################################################
312     """
313 
314     with get_executor("multiprocessing") as executor:
315         try:
316             res_fut = executor.submit(noop)
317             assert res_fut.result() == True, msg
318         except RuntimeError as exc:
319             raise Exception(msg) from exc
320         except EOFError as exc:
321             raise Exception(msg) from exc
322 
323 
324 def get_executor(environment, **kwargs):
325 
326     if environment == "slurm":
327         return SlurmExecutor(**kwargs)
328     elif environment == "pbs":
329         return PBSExecutor(**kwargs)
330     elif environment == "kubernetes":
331         return KubernetesExecutor(**kwargs)
332     elif environment == "multiprocessing":
333         global did_start_test_multiprocessing
334         if not did_start_test_multiprocessing:
335             did_start_test_multiprocessing = True
336             test_valid_multiprocessing()
337 
338         return WrappedProcessPoolExecutor(**kwargs)
339     elif environment == "sequential":
340         return SequentialExecutor(**kwargs)
341     elif environment == "debug_sequential":
342         return DebugSequentialExecutor(**kwargs)
343     elif environment == "test_pickling":
344         return PickleExecutor(**kwargs)
345     raise Exception("Unknown executor: {}".format(environment))
346
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
     """
 
     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 @@
         # 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()

{"golden_diff": "diff --git a/cluster_tools/cluster_tools/__init__.py b/cluster_tools/cluster_tools/__init__.py\n--- a/cluster_tools/cluster_tools/__init__.py\n+++ b/cluster_tools/cluster_tools/__init__.py\n@@ -25,7 +25,7 @@\n     return new_kwargs\n \n \n-PROCESS_POOL_KWARGS_WHITELIST = [\"max_workers\", \"mp_context\", \"initializer\", \"initargs\"]\n+PROCESS_POOL_KWARGS_WHITELIST = [\"max_workers\", \"initializer\", \"initargs\"]\n \n \n class WrappedProcessPoolExecutor(ProcessPoolExecutor):\n@@ -37,31 +37,28 @@\n     \"\"\"\n \n     def __init__(self, **kwargs):\n-        new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)\n+        assert (not \"start_method\" in kwargs or kwargs[\"start_method\"] is None) or (\n+            not \"mp_context\" in kwargs\n+        ), \"Cannot use both `start_method` and `mp_context` kwargs.\"\n \n-        self.did_overwrite_start_method = False\n-        if kwargs.get(\"start_method\", None) is not None:\n-            self.did_overwrite_start_method = True\n-            self.old_start_method = multiprocessing.get_start_method()\n-            start_method = kwargs[\"start_method\"]\n-            logging.info(\n-                f\"Overwriting start_method to {start_method}. Previous value: {self.old_start_method}\"\n-            )\n-            multiprocessing.set_start_method(start_method, force=True)\n+        new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)\n \n-        ProcessPoolExecutor.__init__(self, **new_kwargs)\n+        mp_context = None\n \n-    def shutdown(self, *args, **kwargs):\n+        if \"mp_context\" in kwargs:\n+            mp_context = kwargs[\"mp_context\"]\n+        elif \"start_method\" in kwargs and kwargs[\"start_method\"] is not None:\n+            mp_context = multiprocessing.get_context(kwargs[\"start_method\"])\n+        elif \"MULTIPROCESSING_DEFAULT_START_METHOD\" in os.environ:\n+            mp_context = multiprocessing.get_context(\n+                os.environ[\"MULTIPROCESSING_DEFAULT_START_METHOD\"]\n+            )\n+        else:\n+            mp_context = multiprocessing.get_context(\"spawn\")\n \n-        super().shutdown(*args, **kwargs)\n+        new_kwargs[\"mp_context\"] = mp_context\n \n-        if self.did_overwrite_start_method:\n-            logging.info(\n-                f\"Restoring start_method to original value: {self.old_start_method}.\"\n-            )\n-            multiprocessing.set_start_method(self.old_start_method, force=True)\n-            self.old_start_method = None\n-            self.did_overwrite_start_method = False\n+        ProcessPoolExecutor.__init__(self, **new_kwargs)\n \n     def submit(self, *args, **kwargs):\n \n@@ -88,7 +85,7 @@\n         # where wrapper_fn_1 is called, which eventually calls wrapper_fn_2, which eventually calls actual_fn.\n         call_stack = []\n \n-        if multiprocessing.get_start_method() != \"fork\":\n+        if self._mp_context.get_start_method() != \"fork\":\n             # If a start_method other than the default \"fork\" is used, logging needs to be re-setup,\n             # because the programming context is not inherited in those cases.\n             multiprocessing_logging_setup_fn = get_multiprocessing_logging_setup_fn()\n", "issue": "Always use `spawn` in multiprocessing\nI 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.\r\n\r\nI 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`?\n", "before_files": [{"content": "import logging\nimport multiprocessing\nimport os\nimport shutil\nimport tempfile\nfrom concurrent import futures\nfrom concurrent.futures import ProcessPoolExecutor\nfrom functools import partial\nfrom pathlib import Path\n\nfrom . import pickling\nfrom .multiprocessing_logging_handler import get_multiprocessing_logging_setup_fn\nfrom .schedulers.kube import KubernetesExecutor\nfrom .schedulers.pbs import PBSExecutor\nfrom .schedulers.slurm import SlurmExecutor\nfrom .util import enrich_future_with_uncaught_warning\n\n\ndef get_existent_kwargs_subset(whitelist, kwargs):\n    new_kwargs = {}\n    for arg_name in whitelist:\n        if arg_name in kwargs:\n            new_kwargs[arg_name] = kwargs[arg_name]\n\n    return new_kwargs\n\n\nPROCESS_POOL_KWARGS_WHITELIST = [\"max_workers\", \"mp_context\", \"initializer\", \"initargs\"]\n\n\nclass WrappedProcessPoolExecutor(ProcessPoolExecutor):\n    \"\"\"\n    Wraps the ProcessPoolExecutor to add various features:\n    - map_to_futures and map_unordered method\n    - pickling of job's output (see output_pickle_path_getter and output_pickle_path)\n    - job submission via pickling to circumvent bug in python < 3.8 (see MULTIPROCESSING_VIA_IO_TMP_DIR)\n    \"\"\"\n\n    def __init__(self, **kwargs):\n        new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)\n\n        self.did_overwrite_start_method = False\n        if kwargs.get(\"start_method\", None) is not None:\n            self.did_overwrite_start_method = True\n            self.old_start_method = multiprocessing.get_start_method()\n            start_method = kwargs[\"start_method\"]\n            logging.info(\n                f\"Overwriting start_method to {start_method}. Previous value: {self.old_start_method}\"\n            )\n            multiprocessing.set_start_method(start_method, force=True)\n\n        ProcessPoolExecutor.__init__(self, **new_kwargs)\n\n    def shutdown(self, *args, **kwargs):\n\n        super().shutdown(*args, **kwargs)\n\n        if self.did_overwrite_start_method:\n            logging.info(\n                f\"Restoring start_method to original value: {self.old_start_method}.\"\n            )\n            multiprocessing.set_start_method(self.old_start_method, force=True)\n            self.old_start_method = None\n            self.did_overwrite_start_method = False\n\n    def submit(self, *args, **kwargs):\n\n        output_pickle_path = None\n        if \"__cfut_options\" in kwargs:\n            output_pickle_path = kwargs[\"__cfut_options\"][\"output_pickle_path\"]\n            del kwargs[\"__cfut_options\"]\n\n        if os.environ.get(\"MULTIPROCESSING_VIA_IO\"):\n            # If MULTIPROCESSING_VIA_IO is set, _submit_via_io is used to\n            # workaround size constraints in pythons multiprocessing\n            # implementation. Also see https://github.com/python/cpython/pull/10305/files\n            # This should be fixed in python 3.8\n            submit_fn = self._submit_via_io\n        else:\n            submit_fn = super().submit\n\n        # Depending on the start_method and output_pickle_path, wrapper functions may need to be\n        # executed in the new process context, before the actual code is ran.\n        # These wrapper functions consume their arguments from *args, **kwargs and assume\n        # that the next argument will be another function that is then called.\n        # The call_stack holds all of these wrapper functions and their arguments in the correct order.\n        # For example, call_stack = [wrapper_fn_1, wrapper_fn_1_arg_1, wrapper_fn_2, actual_fn, actual_fn_arg_1]\n        # where wrapper_fn_1 is called, which eventually calls wrapper_fn_2, which eventually calls actual_fn.\n        call_stack = []\n\n        if multiprocessing.get_start_method() != \"fork\":\n            # If a start_method other than the default \"fork\" is used, logging needs to be re-setup,\n            # because the programming context is not inherited in those cases.\n            multiprocessing_logging_setup_fn = get_multiprocessing_logging_setup_fn()\n            call_stack.extend(\n                [\n                    WrappedProcessPoolExecutor._setup_logging_and_execute,\n                    multiprocessing_logging_setup_fn,\n                ]\n            )\n\n        if output_pickle_path is not None:\n            call_stack.extend(\n                [\n                    WrappedProcessPoolExecutor._execute_and_persist_function,\n                    output_pickle_path,\n                ]\n            )\n\n        fut = submit_fn(*call_stack, *args, **kwargs)\n\n        enrich_future_with_uncaught_warning(fut)\n        return fut\n\n    def _submit_via_io(self, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        opt_tmp_dir = os.environ.get(\"MULTIPROCESSING_VIA_IO_TMP_DIR\")\n        if opt_tmp_dir is not None:\n            dirpath = tempfile.mkdtemp(dir=opt_tmp_dir)\n        else:\n            dirpath = tempfile.mkdtemp()\n\n        output_pickle_path = Path(dirpath) / \"jobdescription.pickle\"\n\n        with open(output_pickle_path, \"wb\") as file:\n            pickling.dump((func, args, kwargs), file)\n\n        future = super().submit(\n            WrappedProcessPoolExecutor._execute_via_io, output_pickle_path\n        )\n\n        future.add_done_callback(\n            partial(WrappedProcessPoolExecutor._remove_tmp_file, dirpath)\n        )\n\n        return future\n\n    @staticmethod\n    def _remove_tmp_file(path, _future):\n\n        shutil.rmtree(path)\n\n    @staticmethod\n    def _setup_logging_and_execute(multiprocessing_logging_setup_fn, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        multiprocessing_logging_setup_fn()\n\n        return func(*args, **kwargs)\n\n    @staticmethod\n    def _execute_via_io(serialized_function_info_path):\n\n        with open(serialized_function_info_path, \"rb\") as file:\n            (func, args, kwargs) = pickling.load(file)\n        return func(*args, **kwargs)\n\n    @staticmethod\n    def _execute_and_persist_function(output_pickle_path, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        result = func(*args, **kwargs)\n\n        with open(output_pickle_path, \"wb\") as file:\n            pickling.dump(result, file)\n\n        return result\n\n    def map_unordered(self, func, args):\n\n        futs = self.map_to_futures(func, args)\n\n        # Return a separate generator to avoid that map_unordered\n        # is executed lazily (otherwise, jobs would be submitted\n        # lazily, as well).\n        def result_generator():\n            for fut in futures.as_completed(futs):\n                yield fut.result()\n\n        return result_generator()\n\n    def map_to_futures(self, func, args, output_pickle_path_getter=None):\n\n        if output_pickle_path_getter is not None:\n            futs = [\n                self.submit(\n                    func,\n                    arg,\n                    __cfut_options={\n                        \"output_pickle_path\": output_pickle_path_getter(arg)\n                    },\n                )\n                for arg in args\n            ]\n        else:\n            futs = [self.submit(func, arg) for arg in args]\n\n        return futs\n\n    def forward_log(self, fut):\n        \"\"\"\n        Similar to the cluster executor, this method Takes a future from which the log file is forwarded to the active\n        process. This method blocks as long as the future is not done.\n        \"\"\"\n\n        # Since the default behavior of process pool executors is to show the log in the main process\n        # we don't need to do anything except for blocking until the future is done.\n        return fut.result()\n\n\nclass SequentialExecutor(WrappedProcessPoolExecutor):\n    \"\"\"\n    The same as WrappedProcessPoolExecutor, but always uses only one core. In essence,\n    this is a sequential executor approach, but it still makes use of the standard pool approach.\n    That way, switching between different executors should always work without any problems.\n    \"\"\"\n\n    def __init__(self, **kwargs):\n        kwargs[\"max_workers\"] = 1\n        WrappedProcessPoolExecutor.__init__(self, **kwargs)\n\n\nclass DebugSequentialExecutor(SequentialExecutor):\n    \"\"\"\n    Only use for debugging purposes. This executor does not spawn new processes for its jobs. Therefore,\n    setting breakpoint()'s should be possible without context-related problems.\n    \"\"\"\n\n    def submit(self, *args, **kwargs):\n\n        output_pickle_path = None\n        if \"__cfut_options\" in kwargs:\n            output_pickle_path = kwargs[\"__cfut_options\"][\"output_pickle_path\"]\n            del kwargs[\"__cfut_options\"]\n\n        if output_pickle_path is not None:\n            fut = self._blocking_submit(\n                WrappedProcessPoolExecutor._execute_and_persist_function,\n                output_pickle_path,\n                *args,\n                **kwargs,\n            )\n        else:\n            fut = self._blocking_submit(*args, **kwargs)\n\n        enrich_future_with_uncaught_warning(fut)\n        return fut\n\n    def _blocking_submit(self, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        fut = futures.Future()\n        result = func(*args, **kwargs)\n        fut.set_result(result)\n\n        return fut\n\n\ndef pickle_identity(obj):\n    return pickling.loads(pickling.dumps(obj))\n\n\ndef pickle_identity_executor(func, *args, **kwargs):\n    result = func(*args, **kwargs)\n    return pickle_identity(result)\n\n\nclass PickleExecutor(WrappedProcessPoolExecutor):\n    \"\"\"\n    The same as WrappedProcessPoolExecutor, but always pickles input and output of the jobs.\n    When using this executor for automated tests, it is ensured that using cluster executors in production\n    won't provoke pickling-related problems.\n    \"\"\"\n\n    def submit(self, _func, *_args, **_kwargs):\n\n        (func, args, kwargs) = pickle_identity((_func, _args, _kwargs))\n        return super().submit(pickle_identity_executor, func, *args, **kwargs)\n\n\ndef noop():\n    return True\n\n\ndid_start_test_multiprocessing = False\n\n\ndef test_valid_multiprocessing():\n\n    msg = \"\"\"\n    ###############################################################\n    An attempt has been made to start a new process before the\n    current process has finished its bootstrapping phase.\n\n    This probably means that you are not using fork to start your\n    child processes and you have forgotten to use the proper idiom\n    in the main module:\n\n        if __name__ == '__main__':\n            main()\n            ...\n    ###############################################################\n    \"\"\"\n\n    with get_executor(\"multiprocessing\") as executor:\n        try:\n            res_fut = executor.submit(noop)\n            assert res_fut.result() == True, msg\n        except RuntimeError as exc:\n            raise Exception(msg) from exc\n        except EOFError as exc:\n            raise Exception(msg) from exc\n\n\ndef get_executor(environment, **kwargs):\n\n    if environment == \"slurm\":\n        return SlurmExecutor(**kwargs)\n    elif environment == \"pbs\":\n        return PBSExecutor(**kwargs)\n    elif environment == \"kubernetes\":\n        return KubernetesExecutor(**kwargs)\n    elif environment == \"multiprocessing\":\n        global did_start_test_multiprocessing\n        if not did_start_test_multiprocessing:\n            did_start_test_multiprocessing = True\n            test_valid_multiprocessing()\n\n        return WrappedProcessPoolExecutor(**kwargs)\n    elif environment == \"sequential\":\n        return SequentialExecutor(**kwargs)\n    elif environment == \"debug_sequential\":\n        return DebugSequentialExecutor(**kwargs)\n    elif environment == \"test_pickling\":\n        return PickleExecutor(**kwargs)\n    raise Exception(\"Unknown executor: {}\".format(environment))\n", "path": "cluster_tools/cluster_tools/__init__.py"}], "after_files": [{"content": "import logging\nimport multiprocessing\nimport os\nimport shutil\nimport tempfile\nfrom concurrent import futures\nfrom concurrent.futures import ProcessPoolExecutor\nfrom functools import partial\nfrom pathlib import Path\n\nfrom . import pickling\nfrom .multiprocessing_logging_handler import get_multiprocessing_logging_setup_fn\nfrom .schedulers.kube import KubernetesExecutor\nfrom .schedulers.pbs import PBSExecutor\nfrom .schedulers.slurm import SlurmExecutor\nfrom .util import enrich_future_with_uncaught_warning\n\n\ndef get_existent_kwargs_subset(whitelist, kwargs):\n    new_kwargs = {}\n    for arg_name in whitelist:\n        if arg_name in kwargs:\n            new_kwargs[arg_name] = kwargs[arg_name]\n\n    return new_kwargs\n\n\nPROCESS_POOL_KWARGS_WHITELIST = [\"max_workers\", \"initializer\", \"initargs\"]\n\n\nclass WrappedProcessPoolExecutor(ProcessPoolExecutor):\n    \"\"\"\n    Wraps the ProcessPoolExecutor to add various features:\n    - map_to_futures and map_unordered method\n    - pickling of job's output (see output_pickle_path_getter and output_pickle_path)\n    - job submission via pickling to circumvent bug in python < 3.8 (see MULTIPROCESSING_VIA_IO_TMP_DIR)\n    \"\"\"\n\n    def __init__(self, **kwargs):\n        assert (not \"start_method\" in kwargs or kwargs[\"start_method\"] is None) or (\n            not \"mp_context\" in kwargs\n        ), \"Cannot use both `start_method` and `mp_context` kwargs.\"\n\n        new_kwargs = get_existent_kwargs_subset(PROCESS_POOL_KWARGS_WHITELIST, kwargs)\n\n        mp_context = None\n\n        if \"mp_context\" in kwargs:\n            mp_context = kwargs[\"mp_context\"]\n        elif \"start_method\" in kwargs and kwargs[\"start_method\"] is not None:\n            mp_context = multiprocessing.get_context(kwargs[\"start_method\"])\n        elif \"MULTIPROCESSING_DEFAULT_START_METHOD\" in os.environ:\n            mp_context = multiprocessing.get_context(\n                os.environ[\"MULTIPROCESSING_DEFAULT_START_METHOD\"]\n            )\n        else:\n            mp_context = multiprocessing.get_context(\"spawn\")\n\n        new_kwargs[\"mp_context\"] = mp_context\n\n        ProcessPoolExecutor.__init__(self, **new_kwargs)\n\n    def submit(self, *args, **kwargs):\n\n        output_pickle_path = None\n        if \"__cfut_options\" in kwargs:\n            output_pickle_path = kwargs[\"__cfut_options\"][\"output_pickle_path\"]\n            del kwargs[\"__cfut_options\"]\n\n        if os.environ.get(\"MULTIPROCESSING_VIA_IO\"):\n            # If MULTIPROCESSING_VIA_IO is set, _submit_via_io is used to\n            # workaround size constraints in pythons multiprocessing\n            # implementation. Also see https://github.com/python/cpython/pull/10305/files\n            # This should be fixed in python 3.8\n            submit_fn = self._submit_via_io\n        else:\n            submit_fn = super().submit\n\n        # Depending on the start_method and output_pickle_path, wrapper functions may need to be\n        # executed in the new process context, before the actual code is ran.\n        # These wrapper functions consume their arguments from *args, **kwargs and assume\n        # that the next argument will be another function that is then called.\n        # The call_stack holds all of these wrapper functions and their arguments in the correct order.\n        # For example, call_stack = [wrapper_fn_1, wrapper_fn_1_arg_1, wrapper_fn_2, actual_fn, actual_fn_arg_1]\n        # where wrapper_fn_1 is called, which eventually calls wrapper_fn_2, which eventually calls actual_fn.\n        call_stack = []\n\n        if self._mp_context.get_start_method() != \"fork\":\n            # If a start_method other than the default \"fork\" is used, logging needs to be re-setup,\n            # because the programming context is not inherited in those cases.\n            multiprocessing_logging_setup_fn = get_multiprocessing_logging_setup_fn()\n            call_stack.extend(\n                [\n                    WrappedProcessPoolExecutor._setup_logging_and_execute,\n                    multiprocessing_logging_setup_fn,\n                ]\n            )\n\n        if output_pickle_path is not None:\n            call_stack.extend(\n                [\n                    WrappedProcessPoolExecutor._execute_and_persist_function,\n                    output_pickle_path,\n                ]\n            )\n\n        fut = submit_fn(*call_stack, *args, **kwargs)\n\n        enrich_future_with_uncaught_warning(fut)\n        return fut\n\n    def _submit_via_io(self, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        opt_tmp_dir = os.environ.get(\"MULTIPROCESSING_VIA_IO_TMP_DIR\")\n        if opt_tmp_dir is not None:\n            dirpath = tempfile.mkdtemp(dir=opt_tmp_dir)\n        else:\n            dirpath = tempfile.mkdtemp()\n\n        output_pickle_path = Path(dirpath) / \"jobdescription.pickle\"\n\n        with open(output_pickle_path, \"wb\") as file:\n            pickling.dump((func, args, kwargs), file)\n\n        future = super().submit(\n            WrappedProcessPoolExecutor._execute_via_io, output_pickle_path\n        )\n\n        future.add_done_callback(\n            partial(WrappedProcessPoolExecutor._remove_tmp_file, dirpath)\n        )\n\n        return future\n\n    @staticmethod\n    def _remove_tmp_file(path, _future):\n\n        shutil.rmtree(path)\n\n    @staticmethod\n    def _setup_logging_and_execute(multiprocessing_logging_setup_fn, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        multiprocessing_logging_setup_fn()\n\n        return func(*args, **kwargs)\n\n    @staticmethod\n    def _execute_via_io(serialized_function_info_path):\n\n        with open(serialized_function_info_path, \"rb\") as file:\n            (func, args, kwargs) = pickling.load(file)\n        return func(*args, **kwargs)\n\n    @staticmethod\n    def _execute_and_persist_function(output_pickle_path, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        result = func(*args, **kwargs)\n\n        with open(output_pickle_path, \"wb\") as file:\n            pickling.dump(result, file)\n\n        return result\n\n    def map_unordered(self, func, args):\n\n        futs = self.map_to_futures(func, args)\n\n        # Return a separate generator to avoid that map_unordered\n        # is executed lazily (otherwise, jobs would be submitted\n        # lazily, as well).\n        def result_generator():\n            for fut in futures.as_completed(futs):\n                yield fut.result()\n\n        return result_generator()\n\n    def map_to_futures(self, func, args, output_pickle_path_getter=None):\n\n        if output_pickle_path_getter is not None:\n            futs = [\n                self.submit(\n                    func,\n                    arg,\n                    __cfut_options={\n                        \"output_pickle_path\": output_pickle_path_getter(arg)\n                    },\n                )\n                for arg in args\n            ]\n        else:\n            futs = [self.submit(func, arg) for arg in args]\n\n        return futs\n\n    def forward_log(self, fut):\n        \"\"\"\n        Similar to the cluster executor, this method Takes a future from which the log file is forwarded to the active\n        process. This method blocks as long as the future is not done.\n        \"\"\"\n\n        # Since the default behavior of process pool executors is to show the log in the main process\n        # we don't need to do anything except for blocking until the future is done.\n        return fut.result()\n\n\nclass SequentialExecutor(WrappedProcessPoolExecutor):\n    \"\"\"\n    The same as WrappedProcessPoolExecutor, but always uses only one core. In essence,\n    this is a sequential executor approach, but it still makes use of the standard pool approach.\n    That way, switching between different executors should always work without any problems.\n    \"\"\"\n\n    def __init__(self, **kwargs):\n        kwargs[\"max_workers\"] = 1\n        WrappedProcessPoolExecutor.__init__(self, **kwargs)\n\n\nclass DebugSequentialExecutor(SequentialExecutor):\n    \"\"\"\n    Only use for debugging purposes. This executor does not spawn new processes for its jobs. Therefore,\n    setting breakpoint()'s should be possible without context-related problems.\n    \"\"\"\n\n    def submit(self, *args, **kwargs):\n\n        output_pickle_path = None\n        if \"__cfut_options\" in kwargs:\n            output_pickle_path = kwargs[\"__cfut_options\"][\"output_pickle_path\"]\n            del kwargs[\"__cfut_options\"]\n\n        if output_pickle_path is not None:\n            fut = self._blocking_submit(\n                WrappedProcessPoolExecutor._execute_and_persist_function,\n                output_pickle_path,\n                *args,\n                **kwargs,\n            )\n        else:\n            fut = self._blocking_submit(*args, **kwargs)\n\n        enrich_future_with_uncaught_warning(fut)\n        return fut\n\n    def _blocking_submit(self, *args, **kwargs):\n\n        func = args[0]\n        args = args[1:]\n\n        fut = futures.Future()\n        result = func(*args, **kwargs)\n        fut.set_result(result)\n\n        return fut\n\n\ndef pickle_identity(obj):\n    return pickling.loads(pickling.dumps(obj))\n\n\ndef pickle_identity_executor(func, *args, **kwargs):\n    result = func(*args, **kwargs)\n    return pickle_identity(result)\n\n\nclass PickleExecutor(WrappedProcessPoolExecutor):\n    \"\"\"\n    The same as WrappedProcessPoolExecutor, but always pickles input and output of the jobs.\n    When using this executor for automated tests, it is ensured that using cluster executors in production\n    won't provoke pickling-related problems.\n    \"\"\"\n\n    def submit(self, _func, *_args, **_kwargs):\n\n        (func, args, kwargs) = pickle_identity((_func, _args, _kwargs))\n        return super().submit(pickle_identity_executor, func, *args, **kwargs)\n\n\ndef noop():\n    return True\n\n\ndid_start_test_multiprocessing = False\n\n\ndef test_valid_multiprocessing():\n\n    msg = \"\"\"\n    ###############################################################\n    An attempt has been made to start a new process before the\n    current process has finished its bootstrapping phase.\n\n    This probably means that you are not using fork to start your\n    child processes and you have forgotten to use the proper idiom\n    in the main module:\n\n        if __name__ == '__main__':\n            main()\n            ...\n    ###############################################################\n    \"\"\"\n\n    with get_executor(\"multiprocessing\") as executor:\n        try:\n            res_fut = executor.submit(noop)\n            assert res_fut.result() == True, msg\n        except RuntimeError as exc:\n            raise Exception(msg) from exc\n        except EOFError as exc:\n            raise Exception(msg) from exc\n\n\ndef get_executor(environment, **kwargs):\n\n    if environment == \"slurm\":\n        return SlurmExecutor(**kwargs)\n    elif environment == \"pbs\":\n        return PBSExecutor(**kwargs)\n    elif environment == \"kubernetes\":\n        return KubernetesExecutor(**kwargs)\n    elif environment == \"multiprocessing\":\n        global did_start_test_multiprocessing\n        if not did_start_test_multiprocessing:\n            did_start_test_multiprocessing = True\n            test_valid_multiprocessing()\n\n        return WrappedProcessPoolExecutor(**kwargs)\n    elif environment == \"sequential\":\n        return SequentialExecutor(**kwargs)\n    elif environment == \"debug_sequential\":\n        return DebugSequentialExecutor(**kwargs)\n    elif environment == \"test_pickling\":\n        return PickleExecutor(**kwargs)\n    raise Exception(\"Unknown executor: {}\".format(environment))\n", "path": "cluster_tools/cluster_tools/__init__.py"}]}