FLAML 2.3.4__tar.gz → 2.3.5__tar.gz

{flaml-2.3.4 → FLAML-2.3.5}/FLAML.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: FLAML
-Version: 2.3.4
+Version: 2.3.5
 Summary: A fast library for automated machine learning and tuning
 Home-page: https://github.com/microsoft/FLAML
 Author: Microsoft Corporation
@@ -8,144 +8,34 @@ Author-email: hpo@microsoft.com
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-License-File: LICENSE
-License-File: NOTICE.md
-Requires-Dist: NumPy>=1.17
 Provides-Extra: automl
-Requires-Dist: lightgbm>=2.3.1; extra == "automl"
-Requires-Dist: xgboost<3.0.0,>=0.90; extra == "automl"
-Requires-Dist: scipy>=1.4.1; extra == "automl"
-Requires-Dist: pandas>=1.1.4; extra == "automl"
-Requires-Dist: scikit-learn>=1.0.0; extra == "automl"
 Provides-Extra: notebook
-Requires-Dist: jupyter; extra == "notebook"
 Provides-Extra: spark
-Requires-Dist: pyspark>=3.2.0; extra == "spark"
-Requires-Dist: joblibspark>=0.5.0; extra == "spark"
-Requires-Dist: joblib<=1.3.2; extra == "spark"
 Provides-Extra: test
-Requires-Dist: jupyter; extra == "test"
-Requires-Dist: lightgbm>=2.3.1; extra == "test"
-Requires-Dist: xgboost<2.0.0,>=0.90; extra == "test"
-Requires-Dist: scipy>=1.4.1; extra == "test"
-Requires-Dist: pandas<2.0.0,>=1.1.4; python_version < "3.10" and extra == "test"
-Requires-Dist: pandas>=1.1.4; python_version >= "3.10" and extra == "test"
-Requires-Dist: scikit-learn>=1.0.0; extra == "test"
-Requires-Dist: thop; extra == "test"
-Requires-Dist: pytest>=6.1.1; extra == "test"
-Requires-Dist: coverage>=5.3; extra == "test"
-Requires-Dist: pre-commit; extra == "test"
-Requires-Dist: torch; extra == "test"
-Requires-Dist: torchvision; extra == "test"
-Requires-Dist: catboost<1.2,>=0.26; python_version < "3.11" and extra == "test"
-Requires-Dist: catboost>=0.26; python_version >= "3.11" and extra == "test"
-Requires-Dist: rgf-python; extra == "test"
-Requires-Dist: optuna<=3.6.1,>=2.8.0; extra == "test"
-Requires-Dist: openml; extra == "test"
-Requires-Dist: statsmodels>=0.12.2; extra == "test"
-Requires-Dist: psutil==5.8.0; extra == "test"
-Requires-Dist: dataclasses; extra == "test"
-Requires-Dist: transformers[torch]==4.26; extra == "test"
-Requires-Dist: datasets; extra == "test"
-Requires-Dist: nltk<=3.8.1; extra == "test"
-Requires-Dist: rouge_score; extra == "test"
-Requires-Dist: hcrystalball==0.1.10; extra == "test"
-Requires-Dist: seqeval; extra == "test"
-Requires-Dist: pytorch-forecasting<=0.10.1,>=0.9.0; python_version < "3.11" and extra == "test"
-Requires-Dist: mlflow==2.15.1; extra == "test"
-Requires-Dist: joblibspark>=0.5.0; extra == "test"
-Requires-Dist: joblib<=1.3.2; extra == "test"
-Requires-Dist: nbconvert; extra == "test"
-Requires-Dist: nbformat; extra == "test"
-Requires-Dist: ipykernel; extra == "test"
-Requires-Dist: pytorch-lightning<1.9.1; extra == "test"
-Requires-Dist: tensorboardX==2.6; extra == "test"
-Requires-Dist: requests<2.29.0; extra == "test"
-Requires-Dist: packaging; extra == "test"
-Requires-Dist: pydantic==1.10.9; extra == "test"
-Requires-Dist: sympy; extra == "test"
-Requires-Dist: wolframalpha; extra == "test"
-Requires-Dist: dill; extra == "test"
 Provides-Extra: catboost
