diff --git a/flaml/automl.py b/flaml/automl.py index 381a61a0a..66468237a 100644 --- a/flaml/automl.py +++ b/flaml/automl.py @@ -3069,7 +3069,9 @@ class AutoML(BaseEstimator): if mlflow is not None and mlflow.active_run(): with mlflow.start_run(nested=True): mlflow.log_metric("iter_counter", self._track_iter) - if "intermediate_results" in search_state.metric_for_logging: + if (search_state.metric_for_logging is not None) and ( + "intermediate_results" in search_state.metric_for_logging + ): for each_entry in search_state.metric_for_logging[ "intermediate_results" ]: @@ -3079,7 +3081,8 @@ class AutoML(BaseEstimator): "iter_counter", self._iter_per_learner[estimator] ) del search_state.metric_for_logging["intermediate_results"] - mlflow.log_metrics(search_state.metric_for_logging) + if search_state.metric_for_logging: + mlflow.log_metrics(search_state.metric_for_logging) mlflow.log_metric("trial_time", search_state.trial_time) mlflow.log_metric("wall_clock_time", self._state.time_from_start) mlflow.log_metric("validation_loss", search_state.val_loss) diff --git a/flaml/version.py b/flaml/version.py index 9eb1ebec5..bd538f76e 100644 --- a/flaml/version.py +++ b/flaml/version.py @@ -1 +1 @@ -__version__ = "1.0.11" +__version__ = "1.0.12" diff --git a/test/automl/test_notebook_example.py b/test/automl/test_notebook_example.py index 53b94d4d6..5f2382454 100644 --- a/test/automl/test_notebook_example.py +++ b/test/automl/test_notebook_example.py @@ -154,6 +154,19 @@ def test_mlflow(): pass # subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "mlflow"]) + from sklearn.datasets import load_iris + + with mlflow.start_run(): + automl = AutoML() + automl_settings = { + "time_budget": 2, # in seconds + "metric": "accuracy", + "task": "classification", + "log_file_name": "iris.log", + } + X_train, y_train = load_iris(return_X_y=True) + automl.fit(X_train=X_train, y_train=y_train, **automl_settings) + if __name__ == "__main__": test_automl(600)