Multi-account model deployment with Amazon SageMaker Pipelines

by
Amazon AWS

Amazon SageMaker Pipelines is the first purpose-built CI/CD service for machine learning (ML). It helps you build, automate, manage, and scale end-to-end ML workflows and apply DevOps best practices of CI/CD to ML (also known as MLOps).

Creating multiple accounts to organize all the resources of your organization is a good DevOps practice. A multi-account strategy is important not only to improve governance but also to increase security and control of the resources that support your organization’s business. This strategy allows many different teams inside your organization, to experiment, innovate, and integrate faster, while keeping the production environment safe and available for your customers.

Pipelines makes it easy to apply the same strategy to deploying ML models. Imagine a use case in which you have three different AWS accounts, one for each environment: data science, staging, and production. The data scientist has the freedom to run experiments and train and optimize different models any time in their own account. When a model is good enough to be deployed in production, the data scientist just needs to flip the model approval status to Approved. After that, an automated process deploys the model on the staging account. Here you can automate testing of the model with unit tests or integration tests or test the model manually. After a manual or automated approval, the model is deployed to the production account, which is a more tightly controlled environment used to serve inferences on real-world data. With Pipelines, you can implement a ready-to-use multi-account environment.

In this post, you learn how to use Pipelines to implement your own multi-account ML pipeline. First, you learn how to configure your environment and prepare it to use a predefined template as a SageMaker project for training and deploying a model in two different accounts: staging and production. Then, you see in detail how this custom template was created and how to create and customize templates for your own SageMaker projects.

Preparing the environment

In this section, you configure three different AWS accounts and use SageMaker Studio to create a project that integrates a CI/CD pipeline with the ML pipeline created by a data scientist. The following diagram shows the reference architecture of the environment that is created by the SageMaker custom project and how AWS Organizations integrates the different accounts.

The following diagram shows the reference architecture of the environment that is created by the SageMaker custom project and how AWS Organizations integrates the different accounts.

The diagram contains three different accounts, managed by Organizations. Also, three different user roles (which may be the same person) operate this environment:

  • ML engineer – Responsible for provisioning the SageMaker Studio project that creates the CI/CD pipeline, model registry, and other resources
  • Data scientist – Responsible for creating the ML pipeline that ends with a trained model registered to the model group (also referred to as model package group)
  • Approver – Responsible for testing the model deployed to the staging account and approving the production deployment

It’s possible to run a similar solution without Organizations, if you prefer (although not recommended). But you need to prepare the permissions and the trust relationship between your accounts manually and modify the template to remove the Organizations dependency. Also, if you’re an enterprise with multiple AWS accounts and teams, it’s highly recommended that you use AWS Control Tower for provisioning the accounts and Organizations. AWS Control Tower provides the easiest way to set up and govern a new and secure multi-account AWS environment. For this post, we only discuss implementing the solution with Organizations.

But before you move on, you need to complete the following steps, which are detailed in the next sections:

  1. Create an AWS account to be used by the data scientists (data science account).
  2. Create and configure a SageMaker Studio domain in the data science account.
  3. Create two additional accounts for production and staging.
  4. Create an organizational structure using Organizations, then invite and integrate the additional accounts.
  5. Configure the permissions required to run the pipelines and deploy models on external accounts.
  6. Import the SageMaker project template for deploying models in multiple accounts and make it available for SageMaker Studio.

Configuring SageMaker Studio in your account

Pipelines provides built-in support for MLOps templates to make it easy for you to use CI/CD for your ML projects. These MLOps templates are defined as Amazon CloudFormation templates and published via AWS Service Catalog. These are made available to data scientists via SageMaker Studio, an IDE for ML. To configure Studio in your account, complete the following steps:

  1. Prepare your SageMaker Studio domain.
  2. Enable SageMaker project templates and SageMaker JumpStart for this account and Studio users.

If you have an existing domain, you can simply edit the settings for the domain or individual users to enable this option. Enabling this option creates two different AWS Identity and Account Management (IAM) roles in your AWS account:

  • AmazonSageMakerServiceCatalogProductsLaunchRole – Used by SageMaker to run the project templates and create the required infrastructure resources
  • AmazonSageMakerServiceCatalogProductsUseRole – Used by the CI/CD pipeline to run a job and deploy the models on the target accounts

