Repository: mlflow/mlflow

Issue:
[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](https://xgboost.readthedocs.io/en/stable/tutorials/saving_model.html). A new argument to specify saving format should be added and the line above should be fixed as follows:

```diff
- 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
- [X] `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


Discussion:
What about add a param to log_model like "model_format" ? Available values can be "json" or "binary". And we can consider add this param to autolog.
@WeichenXu123 That sounds good. I'll fix the PR title and description.
<!-- assign-maintainer -->
@BenWilson2 @dbczumar @harupy @WeichenXu123 Please assign a maintainer and start triaging this issue.
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.
```
@olbapjose Awesome! Would you be interested in contributing support for this feature?
@dbczumar Thanks a lot but unfortunately I do not have time these weeks... despite I have been struggling very hard to log an XGBoost model and then retrieve it. No way with Spark XGBoost (it was possible to log it but not to retrieve it when it is a multi-class model), and no way with python xgboost as it contains category columns...
@olbapjose Got it! Understood! Thank you for your feedback; we're excited about getting this issue resolved.
Hey @dbczumar @harupy can I work on this issue?
@AvikantSrivastava Yes, go ahead!

Resolving patch:
diff --git a/mlflow/xgboost/__init__.py b/mlflow/xgboost/__init__.py
--- a/mlflow/xgboost/__init__.py
+++ b/mlflow/xgboost/__init__.py
@@ -103,6 +103,7 @@ def save_model(
     input_example: ModelInputExample = None,
     pip_requirements=None,
     extra_pip_requirements=None,
+    model_format="xgb",
 ):
     """
     Save an XGBoost model to a path on the local file system.
@@ -136,6 +137,7 @@ def save_model(
                           base64-encoded.
     :param pip_requirements: {{ pip_requirements }}
     :param extra_pip_requirements: {{ extra_pip_requirements }}
+    :param model_format: File format in which the model is to be saved.
     """
     import xgboost as xgb
 
@@ -151,7 +153,7 @@ def save_model(
         mlflow_model.signature = signature
     if input_example is not None:
         _save_example(mlflow_model, input_example, path)
-    model_data_subpath = "model.xgb"
+    model_data_subpath = f"model.{model_format}"
     model_data_path = os.path.join(path, model_data_subpath)
 
     # Save an XGBoost model
@@ -169,6 +171,7 @@ def save_model(
         xgb_version=xgb.__version__,
         data=model_data_subpath,
         model_class=xgb_model_class,
+        model_format=model_format,
         code=code_dir_subpath,
     )
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
@@ -219,6 +222,7 @@ def log_model(
     await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
     pip_requirements=None,
     extra_pip_requirements=None,
+    model_format=None,
     **kwargs,
 ):
     """
@@ -258,6 +262,7 @@ def log_model(
                             waits for five minutes. Specify 0 or None to skip waiting.
     :param pip_requirements: {{ pip_requirements }}
     :param extra_pip_requirements: {{ extra_pip_requirements }}
+    :param model_format: File format in which the model is to be saved.
     :param kwargs: kwargs to pass to `xgboost.Booster.save_model`_ method.
     :return: A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
              metadata of the logged model.
@@ -267,6 +272,7 @@ def log_model(
         flavor=mlflow.xgboost,
         registered_model_name=registered_model_name,
         xgb_model=xgb_model,
+        model_format=model_format,
         conda_env=conda_env,
         code_paths=code_paths,
         signature=signature,


Test patch:
diff --git a/tests/xgboost/test_xgboost_model_export.py b/tests/xgboost/test_xgboost_model_export.py
--- a/tests/xgboost/test_xgboost_model_export.py
+++ b/tests/xgboost/test_xgboost_model_export.py
@@ -519,3 +519,14 @@ def test_log_model_with_code_paths(xgb_model):
         _compare_logged_code_paths(__file__, model_uri, mlflow.xgboost.FLAVOR_NAME)
         mlflow.xgboost.load_model(model_uri=model_uri)
         add_mock.assert_called()
+
+
+@pytest.mark.parametrize("model_format", ["xgb", "json", "ubj"])
+def test_log_model_with_model_format(xgb_model, model_format):
+    with mlflow.start_run():
+        model_info = mlflow.xgboost.log_model(xgb_model.model, "model", model_format=model_format)
+        loaded_model = mlflow.xgboost.load_model(model_info.model_uri)
+        np.testing.assert_array_almost_equal(
+            xgb_model.model.predict(xgb_model.inference_dmatrix),
+            loaded_model.predict(xgb_model.inference_dmatrix),
+        )