Repository: automl/auto-sklearn

Issue:
On exception during fit script doesn't terminate cleanly
## Describe the bug ##
If an exception is triggered during the fit process of the `auto-sklearn` classifier the python interpreter process does NOT exit cleanly (needs a manual keyboard interrupt).

## To Reproduce ##
Steps to reproduce the behavior:
1. Write simple script to trigger exception 
```
import autosklearn.classification
import pandas as pd
automl = autosklearn.classification.AutoSklearnClassifier()
a = pd.DataFrame(columns=["id"])
b = pd.DataFrame(columns=["id", "y"])
a.set_index("id", inplace=True)
b.set_index("id", inplace=True)
combined_data = pd.merge(a, b, left_index=True, right_index=True)
x = combined_data.drop('y', axis=1)
y = combined_data['y']
automl.fit(x, y)
```
2. Execute script `python script.py`
3. Script shows this backtrace:
```
Traceback (most recent call last):
  File "script2.py", line 11, in <module>
    automl.fit(x, y)
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/estimators.py", line 1045, in fit
    super().fit(
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/estimators.py", line 375, in fit
    self.automl_.fit(load_models=self.load_models, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/automl.py", line 2056, in fit
    return super().fit(
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/automl.py", line 639, in fit
    self.InputValidator.fit(X_train=X, y_train=y, X_test=X_test, y_test=y_test)
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/data/validation.py", line 147, in fit
    self.target_validator.fit(y_train, y_test)
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/data/target_validator.py", line 119, in fit
    self._fit(y_train, y_test)
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/data/target_validator.py", line 182, in _fit
    self.encoder.fit(y_train.to_numpy().reshape(-1, 1))
  File "/usr/local/lib/python3.8/dist-packages/sklearn/preprocessing/_encoders.py", line 779, in fit
    self._fit(X)
  File "/usr/local/lib/python3.8/dist-packages/sklearn/preprocessing/_encoders.py", line 77, in _fit
    X_list, n_samples, n_features = self._check_X(
  File "/usr/local/lib/python3.8/dist-packages/sklearn/preprocessing/_encoders.py", line 44, in _check_X
    X_temp = check_array(X, dtype=None,
  File "/usr/local/lib/python3.8/dist-packages/sklearn/utils/validation.py", line 63, in inner_f
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/sklearn/utils/validation.py", line 726, in check_array
    raise ValueError("Found array with %d sample(s) (shape=%s) while a"
ValueError: Found array with 0 sample(s) (shape=(0, 1)) while a minimum of 1 is required.
```

## Expected behavior ##
After showing the stack trace, script should terminate, I expect the python interpreter to quit.

## Actual behavior, stacktrace or logfile ##
After showing the exception the script doesn't terminate, it's stuck processing some exit handlers.
If I send keyboard interrupt I get the following stack trace:
```
^CError in atexit._run_exitfuncs:
Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/popen_fork.py", line 27, in poll
Process ForkProcess-1:
    pid, sts = os.waitpid(self.pid, flag)
KeyboardInterrupt
Traceback (most recent call last):
  File "/usr/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/usr/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/util/logging_.py", line 320, in start_log_server
    receiver.serve_until_stopped()
  File "/usr/local/lib/python3.8/dist-packages/autosklearn/util/logging_.py", line 350, in serve_until_stopped
    rd, wr, ex = select.select([self.socket.fileno()],
KeyboardInterrupt
```

## Environment and installation: ##

Please give details about your installation:

- Ubuntu 20.04, virtual env, python 3.10
- Docker (https://hub.docker.com/r/mfeurer/auto-sklearn/ which is same Ubuntu version), python 3.8
- auto-sklearn version 0.14.7


Discussion:
Hey thanks for the report, you're right that this does happen and sorry for that. I guess this easiest fix is a big `try: except` in the main `fit` method such that we can properly close the logging server when an error occurs.

It might take a bit before we can put this in due to some other deadlines we are reaching for right now but it seems like an easy fix (hopefully). The code snippet will help test this a lot easier :)