If you created your SageMaker Studio domain before re:Invent 2020, it’s recommended that you refresh your environment by saving all the work in progress. On the File menu, choose Shutdown, and confirm your choice.

  1. Create and prepare two other AWS accounts for staging and production, if you don’t have them yet.

Configuring Organizations

You need to add the data science account and the two additional accounts to a structure in Organizations. Organizations helps you to centrally manage and govern your environment as you grow and scale your AWS resources. It’s free and benefits your governance strategy.

Each account must be added to a different organizational unit (OU).

  1. On the Organizations console, create a structure of OUs like the following:
  • Root
    • multi-account-deployment (OU)
      • 111111111111 (data science account—SageMaker Studio)
      • production (OU)
        • 222222222222 (AWS account)
      • staging (OU)
        • 333333333333 (AWS account)

After configuring the organization, each account owner receives an invite. The owners need to accept the invites, otherwise the accounts aren’t included in the organization.

  1. Now you need to enable trusted access with AWS organizations (“Enable all features” and “Enable trusted access in the StackSets”).

This process allows your data science account to provision resources in the target accounts. If you don’t do that, the deployment process fails. Also, this feature set is the preferred way to work with Organizations, and it includes consolidating billing features.

  1. Next, on the Organizations console, choose Organize accounts.
  2. Choose staging.
  3. Note down the OU ID.
  4. Repeat this process for the production OU.

Repeat this process for the production OU.

Configuring the permissions

You need to create a SageMaker execution role in each additional account. These roles are assumed by AmazonSageMakerServiceCatalogProductsUseRole in the data science account to deploy the endpoints in the target accounts and test them.

  1. Sign in to the AWS Management Console with the staging account.
  2. Run the following CloudFormation template.

This template creates a new SageMaker role for you.

  1. Provide the following parameters:
    1. SageMakerRoleSuffix – A short string (maximum 10 lowercase with no spaces or alpha-numeric characters) that is added to the role name after the following prefix: sagemaker-role-. The final role name is sagemaker-role-<<sagemaker_role_suffix>>.
    2. PipelineExecutionRoleArn – The ARN of the role from the data science account that assumes the SageMaker role you’re creating. To find the ARN, sign in to the console with the data science account. On the IAM console, choose Roles and search for AmazonSageMakerServiceCatalogProductsUseRole. Choose this role and copy the ARN (arn:aws:iam::<<data_science_acccount_id>>:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole).
  2. After creating this role in the staging account, repeat this process for the production account.

In the data science account, you now configure the policy of the Amazon Simple Storage Service (Amazon S3) bucket used to store the trained model. For this post, we use the default SageMaker bucket of the current Region. It has the following name format: sagemaker-<<region>>-<<aws_account_id>>.

  1. On the Amazon S3 console, search for this bucket, providing the Region you’re using and the ID of the data science account.

If you don’t find it, create a new bucket following this name format.

  1. On the Permissions tab, add the following policy: 
    {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": [
                    "arn:aws:iam::<<staging_account_id>>:root",
                    "arn:aws:iam::<<production_account_id>>:root"
                ]
            },
            "Action": [
                "s3:GetObject",
                "s3:ListBucket"
            ],
            "Resource": [
                "arn:aws:s3:::sagemaker-<<region>>-<<aws_account_id>>",
                "arn:aws:s3:::sagemaker-<<region>>-<<aws_account_id>>/*"
            ]
        }
    ]
}

  1. Save your settings.

The target accounts now have permission to read the trained model during deployment.

The next step is to add new permissions to the roles AmazonSageMakerServiceCatalogProductsUseRole and AmazonSageMakerServiceCatalogProductsLaunchRole.

  1. In the data science account, on the IAM console, choose Roles.
  2. Find the AmazonSageMakerServiceCatalogProductsUseRole role and choose it.
  3. Add a new policy and enter the following JSON code.
  4. Save your changes.
  5. Now, find the AmazonSageMakerServiceCatalogProductsLaunchRole role, choose it and add a new policy with the following content: 
    {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "VisualEditor0",
            "Effect": "Allow",
            "Action": "s3:GetObject",
            "Resource": "arn:aws:s3:::aws-ml-blog/artifacts/sagemaker-pipeline-blog-resources/*"
        }
    ]
}

  1. Save your changes.

That’s it! Your environment is almost ready. You only need one more step and you can start training and deploying models in different accounts.

Importing the custom SageMaker Studio project template

In this step, you import your custom project template.

  1. Sign in to the console with the data science account.
  2. On the AWS Service Catalog console, under Administration, choose Portfolios.
  3. Choose Create a new portfolio.
  4. Name the portfolio SageMaker Organization Templates.
  5. Download the following template to your computer.
  6. Choose the new portfolio.
  7. Choose Upload a new product.
  8. For Product name¸ enter Multi Account Deployment.
  9. For Description, enter Multi account deployment project.
  10. For Owner, enter your name.
  11. Under Version details, for Method, choose Use a template file.
  12. Choose Upload a template.
  13. Upload the template you downloaded.
  14. For Version title, choose 1.0.

The remaining parameters are optional.

  1. Choose Review.
  2. Review your settings and choose Create product.
  3. Choose Refresh to list the new product.
  4. Choose the product you just created.
  5. On the Tags tab, add the following tag to the product:
    1. Key – sagemaker:studio-visibility
    2. Value – True

Back in the portfolio details, you see something similar to the following screenshot (with different IDs).

Back in the portfolio details, you see something similar to the following screenshot (with different IDs).

  1. On the Constraints tab, choose Create constraint.
  2. For Product, choose Multi Account Deployment (the product you just created).
  3. For Constraint type, choose Launch.
  4. Under Launch Constraint, for Method, choose Select IAM role.
  5. Choose AmazonSageMakerServiceCatalogProductsLaunchRole.
  6. Choose Create.
  7. On the Groups, roles, and users tab, choose Add groups, roles, users.
  8. On the Roles tab, select the role you used when configuring your SageMaker Studio domain.
  9. Choose Add access.

If you don’t remember which role you selected, in your data science account, go to the SageMaker console and choose Amazon SageMaker Studio. In the Studio Summary section, locate the attribute Execution role. Search for the name of this role in the previous step.

You’re done! Now it’s time to create a project using this template.

Creating your project

In the previous sections, you prepared the multi-account environment. The next step is to create a project using your new template.

  1. Sign in to the console with the data science account.
  2. On the SageMaker console, open SageMaker Studio with your user.
  3. Choose the Components and registries
  4. On the drop-down menu, choose Projects.
  5. Choose Create project.

Choose Create project.

On the Create project page, SageMaker templates is chosen by default. This option lists the built-in templates. However, you want to use the template you prepared for the multi-account deployment.

  1. Choose Organization templates.
  2. Choose Multi Account Deployment.
  3. Choose Select project template.

If you can’t see the template, make sure you completed all the steps correctly in the previous section.

If you can’t see the template, make sure you completed all the steps correctly in the previous section.

  1. In the Project details section, for Name, enter iris-multi-01.

The project name must have 15 characters or fewer.

  1. In the Project template parameters, use the names of the roles you created in each target account (staging and production) and provide the following properties:
    1. SageMakerExecutionRoleStagingName
    2. SageMakerExecutionRoleProdName
  2. Retrieve the OU IDs you created earlier for the staging and production OUs and provide the following properties:
    1. OrganizationalUnitStagingId
    2. OrganizationalUnitProdId
  3. Choose Create project.

Choose Create project.

Provisioning all the resources takes a few minutes, after which the project is listed in the Projects section. When you choose the project, a tab opens with the project’s metadata. The Model groups tab chows a model group with the same name as your project. It was also created during the project provisioning.

Provisioning all the resources takes a few minutes, after which the project is listed in the Projects section.

The environment is now ready for the data scientist to start training the model.

Training a model

Now that your project is ready, it’s time to train a model.

  1. Download the example notebook to use for this walkthrough.
  2. Choose the Folder icon to change the work area to file management.
  3. Choose the Create folder
  4. Enter a name for the folder.
  5. Choose the folder name.
  6. Choose the Upload file
  7. Choose the Jupyter notebook you downloaded and upload it to the new directory.
  8. Choose the notebook to open a new tab.

