Enable Amazon SageMaker JumpStart for custom IAM execution roles

With an Amazon SageMaker Domain, you can onboard users with an AWS Identity and Access Management (IAM) execution role different than the Domain execution role. In such case, the onboarded Domain user can’t create projects using templates and Amazon SageMaker JumpStart solutions. This post outlines an automated approach to enable JumpStart for Domain users with a custom execution role. We walk you through two different use cases for enabling JumpStart and how to solve these cases programmatically. The automated solution can help you scale your process to enable JumpStart for Domain users with custom roles, increasing productivity of your data science team and Amazon SageMaker Studio administrators.

JumpStart is a feature within Studio that helps you quickly and easily get started with machine learning (ML). With more and more customers increasingly using ML and adopting Amazon SageMaker, JumpStart is making it easier for data science and ML teams to access and fine-tune more than 150 popular open-source models, such as natural language processing, object detection, and image classification models.

Solution overview

JumpStart requires a SageMaker Domain with project templates enabled for the account and Studio users, as shown in the following screenshot.

If enabled, this setting allows users (configured to use the Domain execution role) to create projects using templates and JumpStart solutions. In the scenario where the user’s execution role is different than the Domain execution role, JumpStart remains disabled for that user even when it’s enabled on the Domain. We address this custom role scenario and the automated solution in the following sections.

In this solution, we address the issue for the following two cases:

Use case 1 – Enabling JumpStart in an automated manner for existing Domain users with custom roles regardless of apps assigned
Use case 2 – Providing a reference script that you can use to programmatically enable JumpStart while onboarding a new Domain user with a custom role

Domain user onboarding

After you create a Domain, you can onboard users to launch apps (such as Studio, RStudio, or Canvas). You must assign a default execution role to a Domain user during the creation process, as shown in the following screenshot.

You can choose a role different than the Domain execution role for a user. However, this may disable JumpStart for such users even when it’s enabled on the Domain. This behavior is due to the fact that SageMaker makes no assumption on a custom role and its permission boundary. The required permissions and policies have to be assigned explicitly to access templates and JumpStart solutions published by SageMaker in AWS Service Catalog.

You can enable SageMaker Projects and JumpStart manually for every user by selecting the user profile on the SageMaker Domain control panel. However, this process can be time-consuming if a user already has some apps assigned. The Edit button at bottom right is only enabled when no apps are assigned to that user (see the following screenshot). You have to delete the assigned apps first in order to edit a user profile.

The cause of the disabled JumpStart feature is evident during Step 2 of editing a user profile, where a message states “If there are individual users using custom execution roles in your organization, you need to enable them on the user profile page.”

In the following sections, we walk you through two automated solutions that cover use cases for both existing and new Domain users.

Prerequisites

The steps described as part of this solution have the following prerequisites:

You have created a SageMaker Domain
The SageMaker Domain authentication method is IAM
Custom roles assigned to the SageMaker Domain users have the AmazonSageMakerFullAccess policy attached

In order for JumpStart Solutions to be enabled for users, the AWS Service Catalog portfolio Amazon SageMaker Solutions and ML Ops products must be imported into the account, and this portfolio must be associated with the role that runs SageMaker. The role association is necessary so that Studio can invoke AWS Service Catalog APIs associated with the Solutions portfolio.

As a general best practice, we recommend testing the process in a non-production environment followed by validation tests to make sure everything is configured and operating as per your expectations before making changes to the production environment.

Use case 1: Enable JumpStart for all existing Domain users with a custom role

Let’s first consider the use case for existing users and enable JumpStart for those users in an automated way.

To achieve this, we have created an AWS CloudFormation template that you can run in the same Region where the SageMaker Domain exists.

The CloudFormation stack contained in the attached jumpstart_solutions_resources.template.yaml file has the following components:

AmazonSageMakerServiceCatalogProductsLaunchRole and AmazonSageMakerServiceCatalogProductsUseRole – Creates these two IAM roles, if they don’t already exist.
1PProductUseRolePolicy – Creates this policy used by AmazonSageMakerServiceCatalogProductsUseRole, if this role doesn’t already exist.
setup_solutions_tests_portfolio – An AWS Lambda function that performs the AWS Service Catalog portfolio import and role association by calling Boto3 APIs. This function is called once during CloudFormation stack creation.
LambdaIAMRole role – Used by the function setup_solutions_tests_portfolio for calling AWS Service Catalog and SageMaker APIs.
SetupPortfolioInvoker – Invokes the function setup_solutions_tests_portfolio.

After the Lambda function runs as part of the CloudFormation deployment, it retrofits all the existing SageMaker Domain users to enable JumpStart and Projects for them. For more information on creating and monitoring a CloudFormation stack, refer to How does AWS CloudFormation work.

Use case 2: Enable JumpStart for a single Domain user with a custom role

Many customers prefer to scale the Domain user onboarding process by automating it programmatically. In this section, we provide a Python script reference that you can use as part of the onboarding process to enable JumpStart for a new user with a custom role. This Python script performs the required association for the given user role. The automated process calling this script must have permission to use AWS Service Catalog and SageMaker APIs. See the following code:

sagemaker_client = boto3.client("sagemaker")
sc_client = boto3.client("servicecatalog")

# function to return 'Amazon SageMaker' portfolio id
def get_solutions_portfolio_id(sc_client):
    portfolio_shares = sc_client.list_accepted_portfolio_shares()
    for portfolio in portfolio_shares['PortfolioDetails']:
            if portfolio['ProviderName'] == 'Amazon SageMaker':
                    return(portfolio['Id'])

portfolio_id = get_solutions_portfolio_id(sc_client)
# import Solutions Service Catalog Portfolio 
sagemaker_client.enable_sagemaker_servicecatalog_portfolio()
    	
sc_client.associate_principal_with_portfolio(
                    PortfolioId=portfolio_id,
                    PrincipalARN=, # custom role ARN
                    PrincipalType='IAM'
                    )

You can either call the script independently or embed it as a step within an automated process to create a user profile for onboarding to Studio. For more information on using Boto3, refer to Boto3 reference.

Clean up

After all the custom roles are enabled to use JumpStart, we can clean up the resources no longer needed. You can delete the Lambda function setup_solutions_tests_portfolio and the IAM role LambdaIAMRole created by the CloudFormation template. The other two IAM roles, AmazonSageMakerServiceCatalogProductsLaunchRole and AmazonSageMakerServiceCatalogProductsUseRole, and the associated policy 1PProductUseRolePolicy (if created) must not be deleted because they need to exist for accessing JumpStart.

Conclusion

In this post, we shared the steps to enable JumpStart for a custom role for existing users as well as new users programmatically. As always, make sure to validate the steps mentioned in this solution in a non-production environment before deploying to production.

Try it out and let us know if you have any questions in the comments section!

Additional resources

For more information, see the following:

About the Authors

Nikhil Jha is a Senior Technical Account Manager at Amazon Web Services. His focus areas include AI/ML, and analytics. In his spare time, he enjoys playing badminton with his daughter and exploring the outdoors.