Repository: alteryx/evalml (python)

Issue:
AutoML search: Add callback for error handling
**Background**
We [removed the old `raise_errors` flag](https://github.com/alteryx/evalml/pull/936/files) from `AutoMLSearch.search` in July. This flag stopped automl search on the first error.

**Use-case**
Users want to have control over what happens when automl search encounters an exception during pipeline fit/score.

#1047 tracks building an asynchronous API for automl search. But doing this issue would still add value before we get to #1047.

This would [help with perf testing too](https://github.com/FeatureLabs/evalml_looking_glass/issues/123).

**Proposal**

```python
# four different automl error callbacks we could define
def silent_error(automl, exception):
    """No-op"""

def log_error(automl, exception):
    """doesn't throw, but logs. default behavior"""
    logger.error(f'AutoML search encountered an exception: {str(e)}')
automl = AutoMLSearch(automl_error_callback=log_error)
automl.search(X, y)

def raise_error(automl, exception):
    """throws"""
    logger.error(f'AutoML search got fatal exception: {str(e)}')
    raise exception

def log_and_save_error(automl, exception):
    """doesn't throw, but adds exception to a list in results"""
    logger.warning(f'AutoML search encountered an exception: {str(e)}')
    automl._results['errors'] = automl._results.get('errors', [])
    automl._results['errors'].append(e)
automl = AutoMLSearch(automl_error_callback=raise_error)
automl.search(X, y) 

def raise_and_save_error(automl, exception):
    """save exception and throw"""
    logger.warning(f'AutoML search encountered an exception: {str(e)}')
    automl._results['errors'] = [e]
    raise exception

# usage of any of the above:
automl = AutoMLSearch(automl_error_callback=my_callback)
automl.search(X, y) # may now throw, or continue, depending.
automl.results['errors']
```

@jeremyliweishih fyi

Pull request description:
Add callback for error handling in AutoMLSearch
Closes #1217

Questions:
- Should we subclass [Protocol](https://docs.python.org/3/library/typing.html#typing.Protocol) to enforce that kwargs is passed? (https://stackoverflow.com/questions/64071488/callback-protocol-for-functions-with-different-kwargs)

If all objectives fail and `AutoMLException` is raised, we will raise since it indicates that something is likely wrong. This is the same behavior as what is currently implemented.

Patch:
diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst
index 870170bd..c8942ed8 100644
--- a/docs/source/api_reference.rst
+++ b/docs/source/api_reference.rst
@@ -73,6 +73,21 @@ AutoML Algorithm Classes
     IterativeAlgorithm
 
 
+.. currentmodule:: evalml.automl.callbacks
+
+AutoML Callbacks
+~~~~~~~~~~~~~~~~~~~~~~~~
+.. autosummary::
+    :toctree: generated
+    :nosignatures:
+
+    silent_error_callback
+    log_error_callback
+    raise_error_callback
+    log_and_save_error_callback
+    raise_and_save_error_callback
+
+
 .. currentmodule:: evalml.pipelines
 
 Pipelines
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 7e7b41d1..0558ec2a 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -3,7 +3,8 @@ Release Notes
 
 **Future Releases**
     * Enhancements
-        * Added ability to freeze hyperparameters for AutoMLSearch :pr:`1284`
+        * Added ability to freeze hyperparameters for ``AutoMLSearch`` :pr:`1284`
+        * Added callback for error handling in ``AutoMLSearch`` :pr:`1403`
         * Added the index id to the ``explain_predictions_best_worst`` output to help users identify which rows in their data are included :pr:`1365`
         * The top_k features displayed in ``explain_predictions_*`` functions are now determined by the magnitude of shap values as opposed to the ``top_k`` largest and smallest shap values. :pr:`1374`
         * Added a problem type for time series regression :pr:`1386`
diff --git a/evalml/automl/automl_search.py b/evalml/automl/automl_search.py
index 2b9c2cea..43f5140d 100644
--- a/evalml/automl/automl_search.py
+++ b/evalml/automl/automl_search.py
@@ -1,5 +1,7 @@
 import copy
+import sys
 import time
+import traceback
 from collections import OrderedDict, defaultdict
 
 import cloudpickle
@@ -16,6 +18,7 @@ from sklearn.model_selection import (
 from .pipeline_search_plots import PipelineSearchPlots
 
 from evalml.automl.automl_algorithm import IterativeAlgorithm
+from evalml.automl.callbacks import log_error_callback
 from evalml.automl.data_splitters import TrainingValidationSplit
 from evalml.automl.utils import get_default_primary_search_objective
 from evalml.data_checks import (
@@ -82,6 +85,7 @@ class AutoMLSearch:
                  allowed_model_families=None,
                  start_iteration_callback=None,
                  add_result_callback=None,
+                 error_callback=None,
                  additional_objectives=None,
                  random_state=0,
                  n_jobs=-1,
@@ -126,15 +130,20 @@ class AutoMLSearch:
                 to `multiclass` or `regression` depending on the problem type. Note that if allowed_pipelines is provided,
                 this parameter will be ignored.
 
-            data_split (sklearn.model_selection.BaseCrossValidator): data splitting method to use. Defaults to StratifiedKFold.
+            data_split (sklearn.model_selection.BaseCrossValidator): Data splitting method to use. Defaults to StratifiedKFold.
 
-            tuner_class: the tuner class to use. Defaults to scikit-optimize tuner
+            tuner_class: The tuner class to use. Defaults to SKOptTuner.
 
-            start_iteration_callback (callable): function called before each pipeline training iteration.
-                Passed three parameters: pipeline_class, parameters, and the AutoMLSearch object.
+            start_iteration_callback (callable): Function called before each pipeline training iteration.
+                Callback function takes three positional parameters: The pipeline class, the pipeline parameters, and the AutoMLSearch object.
 
-            add_result_callback (callable): function called after each pipeline training iteration.
-                Passed three parameters: A dictionary containing the training results for the new pipeline, an untrained_pipeline containing the parameters used during training, and the AutoMLSearch object.
+            add_result_callback (callable): Function called after each pipeline training iteration.
+                Callback function takes three positional parameters:: A dictionary containing the training results for the new pipeline, an untrained_pipeline containing the parameters used during training, and the AutoMLSearch object.
+
+            error_callback (callable): Function called when `search()` errors and raises an Exception.
+                Callback function takes three positional parameters: the Exception raised, the traceback, and the AutoMLSearch object.
+                Must also accepts kwargs, so AutoMLSearch is able to pass along other appropriate parameters by default.
+                Defaults to None, which will call `log_error_callback`.
 
             additional_objectives (list): Custom set of objectives to score on.
                 Will override default objectives for problem type if not empty.
@@ -162,6 +171,7 @@ class AutoMLSearch:
         self.tuner_class = tuner_class or SKOptTuner
         self.start_iteration_callback = start_iteration_callback
         self.add_result_callback = add_result_callback
+        self.error_callback = error_callback or log_error_callback
         self.data_split = data_split
         self.verbose = verbose
         self.optimize_thresholds = optimize_thresholds
@@ -212,7 +222,8 @@ class AutoMLSearch:
         self.tolerance = tolerance or 0.0
         self._results = {
             'pipeline_results': {},
-            'search_order': []
+            'search_order': [],
+            'errors': []
         }
         self.random_state = get_random_state(random_state)
         self.n_jobs = n_jobs
@@ -674,22 +685,15 @@ class AutoMLSearch:
                 logger.debug(f"\t\t\tFold {i}: {self.objective.name} score: {scores[self.objective.name]:.3f}")
                 score = scores[self.objective.name]
             except Exception as e:
+                if self.error_callback is not None:
+                    self.error_callback(exception=e, traceback=traceback.format_tb(sys.exc_info()[2]), automl=self,
+                                        fold_num=i, pipeline=pipeline)
                 if isinstance(e, PipelineScoreError):
-                    logger.info(f"\t\t\tFold {i}: Encountered an error scoring the following objectives: {', '.join(e.exceptions)}.")
-                    logger.info(f"\t\t\tFold {i}: The scores for these objectives will be replaced with nan.")
-                    logger.info(f"\t\t\tFold {i}: Please check {logger.handlers[1].baseFilename} for the current hyperparameters and stack trace.")
-                    logger.debug(f"\t\t\tFold {i}: Hyperparameters:\n\t{pipeline.hyperparameters}")
-                    logger.debug(f"\t\t\tFold {i}: Exception during automl search: {str(e)}")
                     nan_scores = {objective: np.nan for objective in e.exceptions}
                     scores = {**nan_scores, **e.scored_successfully}
                     scores = OrderedDict({o.name: scores[o.name] for o in [self.objective] + self.additional_objectives})
                     score = scores[self.objective.name]
                 else:
-                    logger.info(f"\t\t\tFold {i}: Encountered an error.")
-                    logger.info(f"\t\t\tFold {i}: All scores will be replaced with nan.")
-                    logger.info(f"\t\t\tFold {i}: Please check {logger.handlers[1].baseFilename} for the current hyperparameters and stack trace.")
-                    logger.debug(f"\t\t\tFold {i}: Hyperparameters:\n\t{pipeline.hyperparameters}")
-                    logger.debug(f"\t\t\tFold {i}: Exception during automl search: {str(e)}")
                     score = np.nan
                     scores = OrderedDict(zip([n.name for n in self.additional_objectives], [np.nan] * len(self.additional_objectives)))
 
diff --git a/evalml/automl/callbacks.py b/evalml/automl/callbacks.py
new file mode 100644
index 00000000..414fb0c6
--- /dev/null
+++ b/evalml/automl/callbacks.py
@@ -0,0 +1,47 @@
+from evalml.exceptions import PipelineScoreError
+from evalml.utils.logger import get_logger
+
+logger = get_logger(__file__)
+
+
+def silent_error_callback(exception, traceback, automl, **kwargs):
+    """No-op."""
+
+
+def raise_error_callback(exception, traceback, automl, **kwargs):
+    """Raises the exception thrown by the AutoMLSearch object. Also logs the exception as an error."""
+    logger.error(f'AutoMLSearch raised a fatal exception: {str(exception)}')
+    logger.error("\n".join(traceback))
+    raise exception
+
+
+def log_and_save_error_callback(exception, traceback, automl, **kwargs):
+    """Logs the exception thrown by the AutoMLSearch object as a warning
+        and adds the exception to the 'errors' list in AutoMLSearch object results."""
+    logger.warning(f'AutoML search encountered an exception: {str(exception)}')
+    logger.warning("\n".join(traceback))
+    automl._results['errors'].append(exception)
+
+
+def raise_and_save_error_callback(exception, traceback, automl, **kwargs):
+    """Raises the exception thrown by the AutoMLSearch object, logs it as an error,
+        and adds the exception to the 'errors' list in AutoMLSearch object results."""
+    logger.error(f'AutoMLSearch raised a fatal exception: {str(exception)}')
+    logger.error("\n".join(traceback))
+    automl._results['errors'].append(exception)
+    raise exception
+
+
+def log_error_callback(exception, traceback, automl, **kwargs):
+    """Logs the exception thrown as an error. Will not throw. This is the default behavior for AutoMLSearch."""
+    fold_num = kwargs.get('fold_num')
+    pipeline = kwargs.get('pipeline')
+    if isinstance(exception, PipelineScoreError):
+        logger.info(f"\t\t\tFold {fold_num}: Encountered an error scoring the following objectives: {', '.join(exception.exceptions)}.")
+        logger.info(f"\t\t\tFold {fold_num}: The scores for these objectives will be replaced with nan.")
+    else:
+        logger.info(f"\t\t\tFold {fold_num}: Encountered an error.")
+        logger.info(f"\t\t\tFold {fold_num}: All scores will be replaced with nan.")
+    logger.info(f"\t\t\tFold {fold_num}: Please check {logger.handlers[1].baseFilename} for the current hyperparameters and stack trace.")
+    logger.debug(f"\t\t\tFold {fold_num}: Hyperparameters:\n\t{pipeline.hyperparameters}")
+    logger.info(f"\t\t\tFold {fold_num}: Exception during automl search: {str(exception)}")


Test patch:
diff --git a/evalml/tests/automl_tests/test_automl.py b/evalml/tests/automl_tests/test_automl.py
index bf80da47..a3f387cb 100644
--- a/evalml/tests/automl_tests/test_automl.py
+++ b/evalml/tests/automl_tests/test_automl.py
@@ -14,6 +14,13 @@ from evalml.automl import (
     TrainingValidationSplit,
     get_default_primary_search_objective
 )
+from evalml.automl.callbacks import (
+    log_and_save_error_callback,
+    log_error_callback,
+    raise_and_save_error_callback,
+    raise_error_callback,
+    silent_error_callback
+)
 from evalml.data_checks import (
     DataCheck,
     DataCheckError,
@@ -63,7 +70,7 @@ def test_search_results(X_y_regression, X_y_binary, X_y_multi, automl_type):
         X, y = X_y_multi
 
     automl.search(X, y)
-    assert automl.results.keys() == {'pipeline_results', 'search_order'}
+    assert automl.results.keys() == {'pipeline_results', 'search_order', 'errors'}
     assert automl.results['search_order'] == [0, 1]
     assert len(automl.results['pipeline_results']) == 2
     for pipeline_id, results in automl.results['pipeline_results'].items():
@@ -935,7 +942,9 @@ def test_results_getter(mock_fit, mock_score, caplog, X_y_binary):
     X, y = X_y_binary
     automl = AutoMLSearch(problem_type='binary', max_iterations=1)
 
-    assert automl.results == {'pipeline_results': {}, 'search_order': []}
+    assert automl.results == {'pipeline_results': {},
+                              'search_order': [],
+                              'errors': []}
 
     mock_score.return_value = {'Log Loss Binary': 1.0}
     automl.search(X, y)
@@ -1625,3 +1634,50 @@ def test_automl_respects_random_state(mock_fit, mock_score, X_y_binary, dummy_cl
                           random_state=expected_random_state, max_iterations=10)
     automl.search(X, y)
     assert DummyPipeline.num_pipelines_different_seed == 0 and DummyPipeline.num_pipelines_init
+
+
+@patch('evalml.pipelines.BinaryClassificationPipeline.score', return_value={"Log Loss Binary": 0.8})
+@patch('evalml.pipelines.BinaryClassificationPipeline.fit')
+def test_automl_error_callback(mock_fit, mock_score, X_y_binary, caplog):
+    X, y = X_y_binary
+    msg = 'all your model are belong to us'
+    mock_fit.side_effect = Exception(msg)
+    automl = AutoMLSearch(problem_type="binary", error_callback=None)
+    automl.search(X, y)
+    assert msg in caplog.text
+
+    caplog.clear()
+    automl = AutoMLSearch(problem_type="binary", error_callback=silent_error_callback)
+    automl.search(X, y)
+    assert msg not in caplog.text
+
+    caplog.clear()
+    automl = AutoMLSearch(problem_type="binary", error_callback=log_error_callback)
+    automl.search(X, y)
+    assert msg in caplog.text
+
+    caplog.clear()
+    automl = AutoMLSearch(problem_type="binary", error_callback=raise_error_callback)
+    with pytest.raises(Exception, match="all your model are belong to us"):
+        automl.search(X, y)
+    assert "AutoMLSearch raised a fatal exception: all your model are belong to us" in caplog.text
+    assert "fit" in caplog.text  # Check stack trace logged
+
+    caplog.clear()
+    automl = AutoMLSearch(problem_type="binary", error_callback=log_and_save_error_callback)
+    automl.search(X, y)
+    assert "AutoML search encountered an exception: all your model are belong to us" in caplog.text
+    assert "fit" in caplog.text  # Check stack trace logged
+    assert len(automl._results['errors']) == 15  # 5 iterations, 3 folds each
+    for e in automl._results['errors']:
+        assert str(e) == msg
+
+    caplog.clear()
+    automl = AutoMLSearch(problem_type="binary", error_callback=raise_and_save_error_callback)
+    with pytest.raises(Exception, match="all your model are belong to us"):
+        automl.search(X, y)
+    assert "AutoMLSearch raised a fatal exception: all your model are belong to us" in caplog.text
+    assert "fit" in caplog.text  # Check stack trace logged
+    assert len(automl._results['errors']) == 1  # Raises exception at first error
+    for e in automl._results['errors']:
+        assert str(e) == msg
diff --git a/evalml/tests/model_understanding_tests/prediction_explanations_tests/test_user_interface.py b/evalml/tests/model_understanding_tests/prediction_explanations_tests/test_user_interface.py
index 6d8e5611..3b18f7b5 100644
--- a/evalml/tests/model_understanding_tests/prediction_explanations_tests/test_user_interface.py
+++ b/evalml/tests/model_understanding_tests/prediction_explanations_tests/test_user_interface.py
@@ -127,7 +127,7 @@ binary = [{"a": [0], "b": [0], "c": [0],
 
 binary_normalized = [{'a': [0.0], 'b': [0.0], 'c': [0.0], 'd': [0.0], 'e': [0.0], 'f': [0.0], 'foo': [-1.0]},
                      {'a': [0.16], 'b': [0.0], 'c': [0.15],
-                     'd': [-0.08], 'e': [-0.35], 'f': [-0.12], 'foo': [-0.14]}]
+                      'd': [-0.08], 'e': [-0.35], 'f': [-0.12], 'foo': [-0.14]}]
 binary_pipeline_features = pd.DataFrame({"a": 2.18, "b": 2.12, "c": 1.0, "d": -1.56, "e": -1.8, "f": -1.9,
                                          "foo": -20}, index=[23])