[FR] Add argument (e.g. `save_format`) to specify model save format to `mlflow.xgboost.log_model`

Willingness to contribute

No. I cannot contribute this feature at this time.

Proposal Summary

The XGBoost model file name is currently hardcoded as model.xgb:

https://github.com/mlflow/mlflow/blob/cf96d39d81932423d70d267f8a4ab6640e14ec7e/mlflow/xgboost/__init__.py#L154

This makes it impossible to save an XGBoost model in JSON format. A new argument to specify saving format should be added and the line above should be fixed as follows:

- model_data_subpath = "model.xgb" 
+ model_data_subpath = f"model.{save_format}" 

References:

• https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html
• https://xgboost.readthedocs.io/en/stable/python/python_api.html#xgboost.XGBClassifier.save_model

Motivation

What is the use case for this feature?

To provide a way to save an XGBoost model in JSON format.

Why is this use case valuable to support for MLflow users in general?

Why is this use case valuable to support for your project(s) or organization?

Why is it currently difficult to achieve this use case?

Details

No response

What component(s) does this bug affect?

• 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/pipelines: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
• 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

What interface(s) does this bug affect?

• 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

What language(s) does this bug affect?

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

What integration(s) does this bug affect?

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

[olbapjose]

Just a comment: maybe this would also fix the error thrown when trying to xgboost.log_model with an XGBoost model with categorical variables using the (experimental) support for DataFrame category data type. The error says:

xgboost.core.XGBoostError: [16:56:47] ../src/tree/tree_model.cc:871: Check failed: !HasCategoricalSplit(): 
Please use JSON/UBJSON for saving models with categorical splits.

[harupy]

@AvikantSrivastava Yes, go ahead!