From 0c543b459f476c54d40e3d169715016960735c8d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:47:01 +0000 Subject: [PATCH 1/4] Initial plan From 052cf398c2b84227cd2eb9e06622ddb30c93f505 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:53:24 +0000 Subject: [PATCH 2/4] Fix MLflow pipeline fitted-state detection --- flaml/automl/model.py | 3 +++ test/automl/test_mlflow.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/flaml/automl/model.py b/flaml/automl/model.py index abf4c95c98..d863dd6fec 100644 --- a/flaml/automl/model.py +++ b/flaml/automl/model.py @@ -155,6 +155,9 @@ def get_params(self, deep=False): params["_estimator_type"] = self._estimator_type return params + def __sklearn_is_fitted__(self): + return self._model is not None + def __sklearn_tags__(self): """Override sklearn tags to respect the _estimator_type attribute. diff --git a/test/automl/test_mlflow.py b/test/automl/test_mlflow.py index 5185da9c8e..525cb5aee8 100644 --- a/test/automl/test_mlflow.py +++ b/test/automl/test_mlflow.py @@ -1,14 +1,49 @@ +import os import pickle import mlflow import mlflow.entities +import numpy as np import pytest from pandas import DataFrame from sklearn.datasets import load_iris +from sklearn.model_selection import train_test_split from flaml import AutoML +def test_autologged_model_round_trip(tmp_path): + original_tracking_uri = mlflow.get_tracking_uri() + original_experiment_id = mlflow.tracking.fluent._active_experiment_id + original_experiment_id_env = os.environ.get("MLFLOW_EXPERIMENT_ID") + mlflow.set_tracking_uri(tmp_path.as_uri()) + try: + mlflow.set_experiment("flaml_round_trip") + X, y = load_iris(return_X_y=True, as_frame=True) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + automl = AutoML() + with mlflow.start_run() as run: + automl.fit( + X_train, + y_train, + task="classification", + estimator_list=["rf"], + max_iter=1, + verbose=0, + ) + + loaded = mlflow.sklearn.load_model(f"runs:/{run.info.run_id}/model") + assert np.array_equal(automl.predict(X_test), loaded.predict(X_test)) + finally: + mlflow.end_run() + mlflow.set_tracking_uri(original_tracking_uri) + mlflow.tracking.fluent._active_experiment_id = original_experiment_id + if original_experiment_id_env is None: + os.environ.pop("MLFLOW_EXPERIMENT_ID", None) + else: + os.environ["MLFLOW_EXPERIMENT_ID"] = original_experiment_id_env + + class TestMLFlowLoggingParam: def test_update_and_install_requirements(self): import mlflow From a6fea01d0956b75c267cfbccb793327a89ef1cdd Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:00:09 +0000 Subject: [PATCH 3/4] Strengthen MLflow fitted-state regression test --- test/automl/test_mlflow.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/automl/test_mlflow.py b/test/automl/test_mlflow.py index 525cb5aee8..acdaed86ac 100644 --- a/test/automl/test_mlflow.py +++ b/test/automl/test_mlflow.py @@ -8,6 +8,7 @@ from pandas import DataFrame from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split +from sklearn.utils.validation import check_is_fitted from flaml import AutoML @@ -33,9 +34,9 @@ def test_autologged_model_round_trip(tmp_path): ) loaded = mlflow.sklearn.load_model(f"runs:/{run.info.run_id}/model") + check_is_fitted(loaded) assert np.array_equal(automl.predict(X_test), loaded.predict(X_test)) finally: - mlflow.end_run() mlflow.set_tracking_uri(original_tracking_uri) mlflow.tracking.fluent._active_experiment_id = original_experiment_id if original_experiment_id_env is None: From 813d2cf0eaf912953dcd61ce2c83dc9d543409f1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:05:49 +0000 Subject: [PATCH 4/4] Simplify MLflow regression test isolation --- test/automl/test_mlflow.py | 48 ++++++++++++++------------------------ 1 file changed, 17 insertions(+), 31 deletions(-) diff --git a/test/automl/test_mlflow.py b/test/automl/test_mlflow.py index acdaed86ac..80f3fcb4c3 100644 --- a/test/automl/test_mlflow.py +++ b/test/automl/test_mlflow.py @@ -1,4 +1,3 @@ -import os import pickle import mlflow @@ -13,36 +12,23 @@ from flaml import AutoML -def test_autologged_model_round_trip(tmp_path): - original_tracking_uri = mlflow.get_tracking_uri() - original_experiment_id = mlflow.tracking.fluent._active_experiment_id - original_experiment_id_env = os.environ.get("MLFLOW_EXPERIMENT_ID") - mlflow.set_tracking_uri(tmp_path.as_uri()) - try: - mlflow.set_experiment("flaml_round_trip") - X, y = load_iris(return_X_y=True, as_frame=True) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) - automl = AutoML() - with mlflow.start_run() as run: - automl.fit( - X_train, - y_train, - task="classification", - estimator_list=["rf"], - max_iter=1, - verbose=0, - ) - - loaded = mlflow.sklearn.load_model(f"runs:/{run.info.run_id}/model") - check_is_fitted(loaded) - assert np.array_equal(automl.predict(X_test), loaded.predict(X_test)) - finally: - mlflow.set_tracking_uri(original_tracking_uri) - mlflow.tracking.fluent._active_experiment_id = original_experiment_id - if original_experiment_id_env is None: - os.environ.pop("MLFLOW_EXPERIMENT_ID", None) - else: - os.environ["MLFLOW_EXPERIMENT_ID"] = original_experiment_id_env +def test_autologged_model_round_trip(): + X, y = load_iris(return_X_y=True, as_frame=True) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + automl = AutoML() + with mlflow.start_run() as run: + automl.fit( + X_train, + y_train, + task="classification", + estimator_list=["rf"], + max_iter=1, + verbose=0, + ) + + loaded = mlflow.sklearn.load_model(f"runs:/{run.info.run_id}/model") + check_is_fitted(loaded) + assert np.array_equal(automl.predict(X_test), loaded.predict(X_test)) class TestMLFlowLoggingParam: