MLflow integration for DataRobot¶

The MLflow integration for DataRobot allows you to export a model from MLflow and import it into the DataRobot Model Registry, creating key values from the training parameters, metrics, tags, and artifacts in the MLflow model.

Prerequisites for the MLflow integration¶

The MLflow integration for DataRobot requires the following:

• Python >= 3.9
• DataRobot >= 9.0

This integration library uses a preview API endpoint; the DataRobot user associated with your API token must have Owner or User permissions for the DataRobot model package.

Install the MLflow integration for DataRobot¶

You can install the datarobot-mlfow integration with pip:

pip install datarobot-mlflow

If you are running the integration on Azure, use the following command:

pip install "datarobot-mlflow[azure]"

Configure command line options¶

The following command line options are available for the drflow_cli:

The following command line operations are available for the --action option:

Set environment variables¶

In addition to the command line options above, you should also provide any environment variables required for your use case:

You can use export to define these environment variables with the information required for your use case:

export MLOPS_API_TOKEN="<dr-api-key>"
export AZURE_TENANT_ID="<tenant-id>"
export AZURE_CLIENT_ID="<client-id>"
export AZURE_CLIENT_SECRET="<secret>"

Run the sync action to import a model from MLflow into DataRobot¶

You can use the command line options and actions defined above to export MLflow model information from MLflow and import it into the DataRobot Model Registry:

DR_MODEL_ID="<MODEL_PACKAGE_ID>"

env PYTHONPATH=./ \
python datarobot_mlflow/drflow_cli.py \
  --mlflow-url http://localhost:8080 \
  --mlflow-model cost-model  \
  --mlflow-model-version 2 \
  --dr-model $DR_MODEL_ID \
  --dr-url https://app.datarobot.com \
  --with-artifacts \
  --verbose \
  --action sync

After you run this command successfully, you can see MLflow information on the Key Values tab of a Registered Model version:

In addition, in the Activity log of the Key Values tab, you can view a record of the key value creation events:

Troubleshoot Azure AD Service Principal credentials¶

To validate Azure AD Service Principal credentials for troubleshooting purposes, you can use the following command line example:

export MLOPS_API_TOKEN="n/a"  # not used for Azure auth check, but the environment variable must be present

env PYTHONPATH=./ \
python datarobot_mlflow/drflow_cli.py \
  --verbose \
  --auth-type azure-service-principal \
  --service-provider-type azure-databricks \
  --action validate-auth

This command should produce the following output if you haven't configured the required environment variables:

Required environment variable is not defined: AZURE_TENANT_ID
Required environment variable is not defined: AZURE_CLIENT_ID
Required environment variable is not defined: AZURE_CLIENT_SECRET
Azure AD Service Principal credentials are not valid; check environment variables

If you see this error, provide the required Azure AD Service Principal credentials as environment variables:

export AZURE_TENANT_ID="<tenant-id>"
export AZURE_CLIENT_ID="<client-id>"
export AZURE_CLIENT_SECRET="<secret>"

When the environment variables for the Azure AD Service Principal credentials are defined, you should see the following output:

Azure AD Service Principal credentials are valid for obtaining access token