[Bug]: MLflow-autolog'd FLAML model reloads as unfitted Pipeline via mlflow.sklearn.load_model

[immu4989]

Describe the bug

The canonical MLflow round-trip example currently published in Best-Practices.md → "Option 1: MLflow logging (recommended for production)" does not work on current main against recent MLflow versions: the model artifact that FLAML's built-in MLflow autologging path writes to runs:/{run_id}/model reloads as an unfitted sklearn.pipeline.Pipeline, raising sklearn.exceptions.NotFittedError on .predict().

Reproduced verbatim from the doc on:

• flaml from current main (post-fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561, merge SHA 50f3d1be)
• mlflow==2.22.1
• scikit-learn==1.8.0
• Python 3.12.13

Steps to reproduce

import warnings, tempfile, numpy as np
warnings.simplefilter("ignore")
import mlflow
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from flaml import AutoML

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)

with tempfile.TemporaryDirectory() as d:
    mlflow.set_tracking_uri(f"file://{d}/mlruns")
    mlflow.set_experiment("flaml_repro")
    automl = AutoML()
    with mlflow.start_run(run_name="r") as run:
        automl.fit(X_train, y_train, task="classification", time_budget=3, verbose=0)
        run_id = run.info.run_id
    loaded = mlflow.sklearn.load_model(f"runs:/{run_id}/model")
    loaded.predict(X_test)  # raises NotFittedError

Actual behavior

sklearn.exceptions.NotFittedError: This Pipeline instance is not fitted yet.
Call 'fit' with appropriate arguments before using this estimator.
  at sklearn/pipeline.py:740 in predict
  at sklearn/utils/validation.py:1705 in check_is_fitted

The loaded artifact is a sklearn.pipeline.Pipeline, not the original flaml.AutoML instance, and it has no fitted state.

Expected behavior

Per the published example, loaded.predict(X_test) should return predictions equal to automl.predict(X_test).

Verified workaround

Explicitly calling mlflow.sklearn.log_model(automl, artifact_path=...) (after suppressing FLAML's autologging with mlflow_logging=False) round-trips correctly:

with mlflow.start_run(run_name="r") as run:
    automl.fit(X_train, y_train, task="classification", time_budget=3, verbose=0, mlflow_logging=False)
    mlflow.sklearn.log_model(automl, artifact_path="flaml_model")
    run_id = run.info.run_id

loaded = mlflow.sklearn.load_model(f"runs:/{run_id}/flaml_model")
assert np.array_equal(automl.predict(X_test), loaded.predict(X_test))

This workaround is documented in §1.2 of the new Production-Deployment guide (PR #1562) as the path that works on current MLflow versions. The autolog-based example in Best-Practices.md should ideally be either updated to the explicit-log pattern or fixed at the FLAML autolog level so the published recommended pattern works again.

Discovered during

Pre-flight verification of the patterns documented in PR #1562 (production-deployment guide).

Environment

• FLAML: current main (verified through eb6adcff, the merge commit of docs: production-deployment guide for trained FLAML models #1562)
• MLflow: 2.22.1
• scikit-learn: 1.8.0
• Python: 3.12.13
• OS: macOS (Darwin)

[thinkall]

Thanks a lot for reporting this, @immu4989 ! Does it work for an older version of mlflow or an older version of scikit-learn?