Streamline custom environment provisioning for Amazon SageMaker Studio: An automated CI/CD pipeline approach

Attaching a custom Docker image to your Amazon SageMaker Studio domain requires several steps. First, you need to build and push an image to Amazon Elastic Container Registry (Amazon ECR). You also need to ensure that your Amazon SageMaker domain execution role has the necessary permissions to pull images from Amazon ECR. After the image is pushed to Amazon ECR, create a SageMaker custom image in the AWS Management Console. Finally, update your SageMaker domain settings to specify the Amazon Resource Name (ARN) for your custom image. This multi-step process must be performed manually every time an end user wants to create a new custom Docker image and make it available in SageMaker Studio.

This post explains how to automate this process. This approach allows you to update SageMaker configurations, provision custom images, and attach them to your SageMaker domain without writing any additional infrastructure code. By implementing this automation, you can implement a consistent and standardized analysis environment across your organization, increasing team productivity and reducing the security risks associated with using one-time images.

The solution described in this post is intended for machine learning (ML) engineers and platform teams responsible for managing and standardizing large custom environments across organizations. For individual data scientists looking for a self-service experience, we recommend using SageMaker Studio’s native Docker support, as described in Accelerating ML Workflows with Amazon SageMaker Studio Local Mode and Docker Support. Masu. This feature allows data scientists to build, test, and deploy custom Docker containers directly within the SageMaker Studio integrated development environment (IDE), allowing them to iterate and experiment with their analytical environments seamlessly within the familiar SageMaker Studio interface.

Solution overview

The following diagram shows the solution architecture.

Deploy your pipeline using AWS CodePipeline. This automates the creation of custom Docker images and attaching them to your SageMaker domain. The pipeline first checks out the code base from a GitHub repository and creates a custom Docker image based on the configuration declared in the configuration file. After successfully creating and pushing a Docker image to Amazon ECR, the pipeline validates the image by scanning and checking for security vulnerabilities in the image. If no critical or high-security vulnerabilities are found, the pipeline advances to the pre-deployment manual approval stage. Once manual approval is complete, the pipeline deploys your SageMaker domain and automatically attaches your custom image to the domain.

Prerequisites

The prerequisites for implementing the solution described in this post are:

Deploy the solution

To implement the solution, follow these steps:

1. Log in to your AWS account using the AWS CLI in a shell terminal (for more information, see Authenticating with AWS CLI Short-Term Credentials).
2. Verify that you have successfully logged in to your AWS account by running the following command.

aws sts get-caller-identity

4. Fork your GitHub repository into your GitHub account.
5. Clone the forked repository to your local workstation using the following command:

git clone <clone_url_of_forked_repo>

6. Log in to the console and create an AWS CodeStar connection to your GitHub repository in the previous step. For instructions, see Create a connection to GitHub (console).
7. Copy the ARN of the connection you created.
8. Go to Terminal and run the following command to change to the repository directory.

cd streamline-sagemaker-custom-images-cicd

9. Install all libraries from npm by running the following command:

10. Run the shell script in Terminal by running the following command: This script takes an AWS account number and AWS Region as input parameters and deploys an AWS CDK stack that deploys components such as CodePipeline, AWS CodeBuild, and an ECR repository. Set the following VPC_ID export variable using an existing VPC. If you don’t have a VPC, create and use one with at least two subnets.

export AWS_ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
export AWS_REGION=<YOUR_AWS_REGION>
export VPC_ID=<VPC_ID_TO_DEPLOY>
export CODESTAR_CONNECTION_ARN=<CODE_STAR_CONNECTION_ARN_CREATED_IN_ABOVE_STEP>
export REPOSITORY_OWNER=<YOUR_GITHUB_LOGIN_ID>

11. Run the following command to deploy AWS infrastructure using AWS CDK V2 and wait for the template to succeed.

cdk deploy PipelineStack --require-approval never

12. In the CodePipeline console, choose: pipeline in the navigation pane.
13. Select the pipeline link named sagemaker-custom-image-pipeline.

14. You can track the progress of your pipeline in the console and provide approval in the manual approval stage to deploy your SageMaker infrastructure. The pipeline takes approximately 5-8 minutes to build the image and move to the manual approval stage
15. Wait until the pipeline completes the deployment stage.

The pipeline creates infrastructure resources in your AWS account using a SageMaker domain and a SageMaker custom image. Also attach your custom image to your SageMaker domain.

16. In the SageMaker console, select: domain under Administrator configuration in the navigation pane.

17. Open the domain named “team-ds” and navigate to: environment

You should see one custom image attached.

How to deploy and connect custom images

