Manage Serverless Machine Learning Workflows with Amazon Step Functions with the example of Email Campaigns 💌
By: Michael Triska
Tags: Serverless, Step Functions, Amaro, Data Engineering, Data Science

AWS Step Functions for Machine Learning Pipeline Automation

“If you aren’t careful you can end up with data scientists literally emailing Python notebooks and models to engineers for production deployment.” — Kyle Gallatin

In order to scale fast, reduce the maintenance cost of production machine learning pipelines, improve engineering productivity, and increase the experimentation rate, the data team at AMARO uses AWS Step Functions for machine learning pipeline automation. AWS Step Functions is an orchestration service that allows you to build resilient serverless workflows at scale, integrating multiple AWS services. Combined with AWS Personalize this gives you a powerful service that orchestrates an automated pipeline to pre-process data, then train and publish API endpoints.

Image 1: Step Function Machine Learning Workflow.

As we can see in the following architecture diagram, we deployed an AWS Step Functions Workflow that synchronously calls a Fargate task and contains AWS Lambda functions to call Amazon Personalize or Slack. Workflows are made up of a series of steps, with the output of one step acting as the input into the next. The high-level process looks like this:

1. Trigger and Input Processing. The Step Function execution receives a JSON string as input from an event trigger and passes that input over to the first state in the workflow. The input string flows from state to state and depending on the input parameter, we can create different models backing a campaign. In this way, the workflow becomes scalable and generic enough for a variety of use cases even outside of email campaigns. For example, if you change the strategy input, the backend code will use a different dataset or algorithm parameter.

2. Preprocessing, Dataset Creation and Testing. We first use Step Functions to execute an ETL job with ECS Fargate container task in the step “Unload and Upload Dataset”. Amazon Personalize can make recommendations based purely on historical user interaction data as well as real-time event data. For automated email campaigns, we use historical data.

In the ETL pipeline, we create a user interaction data set from click events and eliminate data debt, through testing for any obvious issues before unloading data from our data warehouse to S3. Such debts include heavy-tail or head-heavy user and item distributions as well as repeated user-item activities or duplication rates. Personalize uses the concept of “Dataset Groups” to isolate your dataset, schema training model, and campaign endpoints from different use cases. To reduce the complexity of the Step Function State Machine, we create the dataset resources and import the dataset from S3 to Personalize within the container task rather than with additional Step Function states.

3. Create a Solution and Campaign Endpoint. We use AWS Lambda functions to map the rest of the workflow that you would normally do manually in the AWS interface to select a recommendation algorithm, train a model, extract experiment results or create the campaign endpoint.

And Voilà. You have successfully trained a Personalize model and now can get recommendations for your users leveraging your campaign endpoint.