MNIST handwritten digit recognition model running on a local SageMaker endpoint

Key	Value
Environment
Services	S3, SageMaker, Lambda
Integrations	AWS SDK
Categories	Serverless, S3 website, Lambda function URLs, SageMaker, Machine Learning, JavaScript, Python
Level	Intermediate

Introduction

This is a sample application that demonstrates how to use SageMaker on LocalStack. A simple web frontend allows users to draw a digit and submit it to a locally running SageMaker endpoint. The endpoint returns a prediction of the digit, which is then displayed in the web frontend. Request handling is performed by a Lambda function, accessible via a function URL, that uses the SageMaker SDK to invoke the endpoint.

Here's a short summary of AWS service features we use:

• S3 website
• Lambda function URLs
• SageMaker endpoint

Here's the web application in action:

Screen.Recording.2023-04-27.at.16.04.43.mov

Architecture overview

Prerequisites

• A valid LocalStack for AWS license. Your license provides a LOCALSTACK_AUTH_TOKEN to activate LocalStack.
• lstk CLI.
• AWS CLI (required by lstk aws).

Dev environment

Create a virtualenv and install all the development dependencies there:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

If you'd like to perform training locally, you'll need to install the ml dev dependencies as well:

pip install -r ml/requirements.txt

You'll also need npm/node installed to build the web application. Please install according to official guidelines: https://github.com/nvm-sh/nvm

Download the PyTorch inference image

As our inference container, we use the AWS Deep Learning Container for PyTorch inference (2.5.1, CPU). It is several gigabytes, so pull it before deploying. AWS also publishes it on ECR Public, which needs no login:

make pull-image

This pulls public.ecr.aws/deep-learning-containers/pytorch-inference:2.5.1-cpu-py311-ubuntu22.04-sagemaker and tags it as 763104351884.dkr.ecr.eu-central-1.amazonaws.com/pytorch-inference:2.5.1-cpu-py311-ubuntu22.04-sagemaker, the image the model uses.

LocalStack

Start LocalStack with the LOCALSTACK_AUTH_TOKEN pre-configured:

export LOCALSTACK_AUTH_TOKEN=<your-auth-token>
make start

Instructions

Install the dependencies (lstk, boto3 and the web application's packages):

source .venv/bin/activate
make install

You can then create the AWS infrastructure on LocalStack by running the deploy/deploy_app.py script:

make deploy

This script will create the SageMaker endpoint with the model, which it first uploads to a bucket. The script will also create a lambda function that will be used to invoke the endpoint. Finally, the script will build the web application with the function URL and then create a s3 website to host it.

Using the application

Once deployed, visit http://mnist-website.s3-website.localhost.localstack.cloud:4566

Draw something in the canvas and click on the button that says Predict.

After a few moments the resulting prediction should be displayed in the box to the right.

Serverless SageMaker Endpoint

To switch to a serverless SageMaker endpoint you can also execute the deployment script with the additional -s or --serverless flag:

make deploy DEPLOY_ARGS=--serverless

Testing

make test waits for the endpoint to be InService, then predicts the digits in tests/fixtures (0, 3 and 5, drawn as 28x28 images like the web app sends them):

• by invoking the SageMaker endpoint directly
• through the Lambda function URL that the web app calls
• and checks the S3 website serves the web app, built with the function URL

The GitHub Actions workflow runs this for both the instance and the serverless endpoint.

License

The code of this sample application is published under the Apache 2.0 license (see LICENSE).

Contributing

We appreciate your interest in contributing to our project and are always looking for new ways to improve the developer experience. We welcome feedback, bug reports, and even feature ideas from the community. Please refer to the contributing file for more details on how to get started.