DataRobot Provider for Apache Airflow

This package provides operators, sensors, and a hook to integrate DataRobot into Apache Airflow. Using these components, you should be able to build the essential DataRobot pipeline - create a project, train models, deploy a model, and score predictions against the model deployment.

Install the Airflow provider

The DataRobot provider for Apache Airflow requires an environment with the following dependencies installed:

• Apache Airflow >= 2.3

• DataRobot Python API Client >= 3.2.0

To install the DataRobot provider, run the following command:

Copy
pip install airflow-provider-datarobot

Create a connection from Airflow to DataRobot

The next step is to create a connection from Airflow to DataRobot:

1. In the Airflow user interface, click Admin > Connections to add an Airflow connection.

2. On the List Connection page, click + Add a new record.

3. In the Add Connection dialog box, configure the following fields:

   Field	Description
   Connection Id	datarobot_default (this name is used by default in all operators)
   Connection Type	DataRobot
   API Key	A DataRobot API key, created in the DataRobot Developer Tools, from the API Keys section.
   DataRobot endpoint URL	https://app.datarobot.com/api/v2 by default

4. Click Test to establish a test connection between Airflow and DataRobot.

5. When the connection test is successful, click Save.

Create preconfigured connections to DataRobot

You can create preconfigured connections to store and manage credentials to use with Airflow Operators, replicating the connection on the DataRobot side.

Currently, the supported credential types are:

After creating a preconfigured connection through the Airflow UI or API, you can access your stored credentials with GetOrCreateCredentialOperator or GetOrCreateDataStoreOperator to replicate them in DataRobot and retrieve the corresponding credentials_id or datastore_id.

JSON configuration for the DAG run

Operators and sensors use parameters from the config JSON submitted when triggering the DAG; for example:

Copy
{
    "training_data": "s3-presigned-url-or-local-path-to-training-data",
    "project_name": "Project created from Airflow",
    "autopilot_settings": {
        "target": "readmitted"
    },
    "deployment_label": "Deployment created from Airflow",
    "score_settings": {
        "intake_settings": {
            "type": "s3",
            "url": "s3://path/to/scoring-data/Diabetes10k.csv",
            "credential_id": "<credential_id>"
        },
        "output_settings": {
            "type": "s3",
            "url": "s3://path/to/results-dir/Diabetes10k_predictions.csv",
            "credential_id": "<credential_id>"
        }
    }
}

These config values are accessible in the execute() method of any operator in the DAG through the context["params"] variable; for example, to get training data, you could use the following:

Copy
def execute(self, context: Dict[str, Any]) -> str:
    ...
    training_data = context["params"]["training_data"]
    ...

Modules

Operators

GetOrCreateCredentialOperator

Fetches a credential by name. This operator attempts to find a DataRobot credential with the provided name. If the credential doesn't exist, the operator creates it using the Airflow preconfigured connection with the same connection name.

Returns a credential ID.

Required config parameters:

GetOrCreateDataStoreOperator

Fetches a DataStore by Connection name. If the DataStore does not exist, the operator attempts to create it using Airflow preconfigured connection with the same connection name.

Returns a credential ID.

Required config params:

CreateDatasetFromDataStoreOperator

Loads a dataset from a JDBC Connection to the DataRobot AI Catalog.

Returns a dataset ID.

Required config params:

UploadDatasetOperator

Uploads a local file to the DataRobot AI Catalog.

Returns a dataset ID.

Required config params:

UpdateDatasetFromFileOperator

Creates a new dataset version from a file.

Returns a dataset version ID when the new version uploads successfully.

Required config params:

CreateDatasetVersionOperator

Creates a new version of the existing dataset in the AI Catalog.

Returns a dataset version ID.

Required config params:

CreateOrUpdateDataSourceOperator

Creates a data source or updates it if it already exists.

Returns a DataRobot DataSource ID.

Required config params:

CreateProjectOperator

Creates a DataRobot project.

Returns a project ID.

Several options of source dataset supported:

Create a project directly from a local file or a pre-signed S3 URL.

Required config params:

Note: In case of an S3 input, the training_data value must be a pre-signed AWS S3 URL.

Create a project from an existing dataset in the DataRobot AI Catalog using a dataset ID defined in the config file.

Required config params:

Create a project from an existing dataset in the DataRobot AI Catalog using a dataset ID from the previous operator. In this case, your previous operator must return a valid dataset ID (for example UploadDatasetOperator) and you should use this output value as a dataset_id argument in the CreateProjectOperator object creation step.

