Simple MLOps with Amazon SageMaker, Lambda and AWS Step Functions Data Science SDK | by Stefan Natu | Towards Data Science

By Stefan Natu, Shreyas Subramanian, and Qingwei Li
As the machine learning space matures, there is an increasing need for simple ways to automate and deploy ML pipelines into production. With the explosion of data science platforms, companies and teams are often using diverse machine learning platforms for data exploration, Extract-Transform-Load (ETL) jobs, model training and deployment. In this blog, we describe how users can bring their own algorithm code to build a training and inference image using Docker, train and host their model using Amazon SageMaker and AWS StepFunctions. We use Mask R-CNN, a highly popular instance segmentation model used for numerous computer vision use cases as such as [1]. Readers who wish to get hands-on with the contents of this blog can refer to our Github [2].
AWS CodeBuild: AWS CodeBuild is a fully managed continuous integration (CI) service that allows users to compile and package code into deployable artifacts. Here we will use CodeBuild to package our custom Mask R-CNN container into a Docker image, which we upload into Amazon Elastic Container Registry (ECR)  AWS Lambda: AWS Lambda is a service that lets you run code without provisioning servers. Here we will use AWS Lambda to deploy a CodeBuild job.  AWS StepFunctions: AWS StepFuntions is an orchestration tool that allows users to build pipelines and coordinate microservices into a workflow. AWS StepFunctions can then trigger that workflow in an event driven manner without having the user manage any underlying servers or compute.  Amazon SageMaker: Amazon SageMaker is a fully managed machine learning platform for building, training and deploying machine learning models. Here we use Amazon SageMaker to author training and model deployment jobs, as well as SageMaker Jupyter notebooks to author a StepFunctions workflow.
Because we run the same image in training or hosting, Amazon SageMaker runs your container with the argument train or serve. When Amazon SageMaker runs training, your train script is run just like a regular Python program. Hosting has a very different model than training because hosting is responding to inference requests that come in via HTTP. In this example, we use our recommended Python serving stack to provide robust and scalable serving of inference requests:
In the `container` directory are all the components you need to package the sample algorithm for Amazon SageMager.
├── Dockerfile    ├── build_and_push.sh    └── mask_r_cnn        ├── nginx.conf        ├── predictor.py        ├── serve        ├── wsgi.py        ├── transforms.py        ├── utils.py        ├── coco_eval.py        ├── coco_utils.py        ├── engine.py        └── helper.py
Let's discuss each of these in turn:
Dockerfile describes how to build your Docker container image. More details below.
build_and_push.sh is a script that uses the Dockerfile to build your container images and then pushes it to ECR. We'll invoke the commands directly later in this notebook, but you can just copy and run the script for your own algorithms.
mask_r_cnn is the directory which contains the files that will be installed in the container.
In this simple application, we only install five files in the container. The files that we'll put in the container are:
nginx.conf is the configuration file for the nginx front-end. Generally, you should be able to take this file as-is.
predictor.py is the program that actually implements the Flask web server and the decision tree predictions for this app.
serve is the program started when the container is started for hosting. It simply launches the gunicorn server which runs multiple instances of the Flask app defined in predictor.py. You should be able to take this file as-is.
train is the program that is invoked when the container is run for training.
wsgi.py is a small wrapper used to invoke the Flask app. You should be able to take this file as-is.
We have customized train.py and predictor.py for fine tuning Mask R-CNN during training, and for loading tuned model, deserializing request data, making prediction, and sending back the serialized results.
The Dockerfile describes the image that we want to build. We will start from a standard Ubuntu installation and run the normal tools to install the things needed such as python, torch, torchvision, and Pillow. Finally, we add the code that implements our specific algorithm to the container and set up the right environment to run under. The Dockerfile looks like below,
FROM ubuntu:16.04MAINTAINER Amazon AI <sage-learner@amazon.com>RUN apt-get -y update && apt-get install -y --no-install-recommends \         wget \         gcc\         g++\         python3 \         python3-dev\         nginx \         ca-certificates \    && rm -rf /var/lib/apt/lists/*RUN wget https://bootstrap.pypa.io/get-pip.py && python3 get-pip.py && \    pip install cython numpy==1.16.2 scipy==1.2.1 pandas flask gevent gunicorn && \        (cd /usr/local/lib/python3.5/dist-packages/scipy/.libs; rm *; ln ../../numpy/.libs/* .) && \        rm -rf /root/.cacheRUN pip install  torch torchvision fastai thinc  PillowENV PYTHONUNBUFFERED=TRUEENV PYTHONDONTWRITEBYTECODE=TRUEENV PATH="/opt/program:${PATH}"# Set up the program in the imageCOPY mask_r_cnn /opt/programWORKDIR /opt/program
For building SageMaker-ready containers, such as the one discussed above for Mask R-CNN, we use the open source “SageMaker Containers” project which can be found here https://github.com/aws/sagemaker-containers. SageMaker Containers gives you tools to create SageMaker-compatible Docker containers, and has additional tools for letting you create Frameworks (SageMaker-compatible Docker containers that can run arbitrary Python or shell scripts). Currently, this library is used by the following containers: TensorFlow Script Mode, MXNet, PyTorch, Chainer, and Scikit-learn. To create a Sagemaker compatible container, we require the following components:
train.py file with your training code, andDockerfile, such as the one above
train.py file with your training code, and
Dockerfile, such as the one above
The training script must be located under the folder /opt/ml/code and its relative path is defined in the environment variable SAGEMAKER_PROGRAM. The following scripts are supported:
Python scripts: uses the Python interpreter for any script with .py suffix
Shell scripts: uses the Shell interpreter to execute any other script
When training starts, the interpreter executes the entry point, from the example above:
python train.py
For more information on hyper-parameters and environment variables, please refer to https://github.com/aws/sagemaker-containers#id10.
We will use a Cloud formation template to automate the container build of our Mask R-CNN. The template sets up the following architecture:
The Lambda function contains the train.py file and Dockerfile that can be edited inline. Once triggered (manually, or through a step functions approach as shown in the next section), the Lambda function:
Creates an ECR repository ,if it doesn’t already exist, to store the container images once builtUploads the train.py and Dockerfile to an S3 bucketCreates a Codebuild project and uses the above files with a buildspec.yml to start the process to build a container push the image to ECR.
Creates an ECR repository ,if it doesn’t already exist, to store the container images once built
Uploads the train.py and Dockerfile to an S3 bucket
Creates a Codebuild project and uses the above files with a buildspec.yml to start the process to build a container push the image to ECR.
The Lambda function also contains useful environment variables that can be reconfigured for new builds.
Once the Lambda function is set up, we are now ready to build out Automation pipeline to train and deploy the model to an endpoint. For this, we will use AWS Step Functions, which is an orchestration tool which lets users author state machines as JSON objects and execute them without provisioning or managing any servers. Step Functions also now provides a Data Science Python SDK for authoring machine learning pipelines using python in a familiar Jupyter notebook environment. We refer the reader to the Step Functions Github repository to get started [3].  Here we simply demonstrate the key components of the pipeline that are described in detail in our Github [2]. To author the code, we will use Amazon SageMaker, AWS fully managed machine learning platform.   As with all AWS Services, we first need to provide the appropriate service the IAM permissions to call other AWS Services. For this we first need to allow Amazon SageMaker to call Step Functions APIs, and AWS Step Functions to call SageMaker for model training, and endpoint creation and deployment. Detailed instructions on how to set up proper IAM credentials are described in [3].  Once this is set up, we will first create a Lambda state to run the Lambda function that takes the code and deploys it as a container to host in ECR.
lambda_state = LambdaStep(    state_id="Calls CodeBuild to Build Container",    parameters={          "FunctionName": "Docker_Lambda", #replace with the name of the Lambda function you created        "Payload": {             "input": "HelloWorld"        }    })lambda_state.add_retry(Retry(    error_equals=["States.TaskFailed"],    interval_seconds=15,    max_attempts=2,    backoff_rate=4.0))lambda_state.add_catch(Catch(    error_equals=["States.TaskFailed"],    next_step=Fail("LambdaTaskFailed")))
Retry and Catch steps here are added for error handling. You can modify these with your custom error handling steps, but this lets Step Functions know to end the workflow if the Lambda function fails to deploy the container.  The next steps are to chain together the training job and the model creation from the trained artifacts. Fortunately, Step Functions Data Science SDK provides the logic and APIs to chain these steps together, with any custom branching logic that could be required.
train_step = TrainingStep('Train Step',estimator=maskrcnn,data=os.path.dirname(data_location),job_name=execution_input['JobName'])model_step = ModelStep('Save model',model=train_step.get_expected_model(),model_name=execution_input['ModelName'])endpoint_config_step = EndpointConfigStep(    "Create Endpoint Config",    endpoint_config_name=execution_input['ModelName'],    model_name=execution_input['ModelName'],    initial_instance_count=1,    instance_type='ml.m5.large')endpoint_step = EndpointStep("Create Endpoint",endpoint_name=execution_input['EndpointName'],endpoint_config_name=execution_input['ModelName'])
Here we use our mask-rcnn estimator which is a general SageMaker Estimator object that lets us specify the type of instance we want to train our model on, identify any network or security settings if needed and model hyperparameters, as well as the output paths to the models. The input to the estimator is the container image that is created by our Lambda function above.
maskrcnn = sagemaker.estimator.Estimator(image,                       role, 1, 'ml.p2.xlarge', #feel free to modify with your own. A cost estimate is provided in Readme.                       output_path="s3://{}/{}/output".format(sess.default_bucket(), key),                       sagemaker_session=sess)maskrcnn.set_hyperparameters(num_epochs = 1,                              num_classes = 2)
By using the Chain utility, we can chain all the above steps together to occur sequentially. We can then choose to output the entire workflow as a JSON, that can be used in a much larger Cloud Formation template for example, which also includes information on the provisioning of instances, setting up of network security etc., or run on its own. By creating the workflow and rendering the graph, a state machine will be created in Amazon Step Functions console.
workflow_definition = Chain([    lambda_state,    train_step,    model_step,    endpoint_config_step,    endpoint_step])# Next, we define the workflowworkflow = Workflow(    name="MyWorkflow-BYOC-MaskRCNN-{}".format(uuid.uuid1().hex),    definition=workflow_definition,    role=workflow_execution_role)workflow.render_graph()
This renders the following output:
This completes our Step Function workflow. We can now execute this workflow directly in Amazon SageMaker  and check the progress using the workflow.execute and workflow.render_progress() APIs or from the Step Functions console directly. Executions will be logged in CloudWatch and can be used to send alerts and notifications in a downstream system to users.  With AWS Step Functions Data Science SDK, data scientists and engineers can seamlessly deploy custom containers, train ML models and deploy them into production. For integrating other AWS big data tools such as EMR and AWS Glue into Step Functions workflows, we refer the reader to [4, 5].
[1] He, K., Gkioxari, G., Dollar, P., and Girshick, R., Mask R-CNN, https://arxiv.org/abs/1703.06870. [2] https://github.com/aws-samples/aws-stepfunctions-byoc-mlops-using-data-science-sdk [3] https://github.com/awslabs/amazon-sagemaker-examples/tree/master/step-functions-data-science-sdk [4]https://medium.com/@elesin.olalekan/automating-machine-learning-workflows-with-aws-glue-sagemaker-and-aws-step-functions-data-science-b4ed59e4d7f9 [5]https://aws.amazon.com/blogs/aws/new-using-step-functions-to-orchestrate-amazon-emr-workloads/