Choose the notebook to open a new tab.

You’re prompted to choose a kernel.

  1. Choose Python3 (Data Science).
  2. Choose Select.

Choose Select.

  1. In the second cell of the notebook, replace the project_name variable with the name you gave your project (for this post, iris-multi-01).

You can now run the Jupyter notebook. This notebook creates a very simple pipeline with only two steps: train and register model. It uses the iris dataset and the XGBoost built-in container as the algorithm.

  1. Run the whole notebook.

The process takes some time after you run the cell containing the following code:

start_response = pipeline.start(parameters={
    "TrainingInstanceCount": "1"
})

This starts the training job, which takes approximately 3 minutes to complete. After the training is finished, the next cell of the Jupyter notebook gets the latest version of the model in the model registry and marks it as Approved. Alternatively, you can approve a model from the SageMaker Studio UI. On the Model groups tab, choose the model group and desired version. Choose Update status and Approve before saving.

Choose Update status and Approve before saving

This is the end of the data scientist’s job but the beginning of running the CI/CD pipeline.

Amazon EventBridge monitors the model registry. The listener starts a new deployment job with the provisioned AWS CodePipeline workflow (created with you launched the SageMaker Studio project).

  1. On the CodePipeline console, choose the pipeline starting with the prefix sagemaker-, followed by the name of your project.

On the CodePipeline console, choose the pipeline starting with the prefix sagemaker-, followed by the name of your project.

Shortly after you approve your model, the deployment pipeline starts running. Wait for the pipeline to reach the state DeployStaging. That stage can take approximately 10 minutes to complete. After deploying the first endpoint in the staging account, the pipeline is tested, and then moves to the next step, ApproveDeployment. In this step, it waits for manual approval.

  1. Choose Review.
  2. Enter an approval reason in the text box.
  3. Choose Approve.

The model is now deployed in the production account.

You can also monitor the pipeline on the AWS CloudFormation console, to see the stacks and stack sets the pipeline creates to deploy endpoints in the target accounts. To see the deployed endpoints for each account, sign in to the SageMaker console as either the staging account or production account and choose Endpoints on the navigation pane.

Cleaning up

To clean up all the resources you provisioned in this example, complete the following steps:

  1. Sign in to the console with your main account.
  2. On the AWS CloudFormation console, click on StackSets and delete the following items (endpoints):
    1. Prod – sagemaker-<<sagemaker-project-name>>-<<project-id>>-deploy-prod
    2. Staging – sagemaker-<<sagemaker-project-name>>-<<project-id>>-deploy-staging
  3. In your laptop or workstation terminal, use the AWS Command Line Interface (AWS CLI) and enter the following code to delete your project: 
    aws sagemaker delete-project --project-name iris-multi-01

Make sure you’re using the latest version of the AWS CLI.

Building and customizing a template for your own SageMaker project

SageMaker projects and SageMaker MLOps project templates are powerful features that you can use to automatically create and configure the whole infrastructure required to train, optimize, evaluate, and deploy ML models. A SageMaker project is an AWS Service Catalog provisioned product that enables you to easily create an end-to-end ML solution. For more information, see the AWS Service Catalog Administrator Guide.

A product is a CloudFormation template managed by AWS Service Catalog. For more information about templates and their requirements, see AWS CloudFormation template formats.

ML engineers can design multiple environments and express all the details of this setup as a CloudFormation template, using the concept of infrastructure as code (IaC). You can also integrate these different environments and tasks using a CI/CD pipeline. SageMaker projects provide an easy, secure, and straightforward way of wrapping the infrastructure complexity around in the format of a simple project, which can be launched many times by the other ML engineers and data scientists.

The following diagram illustrates the main steps you need to complete in order to create and publish your custom SageMaker project template.

The following diagram illustrates the main steps you need to complete in order to create and publish your custom SageMaker project template.

We described these steps in more detail in the sections Importing the custom SageMaker Studio Project template and Creating your project.

As an ML engineer, you can design and create a new CloudFormation template for the project, prepare an AWS Service Catalog portfolio, and add a new product to it.

Both data scientists and ML engineers can use SageMaker Studio to create a new project with the custom template. SageMaker invokes AWS Service Catalog and starts provisioning the infrastructure described in the CloudFormation template.

