Databricks MLflow Integration for AI Governance Models

You can integrate an instance of Databricks MLflow with OneTrust to enable automatic mapping of model attributes from Databricks MLflow into corresponding fields in OneTrust. This offers Data Scientists a workstream that is more efficient than manual data entry.

Overview Steps

  1. Create a new compute with all the required details in the Azure MLflow Cloud Platform Console

  2. Create a new Notebook in the workspace in Azure MLflow Cloud Platform Console

  3. Create a new registered model and webhook configurations with relevant Events

  4. Generate static IP address for OneTrust to safelist

  5. Webhook URL will get triggered when a registered model is created

  6. Azure Databricks sends a trigger to OneTrust to create the Model record in AI Governance module using Integration Manager

Prerequisites

  1. Login to Azure Databricks portal and launch workspace.

  2. Submit the static IP address to OneTrust, along with the OneTrust environment domain (e.g. app-eu.onetrust.com). Please specify the duration so OneTrust can safelist the IP for the given timeframe.

  3. Create an Authorization API Key needed to access the integration's webhook URL. This can be copied from the first action in the seeded workflow (see below section).

    For more information on this process, see Managing OAuth 2.0 API Keys.

    CODE 
    Example: 
https://staging.1trust.ninja/api/datasubject/v2/requestqueues/en-us?Authorization=Bearer+MDc5ZWI5Mzc4ZGMxNDViYzllMzRhYmFlNjZlNWFiNzE6WlVzN002ZGlBZWNtT1FJWDcwZFdzN0NaWUQxNVdQdXV=

To create a new model

  1. Go to the Azure Databricks Cloud Console and launch workspace.
  2. Create a new compute with the required details.
  3. Create a new Notebook in the workspace.
  4. Install ! pip install databricks-registry-webhooks
  5. Import the MLflow package from databricks-registry-webhooks using:
    PYTHON 
    import mlflow
from databricks_registry_webhooks import RegistryWebhooksClient, HttpUrlSpec
  6. Add the webhook URL configuration along with authentication details:
    CODE 
    http_url_spec = HttpUrlSpec(
  url="https://<hostname>/integrationmanager/api/v1/webhook/<iD>", authorization="Bearer <token>="
)
http_webhook = RegistryWebhooksClient().create_webhook(
  events=["REGISTERED_MODEL_CREATED"],
  http_url_spec=http_url_spec,
  description="Webhook",
  status="ACTIVE"
)
  7. Create a new registered model:
    PYTHON 
    from sklearn.linear_model import LinearRegression

model = LinearRegression()

mlflow.sklearn.log_model(model, "model", registered_model_name="MLflowTest29")

To create the integration

  1. In the Integrations menu, select Gallery.

    The Integration Gallery screen appears.

  2. Use the search function to locate the MLflow tile.

    Integrations gallery - MLflow
  3. Click the tile.

    The Gallery Details screen appears.

  4. Click the Add button for the Databricks Mlflow Model Registration workflow.

    The Add Workflow Name modal appears.


    Databricks Mlflow Model Registration workflow
  5. Complete the fields, then click the Create button.

    The Integration Builder screen appears.

    Note:

    You can select None for the Credential field.


    Integration builder - Add workflow name
  6. The Built-in Webhook Trigger action will look for the model details using the following sample body:
    JSON 
    {
	"event": "REGISTERED_MODEL_CREATED",
	"event_timestamp": 1702026622188,
	"text": "user@onetrust.com created registered model 'MLFlowTest7'.",
	"model_name": "MLFlowTest29",
	"webhook_id": "f0ae8a35b8eb41f29fa9217bf11df9de"
}
  7. The POST Create an Object with Values Defined in the Request action will return the model's name and organization group ID using our provided payload variables:

    POST Create an Object with Values Defined in the Request - Request tab
  8. Save and activate the integration. Whenever a model is now created and registered, a log will captured for the integration trigger and expected model name.