Resolving patch:
diff --git a/autosklearn/automl.py b/autosklearn/automl.py
--- a/autosklearn/automl.py
+++ b/autosklearn/automl.py
@@ -646,273 +646,282 @@ def fit(
 
         # By default try to use the TCP logging port or get a new port
         self._logger_port = logging.handlers.DEFAULT_TCP_LOGGING_PORT
-        self._logger = self._get_logger(dataset_name)
 
-        # The first thing we have to do is create the logger to update the backend
-        self._backend.setup_logger(self._logger_port)
+        # Once we start the logging server, it starts in a new process
+        # If an error occurs then we want to make sure that we exit cleanly
+        # and shut it down, else it might hang
+        # https://github.com/automl/auto-sklearn/issues/1480
+        try:
+            self._logger = self._get_logger(dataset_name)
 
-        if not only_return_configuration_space:
-            # If only querying the configuration space, we do not save the start time
-            # The start time internally checks for the fit() method to execute only once
-            # But this does not apply when only querying the configuration space
-            self._backend.save_start_time(self._seed)
+            # The first thing we have to do is create the logger to update the backend
+            self._backend.setup_logger(self._logger_port)
 
-        self._stopwatch = StopWatch()
+            if not only_return_configuration_space:
+                # If only querying the configuration space, we do not save the start
+                # time The start time internally checks for the fit() method to execute
+                # only once but this does not apply when only querying the configuration
+                # space
+                self._backend.save_start_time(self._seed)
 
-        # Make sure that input is valid
-        # Performs Ordinal one hot encoding to the target
-        # both for train and test data
-        self.InputValidator = InputValidator(
-            is_classification=is_classification,
-            feat_type=feat_type,
-            logger_port=self._logger_port,
-            allow_string_features=self.allow_string_features,
-        )
-        self.InputValidator.fit(X_train=X, y_train=y, X_test=X_test, y_test=y_test)
-        X, y = self.InputValidator.transform(X, y)
+            self._stopwatch = StopWatch()
 
-        if X_test is not None and y_test is not None:
-            X_test, y_test = self.InputValidator.transform(X_test, y_test)
+            # Make sure that input is valid
+            # Performs Ordinal one hot encoding to the target
+            # both for train and test data
+            self.InputValidator = InputValidator(
+                is_classification=is_classification,
+                feat_type=feat_type,
+                logger_port=self._logger_port,
+                allow_string_features=self.allow_string_features,
+            )
+            self.InputValidator.fit(X_train=X, y_train=y, X_test=X_test, y_test=y_test)
+            X, y = self.InputValidator.transform(X, y)
 
-        # We don't support size reduction on pandas type object yet
-        if (
-            self._dataset_compression is not None
-            and not isinstance(X, pd.DataFrame)
-            and not (isinstance(y, pd.Series) or isinstance(y, pd.DataFrame))
-        ):
-            methods = self._dataset_compression["methods"]
-            memory_allocation = self._dataset_compression["memory_allocation"]
-
-            # Remove precision reduction if we can't perform it
-            if "precision" in methods and X.dtype not in supported_precision_reductions:
-                methods = [method for method in methods if method != "precision"]
-
-            with warnings_to(self._logger):
-                X, y = reduce_dataset_size_if_too_large(
-                    X=X,
-                    y=y,
-                    memory_limit=self._memory_limit,
-                    is_classification=is_classification,
-                    random_state=self._seed,
-                    operations=methods,
-                    memory_allocation=memory_allocation,
-                )
+            if X_test is not None and y_test is not None:
+                X_test, y_test = self.InputValidator.transform(X_test, y_test)
 
-        # Check the re-sampling strategy
-        try:
+            # We don't support size reduction on pandas type object yet
+            if (
+                self._dataset_compression is not None
+                and not isinstance(X, pd.DataFrame)
+                and not (isinstance(y, pd.Series) or isinstance(y, pd.DataFrame))
+            ):
+                methods = self._dataset_compression["methods"]
+                memory_allocation = self._dataset_compression["memory_allocation"]
+
+                # Remove precision reduction if we can't perform it
+                if (
+                    "precision" in methods
+                    and X.dtype not in supported_precision_reductions
+                ):
+                    methods = [method for method in methods if method != "precision"]
+
+                with warnings_to(self._logger):
+                    X, y = reduce_dataset_size_if_too_large(
+                        X=X,
+                        y=y,
+                        memory_limit=self._memory_limit,
+                        is_classification=is_classification,
+                        random_state=self._seed,
+                        operations=methods,
+                        memory_allocation=memory_allocation,
+                    )
+
+            # Check the re-sampling strategy
             self._check_resampling_strategy(
                 X=X,
                 y=y,
                 task=self._task,
             )
-        except Exception as e:
-            self._fit_cleanup()
-            raise e
-
-        # Reset learnt stuff
-        self.models_ = None
-        self.cv_models_ = None
-        self.ensemble_ = None
 
-        # The metric must exist as of this point
-        # It can be provided in the constructor, or automatically
-        # defined in the estimator fit call
-        if isinstance(self._metrics, Sequence):
-            for entry in self._metrics:
-                if not isinstance(entry, Scorer):
-                    raise ValueError(
-                        "Metric {entry} must be instance of autosklearn.metrics.Scorer."
-                    )
-        else:
-            raise ValueError(
-                "Metric must be a sequence of instances of "
-                "autosklearn.metrics.Scorer."
-            )
+            # Reset learnt stuff
+            self.models_ = None
+            self.cv_models_ = None
+            self.ensemble_ = None
 
-        # If no dask client was provided, we create one, so that we can
-        # start a ensemble process in parallel to smbo optimize
-        if self._dask_client is None and (
-            self._ensemble_class is not None
-            or self._n_jobs is not None
-            and self._n_jobs > 1
-        ):
-            self._create_dask_client()
-        else:
-            self._is_dask_client_internally_created = False
+            # The metric must exist as of this point
+            # It can be provided in the constructor, or automatically
+            # defined in the estimator fit call
+            if isinstance(self._metrics, Sequence):
+                for entry in self._metrics:
+                    if not isinstance(entry, Scorer):
+                        raise ValueError(
+                            f"Metric {entry} must be instance of"
+                            " autosklearn.metrics.Scorer."
+                        )
+            else:
+                raise ValueError(
+                    "Metric must be a sequence of instances of "
+                    "autosklearn.metrics.Scorer."
+                )
 
-        self._dataset_name = dataset_name
-        self._stopwatch.start(self._dataset_name)
+            # If no dask client was provided, we create one, so that we can
+            # start a ensemble process in parallel to smbo optimize
+            if self._dask_client is None and (
+                self._ensemble_class is not None
+                or self._n_jobs is not None
+                and self._n_jobs > 1
+            ):
+                self._create_dask_client()
+            else:
+                self._is_dask_client_internally_created = False
 
-        # Take the feature types from the validator
-        self._feat_type = self.InputValidator.feature_validator.feat_type
+            self._dataset_name = dataset_name
+            self._stopwatch.start(self._dataset_name)
 
-        self._log_fit_setup()
+            # Take the feature types from the validator
+            self._feat_type = self.InputValidator.feature_validator.feat_type
 
-        # == Pickle the data manager to speed up loading
-        with self._stopwatch.time("Save Datamanager"):
-            datamanager = XYDataManager(
-                X,
-                y,
-                X_test=X_test,
-                y_test=y_test,
-                task=self._task,
-                feat_type=self._feat_type,
-                dataset_name=dataset_name,
-            )
+            self._log_fit_setup()
 
-            self._backend._make_internals_directory()
-            self._label_num = datamanager.info["label_num"]
-
-            self._backend.save_datamanager(datamanager)
-
-        # = Create a searchspace
-        # Do this before One Hot Encoding to make sure that it creates a
-        # search space for a dense classifier even if one hot encoding would
-        # make it sparse (tradeoff; if one hot encoding would make it sparse,
-        #  densifier and truncatedSVD would probably lead to a MemoryError,
-        # like this we can't use some of the preprocessing methods in case
-        # the data became sparse)
-        with self._stopwatch.time("Create Search space"):
-            self.configuration_space, configspace_path = self._create_search_space(
-                self._backend.temporary_directory,
-                self._backend,
-                datamanager,
-                include=self._include,
-                exclude=self._exclude,
-            )
-
-        if only_return_configuration_space:
-            self._fit_cleanup()
-            return self.configuration_space
-
-        # == Perform dummy predictions
-        with self._stopwatch.time("Dummy predictions"):
-            self.num_run += 1
-            self._do_dummy_prediction()
-
-        # == RUN ensemble builder
-        # Do this before calculating the meta-features to make sure that the
-        # dummy predictions are actually included in the ensemble even if
-        # calculating the meta-features takes very long
-        with self._stopwatch.time("Run Ensemble Builder"):
-
-            elapsed_time = self._stopwatch.time_since(self._dataset_name, "start")
-
-            time_left_for_ensembles = max(0, self._time_for_task - elapsed_time)
-            proc_ensemble = None
-            if time_left_for_ensembles <= 0:
-                # Fit only raises error when an ensemble class is given but
-                # time_left_for_ensembles is zero.
-                if self._ensemble_class is not None:
-                    raise ValueError(
-                        "Not starting ensemble builder because there "
-                        "is no time left. Try increasing the value "
-                        "of time_left_for_this_task."
-                    )
-            elif self._ensemble_class is None:
-                self._logger.info(
-                    "Not starting ensemble builder because no ensemble class is given."
-                )
-            else:
-                self._logger.info(
-                    "Start Ensemble with %5.2fsec time left" % time_left_for_ensembles
-                )
-
-                proc_ensemble = EnsembleBuilderManager(
-                    start_time=time.time(),
-                    time_left_for_ensembles=time_left_for_ensembles,
-                    backend=copy.deepcopy(self._backend),
-                    dataset_name=dataset_name,
+            # == Pickle the data manager to speed up loading
+            with self._stopwatch.time("Save Datamanager"):
+                datamanager = XYDataManager(
+                    X,
+                    y,
+                    X_test=X_test,
+                    y_test=y_test,
                     task=self._task,
-                    metrics=self._metrics,
-                    ensemble_class=self._ensemble_class,
-                    ensemble_kwargs=self._ensemble_kwargs,
-                    ensemble_nbest=self._ensemble_nbest,
-                    max_models_on_disc=self._max_models_on_disc,
-                    seed=self._seed,
-                    precision=self.precision,
-                    max_iterations=self._max_ensemble_build_iterations,
-                    read_at_most=self._read_at_most,
-                    memory_limit=self._memory_limit,
-                    random_state=self._seed,
-                    logger_port=self._logger_port,
-                    pynisher_context=self._multiprocessing_context,
+                    feat_type=self._feat_type,
+                    dataset_name=dataset_name,
                 )
 
-        # kill the datamanager as it will be re-loaded anyways from sub processes
-        try:
-            del self._datamanager
-        except Exception:
-            pass
+                self._backend._make_internals_directory()
+                self._label_num = datamanager.info["label_num"]
+
+                self._backend.save_datamanager(datamanager)
+
+            # = Create a searchspace
+            # Do this before One Hot Encoding to make sure that it creates a
+            # search space for a dense classifier even if one hot encoding would
+            # make it sparse (tradeoff; if one hot encoding would make it sparse,
+            #  densifier and truncatedSVD would probably lead to a MemoryError,
+            # like this we can't use some of the preprocessing methods in case
+            # the data became sparse)
+            with self._stopwatch.time("Create Search space"):
+                self.configuration_space, configspace_path = self._create_search_space(
+                    self._backend.temporary_directory,
+                    self._backend,
+                    datamanager,
+                    include=self._include,
+                    exclude=self._exclude,
+                )
 
-        # => RUN SMAC
-        with self._stopwatch.time("Run SMAC"):
-            elapsed_time = self._stopwatch.time_since(self._dataset_name, "start")
-            time_left = self._time_for_task - elapsed_time
-
-            if self._logger:
-                self._logger.info("Start SMAC with %5.2fsec time left" % time_left)
-            if time_left <= 0:
-                self._logger.warning("Not starting SMAC because there is no time left.")
-                _proc_smac = None
-                self._budget_type = None
-            else:
-                if (
-                    self._per_run_time_limit is None
-                    or self._per_run_time_limit > time_left
-                ):
-                    self._logger.warning(
-                        "Time limit for a single run is higher than total time "
-                        "limit. Capping the limit for a single run to the total "
-                        "time given to SMAC (%f)" % time_left
+            if only_return_configuration_space:
+                return self.configuration_space
+
+            # == Perform dummy predictions
+            with self._stopwatch.time("Dummy predictions"):
+                self.num_run += 1
+                self._do_dummy_prediction()
+
+            # == RUN ensemble builder
+            # Do this before calculating the meta-features to make sure that the
+            # dummy predictions are actually included in the ensemble even if
+            # calculating the meta-features takes very long
+            with self._stopwatch.time("Run Ensemble Builder"):
+
+                elapsed_time = self._stopwatch.time_since(self._dataset_name, "start")
+
+                time_left_for_ensembles = max(0, self._time_for_task - elapsed_time)
+                proc_ensemble = None
+                if time_left_for_ensembles <= 0:
+                    # Fit only raises error when an ensemble class is given but
+                    # time_left_for_ensembles is zero.
+                    if self._ensemble_class is not None:
+                        raise ValueError(
+                            "Not starting ensemble builder because there "
+                            "is no time left. Try increasing the value "
+                            "of time_left_for_this_task."
+                        )
+                elif self._ensemble_class is None:
+                    self._logger.info(
+                        "No ensemble buildin because no ensemble class was given."
                     )
-                    per_run_time_limit = time_left
                 else:
-                    per_run_time_limit = self._per_run_time_limit
+                    self._logger.info(
+                        "Start Ensemble with %5.2fsec time left"
+                        % time_left_for_ensembles
+                    )
+
+                    proc_ensemble = EnsembleBuilderManager(
+                        start_time=time.time(),
+                        time_left_for_ensembles=time_left_for_ensembles,
+                        backend=copy.deepcopy(self._backend),
+                        dataset_name=dataset_name,
+                        task=self._task,
+                        metrics=self._metrics,
+                        ensemble_class=self._ensemble_class,
+                        ensemble_kwargs=self._ensemble_kwargs,
+                        ensemble_nbest=self._ensemble_nbest,
+                        max_models_on_disc=self._max_models_on_disc,
+                        seed=self._seed,
+                        precision=self.precision,
+                        max_iterations=self._max_ensemble_build_iterations,
+                        read_at_most=self._read_at_most,
+                        memory_limit=self._memory_limit,
+                        random_state=self._seed,
+                        logger_port=self._logger_port,
+                        pynisher_context=self._multiprocessing_context,
+                    )
 
-                # Make sure that at least 2 models are created for the ensemble process
-                num_models = time_left // per_run_time_limit
-                if num_models < 2:
-                    per_run_time_limit = time_left // 2
+            # kill the datamanager as it will be re-loaded anyways from sub processes
+            try:
+                del self._datamanager
+            except Exception:
+                pass
+
+            # => RUN SMAC
+            with self._stopwatch.time("Run SMAC"):
+                elapsed_time = self._stopwatch.time_since(self._dataset_name, "start")
+                time_left = self._time_for_task - elapsed_time
+
+                if self._logger:
+                    self._logger.info("Start SMAC with %5.2fsec time left" % time_left)
+                if time_left <= 0:
                     self._logger.warning(
-                        "Capping the per_run_time_limit to {} to have "
-                        "time for a least 2 models in each process.".format(
-                            per_run_time_limit
-                        )
+                        "Not starting SMAC because there is no time left."
                     )
+                    _proc_smac = None
+                    self._budget_type = None
+                else:
+                    if (
+                        self._per_run_time_limit is None
+                        or self._per_run_time_limit > time_left
+                    ):
+                        self._logger.warning(
+                            "Time limit for a single run is higher than total time "
+                            "limit. Capping the limit for a single run to the total "
+                            "time given to SMAC (%f)" % time_left
+                        )
+                        per_run_time_limit = time_left
+                    else:
+                        per_run_time_limit = self._per_run_time_limit
+
+                    # At least 2 models are created for the ensemble process
+                    num_models = time_left // per_run_time_limit
+                    if num_models < 2:
+                        per_run_time_limit = time_left // 2
+                        self._logger.warning(
+                            "Capping the per_run_time_limit to {} to have "
+                            "time for a least 2 models in each process.".format(
+                                per_run_time_limit
+                            )
+                        )
 
-                _proc_smac = AutoMLSMBO(
-                    config_space=self.configuration_space,
-                    dataset_name=self._dataset_name,
-                    backend=self._backend,
-                    total_walltime_limit=time_left,
-                    func_eval_time_limit=per_run_time_limit,
-                    memory_limit=self._memory_limit,
-                    data_memory_limit=self._data_memory_limit,
-                    stopwatch=self._stopwatch,
-                    n_jobs=self._n_jobs,
-                    dask_client=self._dask_client,
-                    start_num_run=self.num_run,
-                    num_metalearning_cfgs=self._initial_configurations_via_metalearning,
-                    config_file=configspace_path,
-                    seed=self._seed,
-                    metadata_directory=self._metadata_directory,
-                    metrics=self._metrics,
-                    resampling_strategy=self._resampling_strategy,
-                    resampling_strategy_args=self._resampling_strategy_arguments,
-                    include=self._include,
-                    exclude=self._exclude,
-                    disable_file_output=self._disable_evaluator_output,
-                    get_smac_object_callback=self._get_smac_object_callback,
-                    smac_scenario_args=self._smac_scenario_args,
-                    scoring_functions=self._scoring_functions,
-                    port=self._logger_port,
-                    pynisher_context=self._multiprocessing_context,
-                    ensemble_callback=proc_ensemble,
-                    trials_callback=self._get_trials_callback,
-                )
+                    n_meta_configs = self._initial_configurations_via_metalearning
+                    _proc_smac = AutoMLSMBO(
+                        config_space=self.configuration_space,
+                        dataset_name=self._dataset_name,
+                        backend=self._backend,
+                        total_walltime_limit=time_left,
+                        func_eval_time_limit=per_run_time_limit,
+                        memory_limit=self._memory_limit,
+                        data_memory_limit=self._data_memory_limit,
+                        stopwatch=self._stopwatch,
+                        n_jobs=self._n_jobs,
+                        dask_client=self._dask_client,
+                        start_num_run=self.num_run,
+                        num_metalearning_cfgs=n_meta_configs,
+                        config_file=configspace_path,
+                        seed=self._seed,
+                        metadata_directory=self._metadata_directory,
+                        metrics=self._metrics,
+                        resampling_strategy=self._resampling_strategy,
+                        resampling_strategy_args=self._resampling_strategy_arguments,
+                        include=self._include,
+                        exclude=self._exclude,
+                        disable_file_output=self._disable_evaluator_output,
+                        get_smac_object_callback=self._get_smac_object_callback,
+                        smac_scenario_args=self._smac_scenario_args,
+                        scoring_functions=self._scoring_functions,
+                        port=self._logger_port,
+                        pynisher_context=self._multiprocessing_context,
+                        ensemble_callback=proc_ensemble,
+                        trials_callback=self._get_trials_callback,
+                    )
 
-                try:
                     (
                         self.runhistory_,
                         self.trajectory_,
@@ -928,42 +937,49 @@ def fit(
                     ]
                     with open(trajectory_filename, "w") as fh:
                         json.dump(saveable_trajectory, fh)
-                except Exception as e:
-                    self._logger.exception(e)
-                    raise
-
-        self._logger.info("Starting shutdown...")
-        # Wait until the ensemble process is finished to avoid shutting down
-        # while the ensemble builder tries to access the data
-        if proc_ensemble is not None:
-            self.ensemble_performance_history = list(proc_ensemble.history)
-
-            if len(proc_ensemble.futures) > 0:
-                # Now we need to wait for the future to return as it cannot be cancelled
-                # while it is running: https://stackoverflow.com/a/49203129
-                self._logger.info(
-                    "Ensemble script still running, waiting for it to finish."
-                )
-                result = proc_ensemble.futures.pop().result()
-                if result:
-                    ensemble_history, _ = result
-                    self.ensemble_performance_history.extend(ensemble_history)
-                self._logger.info("Ensemble script finished, continue shutdown.")
-
-            # save the ensemble performance history file
-            if len(self.ensemble_performance_history) > 0:
-                pd.DataFrame(self.ensemble_performance_history).to_json(
-                    os.path.join(
-                        self._backend.internals_directory, "ensemble_history.json"
+
+            self._logger.info("Starting shutdown...")
+            # Wait until the ensemble process is finished to avoid shutting down
+            # while the ensemble builder tries to access the data
+            if proc_ensemble is not None:
+                self.ensemble_performance_history = list(proc_ensemble.history)
+
+                if len(proc_ensemble.futures) > 0:
+                    # Now we wait for the future to return as it cannot be cancelled
+                    # while it is running: https://stackoverflow.com/a/49203129
+                    self._logger.info(
+                        "Ensemble script still running, waiting for it to finish."
+                    )
+                    result = proc_ensemble.futures.pop().result()
+                    if result:
+                        ensemble_history, _ = result
+                        self.ensemble_performance_history.extend(ensemble_history)
+                    self._logger.info("Ensemble script finished, continue shutdown.")
+
+                # save the ensemble performance history file
+                if len(self.ensemble_performance_history) > 0:
+                    pd.DataFrame(self.ensemble_performance_history).to_json(
+                        os.path.join(
+                            self._backend.internals_directory, "ensemble_history.json"
+                        )
                     )
-                )
 
-        if load_models:
-            self._logger.info("Loading models...")
-            self._load_models()
-            self._logger.info("Finished loading models...")
+            if load_models:
+                self._logger.info("Loading models...")
+                self._load_models()
+                self._logger.info("Finished loading models...")
+
+        # The whole logic above from where we begin the logging server is capture
+        # in a try: finally: so that if something goes wrong, we at least close
+        # down the logging server, preventing it from hanging and not closing
+        # until ctrl+c is pressed
+        except Exception as e:
+            # This will be called before the _fit_cleanup
+            self._logger.exception(e)
+            raise e
+        finally:
+            self._fit_cleanup()
 
-        self._fit_cleanup()
         self.fitted = True
 
         return self


Test patch: