replace fork+mock by spawn+hardcoded patch in main code

Author: mfeurer

autosklearn/ensemble_builder.py

@@ -156,6 +156,7 @@ def build_ensemble(
         self,
         dask_client: dask.distributed.Client,
         pynisher_context: str = 'spawn',
+        unit_test: bool = False
     ) -> None:
 
         # The second criteria is elapsed time
@@ -227,6 +228,7 @@ def build_ensemble(
                     priority=100,
                     pynisher_context=pynisher_context,
                     logger_port=self.logger_port,
+                    unit_test=unit_test,
                 ))
 
                 logger.info(
@@ -265,6 +267,7 @@ def fit_and_return_ensemble(
     return_predictions: bool,
     pynisher_context: str,
     logger_port: int = logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+    unit_test: bool = False,
 ) -> Tuple[
         List[Tuple[int, float, float, float]],
         int,
@@ -321,6 +324,11 @@ def fit_and_return_ensemble(
             Context to use for multiprocessing, can be either fork, spawn or forkserver.
         logger_port: int
             The port where the logging server is listening to.
+        unit_test: bool
+            Turn on unit testing mode. This currently makes fit_ensemble raise a MemoryError.
+            Having this is very bad coding style, but I did not find a way to make
+            unittest.mock work through the pynisher with all spawn contexts. If you know a
+            better solution, please let us know by opening an issue.
 
     Returns
     -------
@@ -343,6 +351,7 @@ def fit_and_return_ensemble(
         read_at_most=read_at_most,
         random_state=random_state,
         logger_port=logger_port,
+        unit_test=unit_test,
     ).run(
         end_at=end_at,
         iteration=iteration,
@@ -354,21 +363,22 @@ def fit_and_return_ensemble(
 
 class EnsembleBuilder(object):
     def __init__(
-            self,
-            backend: Backend,
-            dataset_name: str,
-            task_type: int,
-            metric: Scorer,
-            ensemble_size: int = 10,
-            ensemble_nbest: int = 100,
-            max_models_on_disc: int = 100,
-            performance_range_threshold: float = 0,
-            seed: int = 1,
-            precision: int = 32,
-            memory_limit: Optional[int] = 1024,
-            read_at_most: int = 5,
-            random_state: Optional[Union[int, np.random.RandomState]] = None,
-            logger_port: int = logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+        self,
+        backend: Backend,
+        dataset_name: str,
+        task_type: int,
+        metric: Scorer,
+        ensemble_size: int = 10,
+        ensemble_nbest: int = 100,
+        max_models_on_disc: int = 100,
+        performance_range_threshold: float = 0,
+        seed: int = 1,
+        precision: int = 32,
+        memory_limit: Optional[int] = 1024,
+        read_at_most: int = 5,
+        random_state: Optional[Union[int, np.random.RandomState]] = None,
+        logger_port: int = logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+        unit_test: bool = False,
     ):
         """
             Constructor
@@ -416,6 +426,11 @@ def __init__(
                 read at most n new prediction files in each iteration
             logger_port: int
                 port that receives logging records
+            unit_test: bool
+                Turn on unit testing mode. This currently makes fit_ensemble raise a MemoryError.
+                Having this is very bad coding style, but I did not find a way to make
+                unittest.mock work through the pynisher with all spawn contexts. If you know a
+                better solution, please let us know by opening an issue.
         """
 
         super(EnsembleBuilder, self).__init__()
@@ -454,6 +469,7 @@ def __init__(
         self.memory_limit = memory_limit
         self.read_at_most = read_at_most
         self.random_state = check_random_state(random_state)
+        self.unit_test = unit_test
 
         # Setup the logger
         self.logger_port = logger_port
@@ -1196,6 +1212,9 @@ def fit_ensemble(self, selected_keys: list):
                 trained Ensemble
         """
 
+        if self.unit_test:
+            raise MemoryError()
+
         predictions_train = [self.read_preds[k][Y_ENSEMBLE] for k in selected_keys]
         include_num_runs = [
             (

test/test_ensemble_builder/test_ensemble.py

@@ -764,9 +764,7 @@ def test_ensemble_builder_process_realrun(dask_client_single_worker, ensemble_ba
     assert history[0]['ensemble_test_score'] == 0.9
 
 
-@unittest.mock.patch('autosklearn.ensemble_builder.EnsembleBuilder.fit_ensemble')
 def test_ensemble_builder_nbest_remembered(
-    fit_ensemble,
     ensemble_backend,
     dask_client_single_worker,
 ):
@@ -775,8 +773,6 @@ def test_ensemble_builder_nbest_remembered(
     This way, we can remember it and not waste more time trying big ensemble sizes
     """
 
-    fit_ensemble.side_effect = MemoryError
-
     manager = EnsembleBuilderManager(
         start_time=time.time(),
         time_left_for_ensembles=1000,
@@ -795,15 +791,14 @@ def test_ensemble_builder_nbest_remembered(
         max_iterations=None,
     )
 
-    # Use fork context in the next line to allow for the mock to work
-    manager.build_ensemble(dask_client_single_worker, 'fork')
+    manager.build_ensemble(dask_client_single_worker, unit_test=True)
     future = manager.futures[0]
     dask.distributed.wait([future])  # wait for the ensemble process to finish
     assert future.result() == ([], 5, None, None, None)
     file_path = os.path.join(ensemble_backend.internals_directory, 'ensemble_read_preds.pkl')
     assert not os.path.exists(file_path)
 
-    manager.build_ensemble(dask_client_single_worker, 'fork')
+    manager.build_ensemble(dask_client_single_worker, unit_test=True)
 
     future = manager.futures[0]
     dask.distributed.wait([future])  # wait for the ensemble process to finish