Most Machine Learning models are dynamic. They continuously learn and enhance their performance with additional data. This requires us to constantly update the deployed model.

As we know, deployment is a repetitive process. And if it's repetitive - we’ll automate it. This is where CI/CD pipelines come into play, streamlining the process effectively.

However, doing it for ML models is not as intuitive as you imagine. ML models are fundamentally different from code-based applications (size, type, framework, etc.), and therefore, they have unique requirements for deployment.

Let’s explore how the same tools that helped us in building a continuous training pipeline - Amazon SageMaker, Dagshub, and MLFlow - can help us in Deploying a model.

In this tutorial, we’ll use the image-segmentation model we created in our Kvasir Image Segmentation project and deploy it.

What will this blog cover?

In this tutorial, we'll go through the process of creating a custom deployment automation pipeline with SageMaker, including:

1. Code Structure: We'll create a simple code setup to deploy a real-time inference Endpoint infrastructure, utilizing a CloudFormation template.
2. Stage Configurations: We'll define configuration files for smooth deployment across staging and production stages, ensuring consistency.

How to set up SageMaker for Continuous Deployment?

Project Setup

Projects require a certain structure of files so that they can be sequentially executed by CodePipeline. I’ve created a blog on how to Setup SageMaker for Machine Learning CI/CD pipelines that explains in detail everything you’ll need to set up before starting the deployment process.

Setup Credentials

For a detailed explanation of how to set up IAM credentials take a look at the Setup SageMaker for Machine Learning CI/CD pipelines blog.

SageMaker and CodePipeline are going to be building an environment on your behalf, so we’ll need a service role. If you create the role from the Setup blog, you should already have it set up.

If not, find this role in the IAM console: AmazonSageMakerServiceCatalogProductsUseRole

and edit it with the following policy:

{
            "Effect": "Allow",
            "Action": [
                "s3:CreateBucket",
                "s3:DeleteBucket",
                "s3:GetBucketAcl",
                "s3:GetBucketCors",
                "s3:GetBucketLocation",
                "s3:PutBucketCors",
                "s3:PutObjectTagging",
                "s3:PutObjectAcl"
            ],
            "Resource": [
                "arn:aws:s3:::aws-glue-*",
                "arn:aws:s3:::sagemaker-*",
                "arn:aws:s3:::<your-bucket>/*"
            ]
        },
        {
            "Effect": "Allow",
            "Action": [
                "s3:AbortMultipartUpload",
                "s3:PutObject",
                "s3:PutObjectAcl",
                "s3:GetObject",
                "s3:GetObjectAcl",
                "s3:DeleteObject",
                "s3:GetObjectVersion"
            ],
            "Resource": [
                "arn:aws:s3:::aws-glue-*",
                "arn:aws:s3:::sagemaker-*",
				"arn:aws:s3:::<your-bucket>/*"
            ]
        },

Overview of the Project Structure

The project contains the following files:

• buildspec.yml - Used by CodePipeline’s Build stage to build a CloudFormation template
• build.py - A Python file that contains code to deploy the latest machine learning model and exports staging configuration files. This is invoked from the build stage.
• endpoint-config-template.yml - This CloudFormation template file is packaged by the build step in the CodePipeline and is deployed in different stages.
• staging-config.json/prod-config.json - These files are used to customize the staging/prod stage in the pipeline.

How to Deploy A Model Using MLFlow and SageMaker?

Step 1: Create a Container Image

For a SageMaker Endpoint to function, it needs a container image to manage inferences (predictions). MLFlow offers a way to create and send this container to Amazon's Elastic Container Registry (ECR). This might involve some steps on your computer.

💡Note: Your AWS configuration file is stored in the .aws/ folder on your computer. If you're unsure how to set it up, check out this guide. Make sure the region in the configuration matches your project's region on AWS.

#In your local terminal 

pip install -q mlflow==2.5.0
mlflow sagemaker build-and-push-container

If everything goes well, you'll see a folder named 'mlflow-pyfunc' in AWS ECR, with the version as a sub-folder.

Step 2: Configure Deployment Client

MLFlow provides a handy way to deploy its registered models to Amazon SageMaker using an API. But before deploying, we need to set up MLFlow to access the model. This process is made easier with the help of Dagshub.

DagsHub offers a built-in remote MLFlow server for each repository. This server lets you keep track of experiments and trained models. You can find the specific MLflow configurations for your repository in the remote section of your repo.

Configuring MLFlow involves adding certain credentials. In thebuild.pyfile, the environment variables at the top part do exactly this:

# Defining MLFlow credentials
os.environ['MLFLOW_TRACKING_USERNAME']=args.mlflow_username
os.environ['MLFLOW_TRACKING_PASSWORD']=args.mlflow_token

This connection allows the MLFlow server associated with your repository to write data using the provided username and token.

Once MLFlow is ready, the next step is to create an endpoint. This requires a configuration dictionary with essential details that CodePipeline needs to build an endpoint:

1. Execution Role: The role you set up earlier in "Setup Credentials".
2. Image URI: The ARN of the container image you created in Step 1.
3. Resources: This includes instance type and count.

config = dict(
		# 1
    execution_role_arn=args.model_execution_role,
    bucket_name=args.s3_bucket,
		# 2
    image_url=args.container_image_uri,
    region_name="us-east-2",
    archive=False,
		# 3
    instance_type="ml.m5.4xlarge",
    instance_count=1,
    synchronous=True,
    timeout_seconds=300,
   )

Next, we’ll create the deployment client. Provide the client with a name for your endpoint, the location of your model, and the above configurations:

client = get_deploy_client("sagemaker")
    
client.create_deployment(
  "image-segmentation-test",
  model_uri="runs:/8449085500f84750a3d4590a2b149845/model",
  flavor="python_function",
  config=config
)

Step 3: Update the Automation Script

And now, for the final step, let's configure all the essential parameters in our buildspec.yml file. This file will serve as your go-to configuration.

• Pre-Build Phase: This section includes parameters linked to the model you're deploying. You'll need to customize these settings to initiate a fresh deployment.
• Build Phase: This is where you execute your build script, utilizing the configurations established in the pre-build phase.

CodePipeline is set up in such a way that any changes made to the production (main) branch automatically trigger the pipeline's execution. Once you've registered your model, update the automation script with its version and push the updates. This will initiate the deployment pipeline.

You can keep tabs on the process right in CodePipeline. Once the pipeline successfully completes, you'll find the endpoint in SageMaker under the Inference tab > Endpoints. The endpoint will carry the name you assigned to it.

Go ahead and test out your endpoint, and see if it makes predictions correctly.

Conclusion

Phew!! We did it ! We build an ML model and deployed it using SageMaker!! You can now write test scripts for testing out predictions, or if you’re confident build a web-app to upload images and get predictions. Endless opportunities await, go forth and explore.

And if you do explore, we’d love to know! Share your projects, ideas and thoughts with us. If you have any questions, feel free to reach out. You can join our Discord, where we’ve built a vibrant, helpful, and friendly community.