Required config params:

For more project settings, see the DataRobot documentation.

TrainModelsOperator

Runs DataRobot Autopilot to train models.

Returns None.

Parameters:

Required config params:

Copy
"autopilot_settings": {
    "target": "readmitted"
} 

For more autopilot settings, see the DataRobot documentation.

DeployModelOperator

Deploy a specified model.

Returns a deployment ID.

Parameters:

Required config params:

For more deployment settings, see the DataRobot documentation.

DeployRecommendedModelOperator

Deploys a recommended model.

Returns a deployment ID.

Parameters:

Required config params:

For more deployment settings, see the DataRobot documentation.

ScorePredictionsOperator

Scores batch predictions against the deployment.

Returns a batch prediction job ID.

Prerequisites:

• Use GetOrCreateCredentialOperator to pass a credential_id from the preconfigured DataRobot Credentials (Airflow Connections) or manually set the credential_id parameter in the config.

  Note: You can add S3 credentials to DataRobot via the Python API client.

• Or use a Dataset ID from the DataRobot AI Catalog.

• Or use a DataStore ID for a JDBC source connection; you can use GetOrCreateDataStoreOperator to pass datastore_id from a preconfigured Airflow Connection.

Parameters:

Copy
"score_settings": {
    "intake_settings": {
        "type": "s3",
        "url": "s3://my-bucket/Diabetes10k.csv",
    },
    "output_settings": {
        "type": "s3",
        "url": "s3://my-bucket/Diabetes10k_predictions.csv",
    }
}

Copy
"score_settings": {
    "intake_settings": {
        "type": "s3",
        "url": "s3://my-bucket/Diabetes10k.csv",
        "credential_id": "<credential_id>"
    },
    "output_settings": {
        "type": "s3",
        "url": "s3://my-bucket/Diabetes10k_predictions.csv",
        "credential_id": "<credential_id>"
    }
}

Copy
"score_settings": {
    "intake_settings": {
        "type": "dataset",
        "dataset_id": "<datasetId>",
    },
    "output_settings": { }
}

For more batch prediction settings, see the DataRobot documentation.

GetTargetDriftOperator

Gets the target drift from a deployment.

Returns a dict with the target drift data.

Parameters:

No config params are required; however, the optional params may be passed in the config as follows:

Copy
"target_drift": { }

GetFeatureDriftOperator

Gets the feature drift from a deployment.

Returns a dict with the feature drift data.

Parameters:

No config params are required; however, the optional params may be passed in the config as follows:

Copy
"feature_drift": { }

GetServiceStatsOperator

Gets service stats measurements from a deployment.

Returns a dict with the service stats measurements data.

Parameters:

No config params are required; however, the optional params may be passed in the config as follows:

Copy
"service_stats": { }

GetAccuracyOperator

Gets the accuracy of a deployment’s predictions.

Returns a dict with the accuracy for a deployment.

Parameters:

No config params are required; however, the optional params may be passed in the config as follows:

Copy
"accuracy": { }

GetBiasAndFairnessSettingsOperator

Gets the Bias And Fairness settings for deployment.

Returns a dict with the Bias And Fairness settings for a Deployment.

Parameters:

No config params are required.

UpdateBiasAndFairnessSettingsOperator

Updates the Bias And Fairness settings for deployment.

Parameters:

Sample config params:

Copy
"protected_features": ["attribute1"],
"preferable_target_value": "True",
"fairness_metrics_set": "equalParity",
"fairness_threshold": 0.1,

GetSegmentAnalysisSettingsOperator

Gets the segment analysis settings for a deployment.

Returns a dict with the segment analysis settings for a deployment

Parameters:

No config params are required.

UpdateSegmentAnalysisSettingsOperator

Updates the segment analysis settings for a deployment.

Parameters:

Sample config params:

Copy
"segment_analysis_enabled": True,
"segment_analysis_attributes": ["attribute1", "attribute2"],

GetMonitoringSettingsOperator

Gets the monitoring settings for deployment.

Returns a dict with the config params for a deployment.

Parameters:

No config params are required.

Sample monitoring settings:

Copy
{
"drift_tracking_settings": {  } 
"association_id_settings": {  }
"predictions_data_collection_settings": {  }
}

UpdateMonitoringSettingsOperator

Updates monitoring settings for a deployment.

Parameters:

Sample config params:

Copy
"target_drift_enabled": True,
"feature_drift_enabled": True,
"association_id_column": ["id"],
"required_association_id": False,
"predictions_data_collection_enabled": False,

BatchMonitoringOperator

Creates a batch monitoring job for the deployment.

Returns a batch monitoring job ID.

Prerequisites:

• Use GetOrCreateCredentialOperator to pass a credential_id from the preconfigured DataRobot Credentials (Airflow Connections) or manually set the credential_id parameter in the config.

  Note: You can add S3 credentials to DataRobot via the Python API client.

• Or use a Dataset ID from the DataRobot AI Catalog.

• Or use a DataStore ID for a JDBC source connection; you can use GetOrCreateDataStoreOperator to pass datastore_id from a preconfigured Airflow Connection.

Parameters:

Sample config params:

Copy
"deployment_id": "61150a2fadb5586af4118980",
"monitoring_settings": {
    "intake_settings": {
        "type": "bigquery",
        "dataset": "integration_example_demo",
        "table": "actuals_demo",
        "bucket": "datarobot_demo_airflow",
    },
    "monitoring_columns": {
        "predictions_columns": [
            {"class_name": "True", "column_name": "target_True_PREDICTION"},
            {"class_name": "False", "column_name": "target_False_PREDICTION"},
        ],
        "association_id_column": "id",
        "actuals_value_column": "ACTUAL",
    },
}

Copy
"deployment_id": "61150a2fadb5586af4118980",
"monitoring_settings": {
    "intake_settings": {
        "type": "bigquery",
        "dataset": "integration_example_demo",
        "table": "actuals_demo",
        "bucket": "datarobot_demo_airflow",
        "credential_id": "<credential_id>"
    },
    "monitoring_columns": {
        "predictions_columns": [
            {"class_name": "True", "column_name": "target_True_PREDICTION"},
            {"class_name": "False", "column_name": "target_False_PREDICTION"},
        ],
        "association_id_column": "id",
        "actuals_value_column": "ACTUAL",
    },
}

For more batch monitoring settings, see the DataRobot documentation.

DownloadModelScoringCodeOperator

Downloads scoring code artifact from a model.

Parameters:

Sample config params:

Copy
"source_code": False,

For more scoring code download parameters, see the DataRobot documentation.

DownloadDeploymentScoringCodeOperator

Downloads scoring code artifact from a deployment.

Parameters:

Sample config params:

Copy
"source_code": False,
"include_agent": False,
"include_prediction_explanations": False,
"include_prediction_intervals": False,

For more scoring code download parameters, see the DataRobot documentation.

SubmitActualsFromCatalogOperator

Downloads scoring code artifact from a deployment.

Returns an actuals upload job ID.

Parameters:

Sample config params:

Copy
"association_id_column": "id",
"actual_value_column": "ACTUAL",
"timestamp_column": "timestamp",

StartAutopilotOperator

Triggers DataRobot Autopilot to train a set of models.

Parameters:

Sample config params:

Copy
"autopilot_settings": {
    "target": "column_name",
    "mode": AUTOPILOT_MODE.QUICK,
}

For more analyze_and_model parameters, see the DataRobot documentation.

CreateExecutionEnvironmentOperator

Create an execution environment.

Returns an execution environment ID.

Parameters:

Sample config params:

Copy
"execution_environment_name": "My Demo Env",
"custom_model_description": "This is a custom model created by Airflow",
"programming_language": "python",

For more execution environment creation parameters, see the DataRobot documentation.

CreateExecutionEnvironmentVersionOperator

Create an execution environment version.

Returns a version ID for the newly created execution environment .

Parameters:

For more execution environment version creation parameters, see the DataRobot documentation.

CreateCustomInferenceModelOperator

Create a custom inference model.

Returns the ID for the created custom model.

Parameters:

Sample DAG config params:

Copy
"target_type": - Target type of the custom inference model.
    Values: [`datarobot.TARGET_TYPE.BINARY`, `datarobot.TARGET_TYPE.REGRESSION`,
    `datarobot.TARGET_TYPE.MULTICLASS`, `datarobot.TARGET_TYPE.UNSTRUCTURED`]
"target_name": - Target feature name.
    It is optional (ignored if provided) for `datarobot.TARGET_TYPE.UNSTRUCTURED` target type.
