Arize AI
Search…
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. 1.
    Email generated on monitoring firing
  2. 2.
    SES receives the monitoring email
  3. 3.
    SES triggers lambda function
  4. 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 Airflow
Amazon Managed Workflows for Apache Airflow
Example Kicking off Aiflow in Lambda
1
# Airflow V2
2
import boto3
3
import http.client
4
import base64
5
import ast
6
mwaa_env_name = 'YOUR_ENVIRONMENT_NAME'
7
dag_name = 'YOUR_DAG_NAME'
8
mwaa_cli_command = 'dags trigger'
9
10
client = boto3.client('mwaa')
11
12
def lambda_handler(event, context):
13
# get web token
14
mwaa_cli_token = client.create_cli_token(
15
Name=mwaa_env_name
16
)
17
18
conn = http.client.HTTPSConnection(mwaa_cli_token['WebServerHostname'])
19
payload = "dags trigger " + dag_name
20
headers = {
21
'Authorization': 'Bearer ' + mwaa_cli_token['CliToken'],
22
'Content-Type': 'text/plain'
23
}
24
conn.request("POST", "/aws_mwaa/cli", payload, headers)
25
res = conn.getresponse()
26
data = res.read()
27
dict_str = data.decode("UTF-8")
28
mydata = ast.literal_eval(dict_str)
29
return base64.b64decode(mydata['stdout'])
Copied!
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.