CodePipeline has a stage called . BuildCustomImages It includes automated steps to create a SageMaker custom image using the SageMaker Custom Image CLI and push it to an ECR repository created in your AWS account. The AWS CDK stack for the deployment stage includes the steps necessary to create a SageMaker domain and attach a custom image to the domain. Parameters to create SageMaker domains, custom images, etc. are configured in JSON format and are used by the SageMaker stack under the lib directory. Please refer to sagemakerConfig section of environments/config.json For declared parameters.

Add more custom images

You can now add your own custom Docker image and attach it to the SageMaker domain created by your pipeline. For custom images that are created, see Dockerfile specifications for Docker image specifications.

1. cd to the repository’s images directory in your terminal.

2. Create a new directory (for example, custom) under the images directory.

3. Add your own Dockerfile to this directory. You can use the following Dockerfile configuration for testing.

FROM public.ecr.aws/amazonlinux/amazonlinux:2
ARG NB_USER="sagemaker-user"
ARG NB_UID="1000"
ARG NB_GID="100"
RUN yum update -y && \
    yum install python3 python3-pip shadow-utils -y && \
    yum clean all
RUN yum install --assumeyes python3 shadow-utils && \
    useradd --create-home --shell /bin/bash --gid "${NB_GID}" --uid ${NB_UID} ${NB_USER} && \
    yum clean all && \
    python3 -m pip install jupyterlab
RUN python3 -m pip install --upgrade pip
RUN python3 -m pip install --upgrade urllib3==1.26.6
USER ${NB_UID}
CMD jupyter lab --ip 0.0.0.0 --port 8888 \
--ServerApp.base_url="/jupyterlab/default" \
--ServerApp.token='' \
--ServerApp.allow_origin='*'

4. Update the images section of the json file under your environment directory to include the new image directory name you created.

"images": (
      "repositoryName": "research-platform-ecr",
       "tags":(
         "jlab",
         "custom" << Add here
       )
      }
    )

5. Update same image name customImages Under the created SageMaker domain configuration:

"customImages":(
          "jlab",
          "custom" << Add here
 ),

6. Commit and push your changes to your GitHub repository.
7. You can see that CodePipeline is triggered on push. Track pipeline progress and manually approve deployments.

Once the deployment completes successfully, you should be able to see the custom image you added attached to your domain configuration (see the following screenshot).

cleaning

To clean up your resources, open the AWS CloudFormation console and delete your stack. SagemakerImageStack and PipelineStack In that order. If you encounter errors such as “The S3 bucket is not empty” or “There are images in the ECR repository”, you can manually delete the S3 bucket and ECR repository that were created. Then you can try deleting the CloudFormation stack again.

conclusion

In this post, you learned how to create an automated continuous integration and delivery (CI/CD) pipeline solution to build, scan, and deploy custom Docker images to your SageMaker Studio domain. This solution helps drive consistency in the analytical environment for data science teams across the enterprise. This approach helps achieve machine learning (ML) governance, scalability, and standardization.

About the author

muni anachiis a senior DevOps consultant at AWS with over 10 years of expertise in designing and implementing software systems and cloud platforms. He specializes in guiding nonprofit organizations to adopt DevOps CI/CD architectures while adhering to AWS best practices and the AWS Well-Architected Framework. Apart from his professional endeavors, Muni is an avid sports enthusiast and tries his luck in the kitchen.

Ajay Raghunathan I’m a machine learning engineer at AWS. His current work focuses on designing and implementing large-scale ML solutions. He is a technology enthusiast and builder with core interests: AI/ML, data analytics, serverless, and DevOps. Outside of work, I enjoy spending time with my family, traveling, and playing soccer.

Arun Diasani I’m a Senior Cloud Application Architect at AWS. His current work focuses on designing and implementing innovative software solutions. His role focuses on leveraging his deep knowledge and experience in developing large systems to build robust architectures for complex applications.

Shweta Singh He is a senior product manager on the Amazon SageMaker Machine Learning platform team at AWS, where he leads the SageMaker Python SDK. She has held several product roles at Amazon for over five years. She holds a Bachelor of Science in Computer Engineering and a Master of Science in Financial Engineering from New York University.

jenna ung is a Principal Practice Manager on the Health and Advanced Computing team in AWS Professional Services. Her team focuses on designing and delivering data, ML, and advanced computing solutions for the public sector, including federal, state, and local governments, academic medical centers, nonprofit healthcare organizations, and research institutions.

Meenakshi Pon Shankaran I am a Principal Domain Architect for the Data & ML Professional Services organization at AWS. He has extensive expertise in designing and building large-scale data lakes processing petabytes of data. Currently, his focus is on providing technical leadership to AWS US public sector clients, guiding them to use innovative AWS services to achieve their strategic goals and unlock the full potential of their data. I’m leaving it there.