"programming_language": - Programming language of the custom learning model.
"positive_class_label": - Custom inference model positive class label for binary classification.
"negative_class_label": - Custom inference model negative class label for binary classification.
"prediction_threshold": - Custom inference model prediction threshold.
"class_labels": - Custom inference model class labels for multiclass classification.
"network_egress_policy": - Determines whether the given custom model is isolated, or can access the public network.
"maximum_memory": - The maximum memory that might be allocated by the custom model.
"replicas": - A fixed number of replicas that will be deployed in the cluster.

For more custom inference model creation parameters, see the DataRobot documentation.

CreateCustomModelVersionOperator

Create a custom model version.

Returns the version ID for the created custom model.

Parameters:

Sample DAG config params:

Copy
"is_major_update" - The flag defining if a custom model version will be a minor or a major version.
"files" - The list of tuples, where values in each tuple are the local filesystem path and
            the path the file should be placed in the model.
"files_to_delete" - The list of a file items IDs to be deleted.
"network_egress_policy": - Determines whether the given custom model is isolated, or can access the public network.
"maximum_memory": - The maximum memory that might be allocated by the custom model.
"replicas": - A fixed number of replicas that will be deployed in the cluster.
"required_metadata_values" - Additional parameters required by the execution environment.
"keep_training_holdout_data" - If the version should inherit training and holdout data from the previous version.

For more custom inference model creation parameters, see the DataRobot documentation.

CustomModelTestOperator

Create and start a custom model test.

Returns an ID for the custom model test.

Parameters:

Sample DAG config params:

Copy
"network_egress_policy": - Determines whether the given custom model is isolated, or can access the public network.
"maximum_memory": - The maximum memory that might be allocated by the custom model.
"replicas": - A fixed number of replicas that will be deployed in the cluster.

For more custom model test creation parameters, see the DataRobot documentation.

GetCustomModelTestOverallStatusOperator

Get the overall status for custom model tests.

Returns the custom model test status.

Parameters:

For more custom model test get status parameters, see the DataRobot documentation.

CreateCustomModelDeploymentOperator

Create a deployment from a DataRobot custom model image.

Returns the deployment ID.

Parameters:

For more create_from_custom_model_version parameters, see the DataRobot documentation.

GetDeploymentModelOperator

Gets information about the deployment's current model.

Returns a model information from a deployment.

Parameters:

For more get deployment parameters, see the DataRobot documentation.

ReplaceModelOperator

Replaces the current model for a deployment.

Returns model information for the mode replacing the deployed model.

Parameters:

For more replace_model parameters, see the DataRobot documentation.

ActivateDeploymentOperator

Activate or deactivate a Deployment.

Returns the Deployment status (active or inactive).

Parameters:

For more activate deployment parameters, see the DataRobot documentation.

GetDeploymentStatusOperator

Get the deployment status (active or inactive).

Returns the deployment status.

Parameters:

For more deployment parameters, see the DataRobot documentation.

RelationshipsConfigurationOperator

Creates a relationship configuration.

Returns the relationships configuration ID.

Parameters:

For more Feature Discovery parameters, see the DataRobot documentation.

DatasetDefinitionOperator

Creates a dataset definition for Feature Discovery.

Returns a dataset definition dict.

Parameters:

For more create-dataset-definitions-and-relationships-using-helper-functions, see the DataRobot documentation.

DatasetRelationshipOperator

Create a relationship between datasets defined in DatasetDefinition.

Returns a dataset definition dict.

Parameters:

For more create-dataset-definitions-and-relationships-using-helper-functions, see the DataRobot documentation.

ComputeFeatureImpactOperator

Creates a Feature Impact job in DataRobot.

Returns a Feature Impact job ID.

Parameters:

For more request_feature_impact, see the DataRobot documentation.

ComputeFeatureEffectsOperator

Submit a request to compute Feature Effects for the model.

Returns the Feature Effects job ID.

Parameters:

For more request_feature_impact parameters, see the DataRobot documentation.

ComputeShapOperator

Submit a request to compute a SHAP impact job for the model.

Returns a SHAP impact job ID.

Parameters:

For more shap-impact parameters, see the DataRobot documentation.

CreateExternalModelPackageOperator

Create an external model package in DataRobot MLOps from JSON configuration.

Returns a model package ID of the newly created model package.

Parameters:

Example of JSON configuration for a regression model:

.. code-block:: python

Copy
{
     "name": "Lending club regression",
     "modelDescription": {
             "description": "Regression on lending club dataset"
         }
     "target": {
         "type": "Regression",
         "name": "loan_amnt"
     }
}

Example JSON for a binary classification model:

.. code-block:: python

