Skip to content

Commit c73ff99

Browse files
Copilotthinkall
andauthored
Fix MLflow autologged model reload fitted-state detection for scikit-learn>=1.8.0 (#1573)
* Initial plan * Fix MLflow pipeline fitted-state detection * Strengthen MLflow fitted-state regression test * Simplify MLflow regression test isolation --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Li Jiang <bnujli@gmail.com>
1 parent 91cc4da commit c73ff99

2 files changed

Lines changed: 25 additions & 0 deletions

File tree

flaml/automl/model.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,9 @@ def get_params(self, deep=False):
155155
params["_estimator_type"] = self._estimator_type
156156
return params
157157

158+
def __sklearn_is_fitted__(self):
159+
return self._model is not None
160+
158161
def __sklearn_tags__(self):
159162
"""Override sklearn tags to respect the _estimator_type attribute.
160163

test/automl/test_mlflow.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,35 @@
22

33
import mlflow
44
import mlflow.entities
5+
import numpy as np
56
import pytest
67
from pandas import DataFrame
78
from sklearn.datasets import load_iris
9+
from sklearn.model_selection import train_test_split
10+
from sklearn.utils.validation import check_is_fitted
811

912
from flaml import AutoML
1013

1114

15+
def test_autologged_model_round_trip():
16+
X, y = load_iris(return_X_y=True, as_frame=True)
17+
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
18+
automl = AutoML()
19+
with mlflow.start_run() as run:
20+
automl.fit(
21+
X_train,
22+
y_train,
23+
task="classification",
24+
estimator_list=["rf"],
25+
max_iter=1,
26+
verbose=0,
27+
)
28+
29+
loaded = mlflow.sklearn.load_model(f"runs:/{run.info.run_id}/model")
30+
check_is_fitted(loaded)
31+
assert np.array_equal(automl.predict(X_test), loaded.predict(X_test))
32+
33+
1234
class TestMLFlowLoggingParam:
1335
def test_update_and_install_requirements(self):
1436
import mlflow

0 commit comments

Comments
 (0)