mlflow.xgboost

The mlflow.xgboost module provides an API for logging and loading XGBoost models. This module exports XGBoost models with the following flavors:

XGBoost (native) format

This is the main flavor that can be loaded back into XGBoost.

mlflow.pyfunc

Produced for use by generic pyfunc-based deployment tools and batch inference.

mlflow.xgboost.autolog(importance_types=['weight'])[source]

Note

Experimental: This method may change or be removed in a future release without warning.

Enables automatic logging from XGBoost to MLflow. Logs the following.

  • parameters specified in xgboost.train.

  • metrics on each iteration (if evals specified).

  • metrics at the best iteration (if early_stopping_rounds specified).

  • feature importance as JSON files and plots.

  • trained model.

Note that the scikit-learn API is not supported.

Parameters

importance_types – importance types to log.

mlflow.xgboost.get_default_conda_env()[source]
Returns

The default Conda environment for MLflow Models produced by calls to save_model() and log_model().

mlflow.xgboost.load_model(model_uri)[source]

Load an XGBoost model from a local file or a run.

Parameters

model_uri

The location, in URI format, of the MLflow model. For example:

  • /Users/me/path/to/local/model

  • relative/path/to/local/model

  • s3://my_bucket/path/to/model

  • runs:/<mlflow_run_id>/run-relative/path/to/model

For more information about supported URI schemes, see Referencing Artifacts.

Returns

An XGBoost model (an instance of xgboost.Booster)

mlflow.xgboost.log_model(xgb_model, artifact_path, conda_env=None, registered_model_name=None, **kwargs)[source]

Log an XGBoost model as an MLflow artifact for the current run.

Parameters
  • xgb_model – XGBoost model (an instance of xgboost.Booster) to be saved. Note that models that implement the scikit-learn API are not supported.

  • artifact_path – Run-relative artifact path.

  • conda_env

    Either a dictionary representation of a Conda environment or the path to a Conda environment yaml file. If provided, this describes the environment this model should be run in. At minimum, it should specify the dependencies contained in get_default_conda_env(). If None, the default get_default_conda_env() environment is added to the model. The following is an example dictionary representation of a Conda environment:

    {
        'name': 'mlflow-env',
        'channels': ['defaults'],
        'dependencies': [
            'python=3.7.0',
            'pip': [
                'xgboost==0.90'
            ]
        ]
    }
    

  • registered_model_name – Note:: Experimental: This argument may change or be removed in a future release without warning. If given, create a model version under registered_model_name, also creating a registered model if one with the given name does not exist.

  • kwargs – kwargs to pass to xgboost.Booster.save_model method.

mlflow.xgboost.save_model(xgb_model, path, conda_env=None, mlflow_model=<mlflow.models.Model object>)[source]

Save an XGBoost model to a path on the local file system.

Parameters
  • xgb_model – XGBoost model (an instance of xgboost.Booster) to be saved. Note that models that implement the scikit-learn API are not supported.

  • path – Local path where the model is to be saved.

  • conda_env

    Either a dictionary representation of a Conda environment or the path to a Conda environment yaml file. If provided, this describes the environment this model should be run in. At minimum, it should specify the dependencies contained in get_default_conda_env(). If None, the default get_default_conda_env() environment is added to the model. The following is an example dictionary representation of a Conda environment:

    {
        'name': 'mlflow-env',
        'channels': ['defaults'],
        'dependencies': [
            'python=3.7.0',
            'pip': [
                'xgboost==0.90'
            ]
        ]
    }
    

  • mlflow_modelmlflow.models.Model this flavor is being added to.