As a data scientist, you can now start training the model. After you register it in the model registry, the CI/CD pipeline runs automatically and deploys the model on the target accounts.

If you look at the CloudFormation template from this post in a text editor, you can see that it implements the architecture we outline in this post.

The following code is a snippet of the template:

Description: Toolchain template which provides the resources needed to represent infrastructure as code.
  This template specifically creates a CI/CD pipeline to deploy a given inference image and pretrained Model to two stages in CD -- staging and production.
Parameters:
  SageMakerProjectName:
    Type: String
  SageMakerProjectId:
    Type: String
…
<<other parameters>>
…
Resources:
  MlOpsArtifactsBucket:
    Type: AWS::S3::Bucket
    DeletionPolicy: Retain
    Properties:
      BucketName: …
…
  ModelDeployCodeCommitRepository:
    Type: AWS::CodeCommit::Repository
    Properties:
      RepositoryName: …
      RepositoryDescription: …
      Code:
        S3:
          Bucket: …
          Key: …
…
  ModelDeployBuildProject:
    Type: AWS::CodeBuild::Project
…
  ModelDeployPipeline:
    Type: AWS::CodePipeline::Pipeline
…

The template has two key sections: Parameters (input parameters of the template) and Resources. SageMaker project templates require that you add two input parameters to your template: SageMakerProjectName and SageMakerProjectId. These parameters are used internally by SageMaker Studio. You can add other parameters if needed.

In the Resources section of the snippet, you can see that it creates the following:

  • A new S3 bucket used by the CI/CD pipeline to store the intermediary artifacts passed from one stage to another.
  • An AWS CodeCommit repository to store the artifacts used during the deployment and testing stages.
  • An AWS CodeBuild project to get the artifacts, and validate and configure them for the project. In the multi-account template, this project also creates a new model registry, used by the CI/CD pipeline to deploy new models.
  • A CodePipeline workflow that orchestrates all the steps of the CI/CD pipelines.

Each time you register a new model to the model registry or push a new artifact to the CodeCommit repo, this CodePipeline workflow starts. These events are captured by an EventBridge rule, provisioned by the same template. The CI/CD pipeline contains the following stages:

  • Source – Reads the artifacts from the CodeCommit repository and shares with the other steps.
  • Build – Runs the CodeBuild project to do the following:
    • Verify if a model registry is already created, and create one if needed.
    • Prepare a new CloudFormation template that is used by the next two deployment stages.
  • DeployStaging – Contains the following components:
    • DeployResourcesStaging – Gets the CloudFormation template prepared in the Build step and deploys a new stack. This stack deploys a new SageMaker endpoint in the target account.
    • TestStaging – Invokes a second CodeBuild project that runs a custom Python script that tests the deployed endpoint.
    • ApproveDeployment – A manual approval step. If approved, it moves to the next stage to deploy an endpoint in production, or ends the workflow if not approved.
  • DeployProd – Similar to DeployStaging, it uses the same CloudFormation template but with different input parameters. It deploys a new SageMaker endpoint in the production account. 

You can start a new training process and register your model to the model registry associated with the SageMaker project. Use the Jupyter notebook provided in this post and customize your own ML pipeline to prepare your dataset and train, optimize, and test your models before deploying them. For more information about these features, see Automate MLOps with SageMaker Projects. For more Pipelines examples, see the GitHub repo.

Conclusions and next steps

In this post, you saw how to prepare your own environment to train and deploy ML models in multiple AWS accounts by using SageMaker Pipelines.

With SageMaker projects, the governance and security of your environment can be significantly improved if you start managing your ML projects as a library of SageMaker project templates.

As a next step, try to modify the SageMaker project template and customize it to address your organization’s needs. Add as many steps as you want and keep in mind that you can capture the CI/CD events and notify users or call other services to build comprehensive solutions.

About the Author

Samir Araújo is an AI/ML Solutions Architect at AWS. He helps customers creating AI/ML solutions solve their business challenges using the AWS platform. He has been working on several AI/ML projects related to computer vision, natural language processing, forecasting, ML at the edge, and more. He likes playing with hardware and automation projects in his free time, and he has a particular interest for robotics.

Read More