Autologging functionality for scikit-learn integration with XGBoost (Part 1)

[jwyyy]

Signed-off-by: Junwen Yao jwyiao@gmail.com

What changes are proposed in this pull request?

This is the first PR to add autologging for XGBoost sklearn models. It revises model saving / loading and adds model_class to specify XGBoost model class, such as Booster, or XGBRegressor.

A separate PR will be made to add autologging for XGBoost sklearn models using mlflow.sklearn routine.

(Draft + discussion: #4885)

How is this patch tested?

New test functions are added. See PR files.

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

This PR will enable saving / loading XGBoost models, including sklearn models, with model class specification.
Functions save_model() / load_model() in mlflow.xgboost can be used as before.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[jwyyy]

Hi @dbczumar, as per our discussion, I separate the autologging functionality and model saving / loading with model class specification into two PRs. This is the first one to address model saving / loading. Please let me know if I missed anything. Thanks!

[dbczumar]

@jwyyy Looks awesome! I left a few tiny comments - I think this is very close!

[jwyyy]

Hi @dbczumar @harupy, thank you for your feedback! I updated this PR and doc according to your suggestions. Please let me know if they look good. Thanks!

[jwyyy]

Hi @harupy @dbczumar, thank you for initializing the auto check! According to the error messages, two unsuccessful checks were due to network connection exception, not related to the changes in this PR. I am not sure if I can fix these two problems efficiently. Can you take a look and re-try the check? Thanks! Also please let me know if the latest commit look good.

[jwyyy]

Hi @harupy @dbczumar, I notice that this PR #4997 tries to add a new argument to load_model() in every flavor. This will have an impact on how we handle model loading in mlflow.xgboost (although the revision should be very easy). I am wondering which PR will be addressed first. Or maybe we can directly add the xgboost part in this PR. Please let you know your plan / ideas. Thanks!

(The previous failed check (35/36) was due to Keras 2.7 compatibility, which should be fixed now.)

[harupy]

I notice that this PR #4997 tries to add a new argument to load_model() in every flavor.

Yes it does, but the new argument doesn't change how to load the model, it just allows users to determine where to load the model. Therefore, it should not conflict with this PR.

[harupy]

Looks good to me once #4954 (comment) is addressed!

[jwyyy]

Looks good to me once #4954 (comment) is addressed!

@harupy I just updated the comments in this PR, including #4954 (review).

[jwyyy]

@harupy @dbczumar Thank you very much for your feedback and suggestions! It seems we are at the final stage of this PR. Once it is merged, I will make the second PR which completes the autologging for XGBoost sklearn models based on our previous discussion. It should be ready within the next few days. Meanwhile I will also submit the first PR for LightGBM models (similar to this PR). Thanks again!

[dbczumar]

LGTM! Thanks so much @jwyyy !

[jwyyy]

Two failed tests seem unrelated to this PR. The error messages are the same as unsuccessful checks in recent commits to master #5032 #5031 #4999 #5029.

[harupy]

Two failed tests seem unrelated to this PR.

Yep, we're already investigating them. Let's merge this PR!