Airflow Retrain

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

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.

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

# 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.

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.

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

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.