Threads and forkserver (#1062)

* use threads again

* try the forkserver

* use forkserver pre-load for faster process starting

* streamline code

* de-duplicate code

* add missing file

* Update parallel.py

* Update parallel.py

Author: mfeurer

autosklearn/automl.py

@@ -49,6 +49,7 @@
     get_named_client_logger,
 )
 from autosklearn.util import pipeline, RE_PATTERN
+from autosklearn.util.parallel import preload_modules
 from autosklearn.ensemble_builder import EnsembleBuilderManager
 from autosklearn.ensembles.singlebest_ensemble import SingleBest
 from autosklearn.smbo import AutoMLSMBO
@@ -228,7 +229,7 @@ def __init__(self,
         # examples. Nevertheless, multi-process runs
         # have spawn as requirement to reduce the
         # possibility of a deadlock
-        self._multiprocessing_context = 'spawn'
+        self._multiprocessing_context = 'forkserver'
         if self._n_jobs == 1 and self._dask_client is None:
             self._multiprocessing_context = 'fork'
             self._dask_client = SingleThreadedClient()
@@ -248,11 +249,10 @@ def __init__(self,
 
     def _create_dask_client(self):
         self._is_dask_client_internally_created = True
-        dask.config.set({'distributed.worker.daemon': False})
         self._dask_client = dask.distributed.Client(
             dask.distributed.LocalCluster(
                 n_workers=self._n_jobs,
-                processes=True if self._n_jobs != 1 else False,
+                processes=False,
                 threads_per_worker=1,
                 # We use the temporal directory to save the
                 # dask workers, because deleting workers
@@ -299,8 +299,8 @@ def _get_logger(self, name):
         # under the above logging configuration setting
         # We need to specify the logger_name so that received records
         # are treated under the logger_name ROOT logger setting
-        context = multiprocessing.get_context(
-            self._multiprocessing_context)
+        context = multiprocessing.get_context(self._multiprocessing_context)
+        preload_modules(context)
         self.stop_logging_server = context.Event()
         port = context.Value('l')  # be safe by using a long
         port.value = -1

autosklearn/ensemble_builder.py

@@ -31,6 +31,7 @@
 from autosklearn.ensembles.ensemble_selection import EnsembleSelection
 from autosklearn.ensembles.abstract_ensemble import AbstractEnsemble
 from autosklearn.util.logging_ import get_named_client_logger
+from autosklearn.util.parallel import preload_modules
 
 Y_ENSEMBLE = 0
 Y_VALID = 1
@@ -572,11 +573,11 @@ def __init__(
     def run(
         self,
         iteration: int,
+        pynisher_context: str,
         time_left: Optional[float] = None,
         end_at: Optional[float] = None,
         time_buffer=5,
         return_predictions: bool = False,
-        pynisher_context: str = 'spawn',
     ):
 
         if time_left is None and end_at is None:
@@ -606,6 +607,7 @@ def run(
             if wall_time_in_s < 1:
                 break
             context = multiprocessing.get_context(pynisher_context)
+            preload_modules(context)
 
             safe_ensemble_script = pynisher.enforce_limits(
                 wall_time_in_s=wall_time_in_s,

autosklearn/evaluation/__init__.py

@@ -24,6 +24,7 @@
 import autosklearn.evaluation.test_evaluator
 import autosklearn.evaluation.util
 from autosklearn.util.logging_ import get_named_client_logger
+from autosklearn.util.parallel import preload_modules
 
 
 def fit_predict_try_except_decorator(ta, queue, cost_for_crash, **kwargs):
@@ -97,12 +98,12 @@ def _encode_exit_status(exit_status):
 class ExecuteTaFuncWithQueue(AbstractTAFunc):
 
     def __init__(self, backend, autosklearn_seed, resampling_strategy, metric,
-                 cost_for_crash, abort_on_first_run_crash, port,
+                 cost_for_crash, abort_on_first_run_crash, port, pynisher_context,
                  initial_num_run=1, stats=None,
                  run_obj='quality', par_factor=1, scoring_functions=None,
                  output_y_hat_optimization=True, include=None, exclude=None,
                  memory_limit=None, disable_file_output=False, init_params=None,
-                 budget_type=None, ta=False, pynisher_context='spawn', **resampling_strategy_args):
+                 budget_type=None, ta=False, **resampling_strategy_args):
 
         if resampling_strategy == 'holdout':
             eval_function = autosklearn.evaluation.train_evaluator.eval_holdout
@@ -261,6 +262,7 @@ def run(
     ) -> Tuple[StatusType, float, float, Dict[str, Union[int, float, str, Dict, List, Tuple]]]:
 
         context = multiprocessing.get_context(self.pynisher_context)
+        preload_modules(context)
         queue = context.Queue()
 
         if not (instance_specific is None or instance_specific == '0'):

autosklearn/util/parallel.py

@@ -0,0 +1,20 @@
+import multiprocessing
+import sys
+
+
+def preload_modules(context: multiprocessing.context.BaseContext) -> None:
+    all_loaded_modules = sys.modules.keys()
+    preload = [
+        loaded_module for loaded_module in all_loaded_modules
+        if loaded_module.split('.')[0] in (
+            'smac',
+            'autosklearn',
+            'numpy',
+            'scipy',
+            'pandas',
+            'pynisher',
+            'sklearn',
+            'ConfigSpace',
+        ) and 'logging' not in loaded_module
+    ]
+    context.set_forkserver_preload(preload)

scripts/run_auto-sklearn_for_metadata_generation.py

@@ -151,7 +151,8 @@
                                         include=include,
                                         metric=automl_arguments['metric'],
                                         cost_for_crash=get_cost_of_crash(automl_arguments['metric']),
-                                        abort_on_first_run_crash=False,)
+                                        abort_on_first_run_crash=False,
+                                        pynisher_context='fork')
             run_info, run_value = ta.run_wrapper(
                 RunInfo(
                     config=config,

test/conftest.py

@@ -3,7 +3,6 @@
 import time
 import unittest.mock
 
-import dask
 from dask.distributed import Client, get_client
 import psutil
 import pytest
@@ -125,8 +124,7 @@ def dask_client(request):
     Workers are in subprocesses to not create deadlocks with the pynisher and logging.
     """
 
-    dask.config.set({'distributed.worker.daemon': False})
-    client = Client(n_workers=2, threads_per_worker=1, processes=True)
+    client = Client(n_workers=2, threads_per_worker=1, processes=False)
     print("Started Dask client={}\n".format(client))
 
     def get_finalizer(address):
@@ -151,7 +149,6 @@ def dask_client_single_worker(request):
     it is used very rarely to avoid this issue as much as possible.
     """
 
-    dask.config.set({'distributed.worker.daemon': False})
     client = Client(n_workers=1, threads_per_worker=1, processes=False)
     print("Started Dask client={}\n".format(client))
 

test/test_ensemble_builder/test_ensemble.py

@@ -504,7 +504,7 @@ def test_run_end_at(ensemble_backend):
 
         current_time = time.time()
 
-        ensbuilder.run(end_at=current_time + 10, iteration=1)
+        ensbuilder.run(end_at=current_time + 10, iteration=1, pynisher_context='forkserver')
         # 4 seconds left because: 10 seconds - 5 seconds overhead - very little overhead,
         # but then rounded to an integer
         assert pynisher_mock.call_args_list[0][1]["wall_time_in_s"], 4
@@ -579,7 +579,7 @@ def mtime_mock(filename):
 
         # And then it still runs, but basically won't do anything any more except for raising error
         # messages via the logger
-        ensbuilder.run(time_left=1000, iteration=0)
+        ensbuilder.run(time_left=1000, iteration=0, pynisher_context='fork')
         assert os.path.exists(read_scores_file)
         assert not os.path.exists(read_preds_file)
         assert logger_mock.warning.call_count == 4

test/test_evaluation/test_evaluation.py

@@ -112,6 +112,7 @@ def test_zero_or_negative_cutoff(self, pynisher_mock):
                                     metric=accuracy,
                                     cost_for_crash=get_cost_of_crash(accuracy),
                                     abort_on_first_run_crash=False,
+                                    pynisher_context='forkserver',
                                     )
         self.scenario.wallclock_limit = 5
         self.stats.submitted_ta_runs += 1
@@ -130,6 +131,7 @@ def test_cutoff_lower_than_remaining_time(self, pynisher_mock):
                                     metric=accuracy,
                                     cost_for_crash=get_cost_of_crash(accuracy),
                                     abort_on_first_run_crash=False,
+                                    pynisher_context='forkserver',
                                     )
         self.stats.ta_runs = 1
         ta.run_wrapper(RunInfo(config=config, cutoff=30, instance=None, instance_specific=None,
@@ -224,6 +226,7 @@ def test_eval_with_limits_holdout_fail_timeout(self, pynisher_mock):
                                     metric=accuracy,
                                     cost_for_crash=get_cost_of_crash(accuracy),
                                     abort_on_first_run_crash=False,
+                                    pynisher_context='forkserver',
                                     )
         info = ta.run_wrapper(RunInfo(config=config, cutoff=30, instance=None,
                                       instance_specific=None, seed=1, capped=False))
@@ -259,6 +262,7 @@ def side_effect(**kwargs):
                                     metric=accuracy,
                                     cost_for_crash=get_cost_of_crash(accuracy),
                                     abort_on_first_run_crash=False,
+                                    pynisher_context='forkserver',
                                     )
         info = ta.run_wrapper(RunInfo(config=config, cutoff=30, instance=None,
                                       instance_specific=None, seed=1, capped=False))
@@ -282,6 +286,7 @@ def side_effect(**kwargs):
                                     metric=accuracy,
                                     cost_for_crash=get_cost_of_crash(accuracy),
                                     abort_on_first_run_crash=False,
+                                    pynisher_context='forkserver',
                                     )
         info = ta.run_wrapper(RunInfo(config=config, cutoff=30, instance=None,
                                       instance_specific=None, seed=1, capped=False))