Copy
{
    "name": "Surgical Model",
    "modelDescription": {
        "description": "Binary classification on surgical dataset",
        "location": "/tmp/myModel"
        },
        "target": {
            "type": "Binary",
            "name": "complication",
            "classNames": ["Yes","No"],  # minority/positive class should be listed first
            "predictionThreshold": 0.5
        }
    }
}

Example JSON for a multiclass classification model:

.. code-block:: python

Copy
{
    "name": "Iris classifier",
    "modelDescription": {
    "description": "Classification on iris dataset",
    "location": "/tmp/myModel"
},
    "target": {
        "type": "Multiclass",
        "name": "Species",
        "classNames": [
            "Iris-versicolor",
            "Iris-virginica",
            "Iris-setosa"
        ]
    }
}

DeployModelPackageOperator

Create a deployment from a DataRobot model package.

Returns the created deployment ID.

Parameters:

AddExternalDatasetOperator

Upload a new dataset from the AI Catalog to make predictions for a model.

Returns an external dataset ID for the model,

Parameters:

For more upload_dataset_from_catalog parameters, see the DataRobot documentation.

RequestModelPredictionsOperator

Requests predictions against a previously uploaded dataset.

Returns a model predictions job ID.

Parameters:

For more request_predictions, see the DataRobot documentation.

TrainModelOperator

Submit a job to the queue to train a model from a specific blueprint.

Returns a model training job ID.

Parameters:

Example of DAG config params: { "sample_pct": "scoring_type": "training_row_count": "n_clusters": }

For more start-training-a-model, see the DataRobot documentation.

RetrainModelOperator

Submit a job to the queue to retrain a model on a specific sample size and/or custom feature list.

Returns a model retraining job ID.

Parameters:

Example of DAG config params: { "sample_pct": "scoring_type": "training_row_count": }

For more train-a-model-on-a-different-sample-size, see the DataRobot documentation.

PredictionExplanationsInitializationOperator

Initialize prediction explanations for a model.

Returns a prediction explanations initialization job ID.

Parameters:

For more prediction-explanations, see the DataRobot documentation.

ComputePredictionExplanationsOperator

Create prediction explanations for the specified dataset.

Returns a job ID for the prediction explanations for the specified dataset.

Parameters:

Example of DAG config params:

{ "max_explanations" "threshold_low" "threshold_high" }

For more prediction-explanations, see the DataRobot documentation.

Sensors

AutopilotCompleteSensor

Checks if Autopilot is complete.

Parameters:

ScoringCompleteSensor

Checks if batch scoring is complete.

Parameters:

MonitoringJobCompleteSensor

Checks if a monitoring job is complete.

Parameters:

BaseAsyncResolutionSensor

Checks if the DataRobot Async API call is complete.

Parameters:

DataRobotJobSensor

Checks whether a DataRobot job is complete.

Parameters:

ModelTrainingJobSensor

Checks whether a DataRobot model training job is complete.

Returns False if the job is not yet completed, and returns PokeReturnValue(True, trained_model.id) if model training has completed.

Parameters:

Hooks

DataRobotHook

A hook to initialize the DataRobot Public API client.

Pipeline

The modules described above allow you to construct a standard DataRobot pipeline in an Airflow DAG:

Copy
create_project_op >> train_models_op >> autopilot_complete_sensor >> deploy_model_op >> score_predictions_op >> scoring_complete_sensor

Example DAGS

See the the datarobot_provider/example_dags directory for the example DAGs.

You can find the following examples using a preconfigured connection in the datarobot_provider/example_dags directory:

The advanced end-to-end workflow in DataRobot (advanced_datarobot_pipeline_jdbc_dag.py) contains the following steps:

• Ingest a dataset to the AI Catalog from JDBC datasource
• Create a DataRobot project
• Train models using Autopilot
• Deploy the recommended model
• Change deployment settings (enable monitoring settings, segment analysis, and bias and fairness)
• Run batch scoring using a JDBC datasource
• Upload actuals from a JDBC datasource
• Collect deployment metrics: service statistics, features drift, target drift, accuracy and process it with custom python operator.

Issues

Please submit issues and pull requests in our official repo: https://github.com/datarobot/airflow-provider-datarobot

We are happy to hear from you. Please email any feedback to the authors at support@datarobot.com.

Copyright Notice

Copyright 2023 DataRobot, Inc. and its affiliates.

All rights reserved.

This is proprietary source code of DataRobot, Inc. and its affiliates.

Released under the terms of DataRobot Tool and Utility Agreement.