Airflow Retrain

AWS Example

The following example shows how to kick off an Airflow retrain of a model from AWS.

This setup uses email based triggers based on feedback from security teams. The goal of this approach was to kick off retraining events without requiring direct network IP connectivity between Arize and the customers AWS account.

Arize to AWS event flow

The flow of events on a model retrain:

1. Email generated on monitoring firing

2. SES receives the monitoring email

3. SES triggers lambda function

4. Lambda function contains all logic to kick off airflow retrain

SES (Simple Email Service) Identity

This section walks through setting up the SES service.

Setup an SES Identity, this verifies the ownership of a domain or email address such that Amazon SES can receive emails sent to that address.

Create Identity

Once the identity is setup you can setup a rule to receive emails.

Create Lambda Function for Airflow Retrain

This section walks through creating a lambda function that we will reference in the email reception call.

This example uses a Python Lambda function.

The code of the Lambda function will follow this example:

Invoking DAGs with an AWS Lambda function - Amazon Managed Workflows for Apache AirflowAmazon Managed Workflows for Apache Airflow

Example Kicking off Aiflow in Lambda

# Airflow V2
import boto3
import http.client
import base64
import ast
mwaa_env_name = 'YOUR_ENVIRONMENT_NAME'
dag_name = 'YOUR_DAG_NAME'
mwaa_cli_command = 'dags trigger'
​
client = boto3.client('mwaa')
​
def lambda_handler(event, context):
    # get web token
    mwaa_cli_token = client.create_cli_token(
        Name=mwaa_env_name
    )
    
    conn = http.client.HTTPSConnection(mwaa_cli_token['WebServerHostname'])
    payload = "dags trigger " + dag_name
    headers = {
      'Authorization': 'Bearer ' + mwaa_cli_token['CliToken'],
      'Content-Type': 'text/plain'
    }
    conn.request("POST", "/aws_mwaa/cli", payload, headers)
    res = conn.getresponse()
    data = res.read()
    dict_str = data.decode("UTF-8")
    mydata = ast.literal_eval(dict_str)
    return base64.b64decode(mydata['stdout'])

The above code shows an example of kicking off an Airflow job using a labmda function.

SES (Simple Email Service) Rule

Lastly we are going to create a rule in our SES service that references the lambda function and calls it when the email is received.

In the email receiving section create a rule set and then create a rule.

Create Rule

Name the rule that will call the lambda function when the email is received.

Choose an email address to use in the Arize platform in the recipient section of the rule, emails to this address will execute the rule.

Lambda Reference

The AWS Lambda function reference should be chosen in the action section.

On initial setup when you are debugging, it can be useful to setup a 2nd action with "Deliver to S3 bucket" so you can review received messages as well.

Lambda-airflow-retrain

In the lambda function reference the function name we created in the previous section - lambda-airflow-retrain.

Create the rule.

Finally, the rule does need to be set as active to work, go back to email receiving at set rule as active.

Use Lambda-Airflow Email Address in Arize

Use the email address created for the rule to send for key retrain alerts.