-Requires-Dist: catboost<1.2,>=0.26; python_version < "3.11" and extra == "catboost"
-Requires-Dist: catboost<=1.2.5,>=0.26; python_version >= "3.11" and extra == "catboost"
 Provides-Extra: blendsearch
-Requires-Dist: optuna<=3.6.1,>=2.8.0; extra == "blendsearch"
-Requires-Dist: packaging; extra == "blendsearch"
 Provides-Extra: ray
-Requires-Dist: ray[tune]~=1.13; extra == "ray"
 Provides-Extra: azureml
-Requires-Dist: azureml-mlflow; extra == "azureml"
 Provides-Extra: nni
-Requires-Dist: nni; extra == "nni"
 Provides-Extra: vw
-Requires-Dist: vowpalwabbit<9.0.0,>=8.10.0; extra == "vw"
-Requires-Dist: scikit-learn; extra == "vw"
 Provides-Extra: hf
-Requires-Dist: transformers[torch]==4.26; extra == "hf"
-Requires-Dist: datasets; extra == "hf"
-Requires-Dist: nltk<=3.8.1; extra == "hf"
-Requires-Dist: rouge_score; extra == "hf"
-Requires-Dist: seqeval; extra == "hf"
 Provides-Extra: nlp
-Requires-Dist: transformers[torch]==4.26; extra == "nlp"
-Requires-Dist: datasets; extra == "nlp"
-Requires-Dist: nltk<=3.8.1; extra == "nlp"
-Requires-Dist: rouge_score; extra == "nlp"
-Requires-Dist: seqeval; extra == "nlp"
-Provides-Extra: ts-forecast
-Requires-Dist: holidays<0.14; extra == "ts-forecast"
-Requires-Dist: prophet>=1.0.1; extra == "ts-forecast"
-Requires-Dist: statsmodels>=0.12.2; extra == "ts-forecast"
-Requires-Dist: hcrystalball==0.1.10; extra == "ts-forecast"
+Provides-Extra: ts_forecast
 Provides-Extra: forecast
-Requires-Dist: holidays<0.14; extra == "forecast"
-Requires-Dist: prophet>=1.0.1; extra == "forecast"
-Requires-Dist: statsmodels>=0.12.2; extra == "forecast"
-Requires-Dist: hcrystalball==0.1.10; extra == "forecast"
-Requires-Dist: pytorch-forecasting>=0.9.0; python_version < "3.11" and extra == "forecast"
-Requires-Dist: pytorch-lightning==1.9.0; extra == "forecast"
-Requires-Dist: tensorboardX==2.6; extra == "forecast"
 Provides-Extra: benchmark
-Requires-Dist: catboost>=0.26; extra == "benchmark"
-Requires-Dist: psutil==5.8.0; extra == "benchmark"
-Requires-Dist: xgboost==1.3.3; extra == "benchmark"
-Requires-Dist: pandas==1.1.4; extra == "benchmark"
 Provides-Extra: openai
-Requires-Dist: openai==0.27.8; extra == "openai"
-Requires-Dist: diskcache; extra == "openai"
 Provides-Extra: autogen
-Requires-Dist: openai==0.27.8; extra == "autogen"
-Requires-Dist: diskcache; extra == "autogen"
-Requires-Dist: termcolor; extra == "autogen"
 Provides-Extra: mathchat
-Requires-Dist: openai==0.27.8; extra == "mathchat"
-Requires-Dist: diskcache; extra == "mathchat"
-Requires-Dist: termcolor; extra == "mathchat"
-Requires-Dist: sympy; extra == "mathchat"
-Requires-Dist: pydantic==1.10.9; extra == "mathchat"
-Requires-Dist: wolframalpha; extra == "mathchat"
 Provides-Extra: retrievechat
