We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Warn when multiprocessing fails due to missing if __name__ guard
If wklibs are used with multiprocessing (e.g., calling `downsample` on a layer), python's multiprocessing module will import the main module (if `spawn` is used which is the default on OS  X). If that module has side effects (i.e., it is not guarded with `if __name__ == "__main__":`), weird errors can occur.

On start up, we could spawn a dummy job with the clustertools and if that fails, we could emit

```
            raise AssertionError(
                """
                #########################################
                #########################################
                #########################################
                #########################################
                #########################################
                ## Multiprocessing setup does not work ##
                ## properly. Please check that you are ##
                ## using an if __name__ == "__main__"  ##
                ## guard in your main module.          ##
                #########################################
                #########################################
                #########################################
                #########################################
                #########################################
            """
            ) 
```

to help the user with what's wrong.

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

{"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@@ -286,12 +286,52 @@\n         return super().submit(pickle_identity_executor, func, *args, **kwargs)\n \n \n+def noop():\n+    return True\n+\n+\n+did_start_test_multiprocessing = False\n+\n+\n+def 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+\n def get_executor(environment, **kwargs):\n+\n     if environment == \"slurm\":\n         return SlurmExecutor(**kwargs)\n     elif environment == \"pbs\":\n         return PBSExecutor(**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", "issue": "Warn when multiprocessing fails due to missing if __name__ guard\nIf wklibs are used with multiprocessing (e.g., calling `downsample` on a layer), python's multiprocessing module will import the main module (if `spawn` is used which is the default on OS  X). If that module has side effects (i.e., it is not guarded with `if __name__ == \"__main__\":`), weird errors can occur.\r\n\r\nOn start up, we could spawn a dummy job with the clustertools and if that fails, we could emit\r\n\r\n```\r\n            raise AssertionError(\r\n                \"\"\"\r\n                #########################################\r\n                #########################################\r\n                #########################################\r\n                #########################################\r\n                #########################################\r\n                ## Multiprocessing setup does not work ##\r\n                ## properly. Please check that you are ##\r\n                ## using an if __name__ == \"__main__\"  ##\r\n                ## guard in your main module.          ##\r\n                #########################################\r\n                #########################################\r\n                #########################################\r\n                #########################################\r\n                #########################################\r\n            \"\"\"\r\n            ) \r\n```\r\n\r\nto help the user with what's wrong.\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.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 get_executor(environment, **kwargs):\n    if environment == \"slurm\":\n        return SlurmExecutor(**kwargs)\n    elif environment == \"pbs\":\n        return PBSExecutor(**kwargs)\n    elif environment == \"multiprocessing\":\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.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 == \"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"}]}