-Requires-Dist: openai==0.27.8; extra == "retrievechat"
-Requires-Dist: diskcache; extra == "retrievechat"
-Requires-Dist: termcolor; extra == "retrievechat"
-Requires-Dist: chromadb; extra == "retrievechat"
-Requires-Dist: tiktoken; extra == "retrievechat"
-Requires-Dist: sentence_transformers; extra == "retrievechat"
 Provides-Extra: synapse
-Requires-Dist: joblibspark>=0.5.0; extra == "synapse"
-Requires-Dist: optuna<=3.6.1,>=2.8.0; extra == "synapse"
-Requires-Dist: pyspark>=3.2.0; extra == "synapse"
 Provides-Extra: autozero
-Requires-Dist: scikit-learn; extra == "autozero"
-Requires-Dist: pandas; extra == "autozero"
-Requires-Dist: packaging; extra == "autozero"
+License-File: LICENSE
+License-File: NOTICE.md
 
 [![PyPI version](https://badge.fury.io/py/FLAML.svg)](https://badge.fury.io/py/FLAML)
 ![Conda version](https://img.shields.io/conda/vn/conda-forge/flaml)
@@ -189,7 +79,7 @@ FLAML has a .NET implementation in [ML.NET](http://dot.net/ml), an open-source,
 
 ## Installation
 
-FLAML requires **Python version >= 3.8**. It can be installed from pip:
+FLAML requires **Python version >= 3.9**. It can be installed from pip:
 
 ```bash
 pip install flaml

{flaml-2.3.4 → FLAML-2.3.5}/FLAML.egg-info/SOURCES.txt

@@ -90,6 +90,7 @@ flaml/onlineml/trial.py
 flaml/onlineml/trial_runner.py
 flaml/tune/__init__.py
 flaml/tune/analysis.py
+flaml/tune/logger.py
 flaml/tune/result.py
 flaml/tune/sample.py
 flaml/tune/space.py

{flaml-2.3.4 → FLAML-2.3.5}/FLAML.egg-info/requires.txt

@@ -121,7 +121,7 @@ statsmodels>=0.12.2
 psutil==5.8.0
 dataclasses
 transformers[torch]==4.26
-datasets
+datasets<=3.5.0
 nltk<=3.8.1
 rouge_score
 hcrystalball==0.1.10

{flaml-2.3.4 → FLAML-2.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: FLAML
-Version: 2.3.4
+Version: 2.3.5
 Summary: A fast library for automated machine learning and tuning
 Home-page: https://github.com/microsoft/FLAML
 Author: Microsoft Corporation
@@ -8,144 +8,34 @@ Author-email: hpo@microsoft.com
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-License-File: LICENSE
-License-File: NOTICE.md
-Requires-Dist: NumPy>=1.17
 Provides-Extra: automl
-Requires-Dist: lightgbm>=2.3.1; extra == "automl"
-Requires-Dist: xgboost<3.0.0,>=0.90; extra == "automl"
-Requires-Dist: scipy>=1.4.1; extra == "automl"
-Requires-Dist: pandas>=1.1.4; extra == "automl"
-Requires-Dist: scikit-learn>=1.0.0; extra == "automl"
 Provides-Extra: notebook
-Requires-Dist: jupyter; extra == "notebook"
 Provides-Extra: spark
-Requires-Dist: pyspark>=3.2.0; extra == "spark"
-Requires-Dist: joblibspark>=0.5.0; extra == "spark"
-Requires-Dist: joblib<=1.3.2; extra == "spark"
 Provides-Extra: test
-Requires-Dist: jupyter; extra == "test"
-Requires-Dist: lightgbm>=2.3.1; extra == "test"
-Requires-Dist: xgboost<2.0.0,>=0.90; extra == "test"
-Requires-Dist: scipy>=1.4.1; extra == "test"
-Requires-Dist: pandas<2.0.0,>=1.1.4; python_version < "3.10" and extra == "test"
-Requires-Dist: pandas>=1.1.4; python_version >= "3.10" and extra == "test"
-Requires-Dist: scikit-learn>=1.0.0; extra == "test"
-Requires-Dist: thop; extra == "test"
-Requires-Dist: pytest>=6.1.1; extra == "test"
-Requires-Dist: coverage>=5.3; extra == "test"
-Requires-Dist: pre-commit; extra == "test"
-Requires-Dist: torch; extra == "test"
-Requires-Dist: torchvision; extra == "test"
-Requires-Dist: catboost<1.2,>=0.26; python_version < "3.11" and extra == "test"
-Requires-Dist: catboost>=0.26; python_version >= "3.11" and extra == "test"
-Requires-Dist: rgf-python; extra == "test"
-Requires-Dist: optuna<=3.6.1,>=2.8.0; extra == "test"
-Requires-Dist: openml; extra == "test"
-Requires-Dist: statsmodels>=0.12.2; extra == "test"
-Requires-Dist: psutil==5.8.0; extra == "test"
-Requires-Dist: dataclasses; extra == "test"
-Requires-Dist: transformers[torch]==4.26; extra == "test"
-Requires-Dist: datasets; extra == "test"
-Requires-Dist: nltk<=3.8.1; extra == "test"
-Requires-Dist: rouge_score; extra == "test"
-Requires-Dist: hcrystalball==0.1.10; extra == "test"
-Requires-Dist: seqeval; extra == "test"
-Requires-Dist: pytorch-forecasting<=0.10.1,>=0.9.0; python_version < "3.11" and extra == "test"
-Requires-Dist: mlflow==2.15.1; extra == "test"
-Requires-Dist: joblibspark>=0.5.0; extra == "test"
-Requires-Dist: joblib<=1.3.2; extra == "test"
-Requires-Dist: nbconvert; extra == "test"
-Requires-Dist: nbformat; extra == "test"
-Requires-Dist: ipykernel; extra == "test"
-Requires-Dist: pytorch-lightning<1.9.1; extra == "test"
-Requires-Dist: tensorboardX==2.6; extra == "test"
-Requires-Dist: requests<2.29.0; extra == "test"
-Requires-Dist: packaging; extra == "test"
-Requires-Dist: pydantic==1.10.9; extra == "test"
-Requires-Dist: sympy; extra == "test"
-Requires-Dist: wolframalpha; extra == "test"
-Requires-Dist: dill; extra == "test"
 Provides-Extra: catboost
-Requires-Dist: catboost<1.2,>=0.26; python_version < "3.11" and extra == "catboost"
-Requires-Dist: catboost<=1.2.5,>=0.26; python_version >= "3.11" and extra == "catboost"
 Provides-Extra: blendsearch
-Requires-Dist: optuna<=3.6.1,>=2.8.0; extra == "blendsearch"
-Requires-Dist: packaging; extra == "blendsearch"
 Provides-Extra: ray
-Requires-Dist: ray[tune]~=1.13; extra == "ray"
 Provides-Extra: azureml
-Requires-Dist: azureml-mlflow; extra == "azureml"
 Provides-Extra: nni
-Requires-Dist: nni; extra == "nni"
 Provides-Extra: vw
-Requires-Dist: vowpalwabbit<9.0.0,>=8.10.0; extra == "vw"
-Requires-Dist: scikit-learn; extra == "vw"
 Provides-Extra: hf
-Requires-Dist: transformers[torch]==4.26; extra == "hf"
-Requires-Dist: datasets; extra == "hf"
-Requires-Dist: nltk<=3.8.1; extra == "hf"
-Requires-Dist: rouge_score; extra == "hf"
-Requires-Dist: seqeval; extra == "hf"
 Provides-Extra: nlp
-Requires-Dist: transformers[torch]==4.26; extra == "nlp"
-Requires-Dist: datasets; extra == "nlp"
-Requires-Dist: nltk<=3.8.1; extra == "nlp"
-Requires-Dist: rouge_score; extra == "nlp"
-Requires-Dist: seqeval; extra == "nlp"
-Provides-Extra: ts-forecast
-Requires-Dist: holidays<0.14; extra == "ts-forecast"
-Requires-Dist: prophet>=1.0.1; extra == "ts-forecast"
-Requires-Dist: statsmodels>=0.12.2; extra == "ts-forecast"
-Requires-Dist: hcrystalball==0.1.10; extra == "ts-forecast"
+Provides-Extra: ts_forecast
 Provides-Extra: forecast
-Requires-Dist: holidays<0.14; extra == "forecast"
-Requires-Dist: prophet>=1.0.1; extra == "forecast"
-Requires-Dist: statsmodels>=0.12.2; extra == "forecast"
-Requires-Dist: hcrystalball==0.1.10; extra == "forecast"
-Requires-Dist: pytorch-forecasting>=0.9.0; python_version < "3.11" and extra == "forecast"
-Requires-Dist: pytorch-lightning==1.9.0; extra == "forecast"
-Requires-Dist: tensorboardX==2.6; extra == "forecast"
 Provides-Extra: benchmark
-Requires-Dist: catboost>=0.26; extra == "benchmark"
-Requires-Dist: psutil==5.8.0; extra == "benchmark"
-Requires-Dist: xgboost==1.3.3; extra == "benchmark"
-Requires-Dist: pandas==1.1.4; extra == "benchmark"
 Provides-Extra: openai
-Requires-Dist: openai==0.27.8; extra == "openai"
-Requires-Dist: diskcache; extra == "openai"
 Provides-Extra: autogen
-Requires-Dist: openai==0.27.8; extra == "autogen"
-Requires-Dist: diskcache; extra == "autogen"
-Requires-Dist: termcolor; extra == "autogen"
 Provides-Extra: mathchat
-Requires-Dist: openai==0.27.8; extra == "mathchat"
-Requires-Dist: diskcache; extra == "mathchat"
-Requires-Dist: termcolor; extra == "mathchat"
-Requires-Dist: sympy; extra == "mathchat"
-Requires-Dist: pydantic==1.10.9; extra == "mathchat"
-Requires-Dist: wolframalpha; extra == "mathchat"
 Provides-Extra: retrievechat
-Requires-Dist: openai==0.27.8; extra == "retrievechat"
-Requires-Dist: diskcache; extra == "retrievechat"
-Requires-Dist: termcolor; extra == "retrievechat"
-Requires-Dist: chromadb; extra == "retrievechat"
-Requires-Dist: tiktoken; extra == "retrievechat"
-Requires-Dist: sentence_transformers; extra == "retrievechat"
 Provides-Extra: synapse
-Requires-Dist: joblibspark>=0.5.0; extra == "synapse"
-Requires-Dist: optuna<=3.6.1,>=2.8.0; extra == "synapse"
-Requires-Dist: pyspark>=3.2.0; extra == "synapse"
 Provides-Extra: autozero
-Requires-Dist: scikit-learn; extra == "autozero"
-Requires-Dist: pandas; extra == "autozero"
-Requires-Dist: packaging; extra == "autozero"
+License-File: LICENSE
+License-File: NOTICE.md
 
 [![PyPI version](https://badge.fury.io/py/FLAML.svg)](https://badge.fury.io/py/FLAML)
 ![Conda version](https://img.shields.io/conda/vn/conda-forge/flaml)
@@ -189,7 +79,7 @@ FLAML has a .NET implementation in [ML.NET](http://dot.net/ml), an open-source,
 
 ## Installation
 
-FLAML requires **Python version >= 3.8**. It can be installed from pip:
+FLAML requires **Python version >= 3.9**. It can be installed from pip:
 
 ```bash
 pip install flaml

{flaml-2.3.4 → FLAML-2.3.5}/README.md

@@ -40,7 +40,7 @@ FLAML has a .NET implementation in [ML.NET](http://dot.net/ml), an open-source,
 
 ## Installation
 
-FLAML requires **Python version >= 3.8**. It can be installed from pip:
+FLAML requires **Python version >= 3.9**. It can be installed from pip:
 
 ```bash
 pip install flaml

{flaml-2.3.4 → FLAML-2.3.5}/flaml/automl/automl.py

@@ -10,6 +10,7 @@ import os
 import random
 import sys
 import time
+from concurrent.futures import as_completed
 from functools import partial
 from typing import Callable, List, Optional, Union
 
@@ -187,9 +188,16 @@ class AutoML(BaseEstimator):
             mem_thres: A float of the memory size constraint in bytes.
             pred_time_limit: A float of the prediction latency constraint in seconds.
                 It refers to the average prediction time per row in validation data.
-            train_time_limit: A float of the training time constraint in seconds.
+            train_time_limit: None or a float of the training time constraint in seconds for each trial.
+                Only valid for sequential search.
             verbose: int, default=3 | Controls the verbosity, higher means more
                 messages.
+                verbose=0: logger level = CRITICAL
+                verbose=1: logger level = ERROR
+                verbose=2: logger level = WARNING
+                verbose=3: logger level = INFO
+                verbose=4: logger level = DEBUG
+                verbose>5: logger level = NOTSET
             retrain_full: bool or str, default=True | whether to retrain the
                 selected model on the full training data when using holdout.
                 True - retrain only after search finishes; False - no retraining;
@@ -424,6 +432,8 @@ class AutoML(BaseEstimator):
             If `model_history` was set to True, then the returned model is trained.
         """
         state = self._search_states.get(estimator_name)
+        if state and estimator_name == self._best_estimator:
+            return self.model
         return state and getattr(state, "trained_estimator", None)
 
     @property
@@ -1332,7 +1342,8 @@ class AutoML(BaseEstimator):
             mem_thres: A float of the memory size constraint in bytes.
             pred_time_limit: A float of the prediction latency constraint in seconds.
                 It refers to the average prediction time per row in validation data.
-            train_time_limit: None or a float of the training time constraint in seconds.
+            train_time_limit: None or a float of the training time constraint in seconds for each trial.
+                Only valid for sequential search.
             X_val: None or a numpy array or a pandas dataframe of validation data.
             y_val: None or a numpy array or a pandas series of validation labels.
             sample_weight_val: None or a numpy array of the sample weight of
@@ -1345,6 +1356,12 @@ class AutoML(BaseEstimator):
                 for training data.
             verbose: int, default=3 | Controls the verbosity, higher means more
                 messages.
+                verbose=0: logger level = CRITICAL
+                verbose=1: logger level = ERROR
+                verbose=2: logger level = WARNING
+                verbose=3: logger level = INFO
+                verbose=4: logger level = DEBUG
+                verbose>5: logger level = NOTSET
             retrain_full: bool or str, default=True | whether to retrain the
                 selected model on the full training data when using holdout.
                 True - retrain only after search finishes; False - no retraining;
@@ -1623,6 +1640,13 @@ class AutoML(BaseEstimator):
             _ch.setFormatter(logger_formatter)
             logger.addHandler(_ch)
 
+        if model_history:
+            logger.warning(
+                "With `model_history` set to `True` by default, all intermediate models are retained in memory, "
+                "which may significantly increase memory usage and slow down training. "
+                "Consider setting `model_history=False` to optimize memory and accelerate the training process."
+            )
+
         if not use_ray and not use_spark and n_concurrent_trials > 1:
             if ray_available:
                 logger.warning(
@@ -1708,7 +1732,7 @@ class AutoML(BaseEstimator):
                 if not (mlflow.active_run() is not None or is_autolog_enabled()):
                     self.mlflow_integration.only_history = True
             except KeyError:
-                print("Not in Fabric, Skipped")
+                logger.info("Not in Fabric, Skipped")
         task.validate_data(
             self,
             self._state,
@@ -2529,6 +2553,21 @@ class AutoML(BaseEstimator):
             self._selected = state = self._search_states[estimator]
             state.best_config_sample_size = self._state.data_size[0]
             state.best_config = state.init_config[0] if state.init_config else {}
+            self._track_iter = 0
+            self._config_history[self._track_iter] = (estimator, state.best_config, self._state.time_from_start)
+            self._best_iteration = self._track_iter
+            state.val_loss = getattr(state, "val_loss", float("inf"))
+            state.best_loss = getattr(state, "best_loss", float("inf"))
+            state.config = getattr(state, "config", state.best_config.copy())
+            state.metric_for_logging = getattr(state, "metric_for_logging", None)
+            state.sample_size = getattr(state, "sample_size", self._state.data_size[0])
+            state.learner_class = getattr(state, "learner_class", self._state.learner_classes.get(estimator))
+            if hasattr(self, "mlflow_integration") and self.mlflow_integration:
+                self.mlflow_integration.record_state(
+                    automl=self,
+                    search_state=state,
+                    estimator=estimator,
+                )
         elif self._use_ray is False and self._use_spark is False:
             self._search_sequential()
         else:
@@ -2700,16 +2739,47 @@ class AutoML(BaseEstimator):
                             ):
                                 if mlflow.active_run() is None:
                                     mlflow.start_run(run_id=self.mlflow_integration.parent_run_id)
-                                self.mlflow_integration.log_model(
-                                    self._trained_estimator.model,
-                                    self.best_estimator,
-                                    signature=self.estimator_signature,
-                                )
-                                self.mlflow_integration.pickle_and_log_automl_artifacts(
-                                    self, self.model, self.best_estimator, signature=self.pipeline_signature
-                                )
+                                if self.best_estimator.endswith("_spark"):
+                                    self.mlflow_integration.log_model(
+                                        self._trained_estimator.model,
+                                        self.best_estimator,
+                                        signature=self.estimator_signature,
+                                        run_id=self.mlflow_integration.parent_run_id,
+                                    )
+                                else:
+                                    self.mlflow_integration.pickle_and_log_automl_artifacts(
+                                        self,
+                                        self.model,
+                                        self.best_estimator,
+                                        signature=self.pipeline_signature,
+                                        run_id=self.mlflow_integration.parent_run_id,
+                                    )
                 else:
-                    logger.info("not retraining because the time budget is too small.")
+                    logger.warning("not retraining because the time budget is too small.")
+        self.wait_futures()
+
+    def wait_futures(self):
+        if self.mlflow_integration is not None:
+            logger.debug("Collecting results from submitted record_state tasks")
+            t1 = time.perf_counter()
+            for future in as_completed(self.mlflow_integration.futures):
+                _task = self.mlflow_integration.futures[future]
+                try:
+                    result = future.result()
+                    logger.debug(f"Result for record_state task {_task}: {result}")
+                except Exception as e:
+                    logger.warning(f"Exception for record_state task {_task}: {e}")
+            for future in as_completed(self.mlflow_integration.futures_log_model):
+                _task = self.mlflow_integration.futures_log_model[future]
+                try:
+                    result = future.result()
+                    logger.debug(f"Result for log_model task {_task}: {result}")
+                except Exception as e:
+                    logger.warning(f"Exception for log_model task {_task}: {e}")
+            t2 = time.perf_counter()
+            logger.debug(f"Collecting results from tasks submitted to executors costs {t2-t1} seconds.")
+        else:
+            logger.debug("No futures to wait for.")
 
     def __del__(self):
         if (

{flaml-2.3.4 → FLAML-2.3.5}/flaml/automl/contrib/histgb.py

@@ -1,7 +1,7 @@
 try:
     from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor
-except ImportError:
-    pass
+except ImportError as e:
+    print(f"scikit-learn is required for HistGradientBoostingEstimator. Please install it; error: {e}")
 
 from flaml import tune
 from flaml.automl.model import SKLearnEstimator