Categories: FAANG

Manage Amazon SageMaker JumpStart foundation model access with private hubs

js private hub.sm

Amazon SageMaker JumpStart is a machine learning (ML) hub offering pre-trained models and pre-built solutions. It provides access to hundreds of foundation models (FMs). A private hub is a feature in SageMaker JumpStart that allows an organization to share their models and notebooks so as to centralize model artifacts, facilitate discoverability, and increase the reuse within the organization. With new models released daily, many enterprise admins want more control over the FMs that can be discovered and used by users within their organization (for example, only allowing models based on pytorch framework to be discovered).

Now enterprise admins can effortlessly configure granular access control over the FMs that SageMaker JumpStart provides out of box so that only allowed models can be accessed by users within their organizations. In this post, we discuss the steps required for an administrator to configure granular access control of models in SageMaker JumpStart using a private hub, as well as the steps for users to access and consume models from the private hub.

Solution overview

Starting today, with SageMaker JumpStart and its private hub feature, administrators can create repositories for a subset of models tailored to different teams, use cases, or license requirements using the Amazon SageMaker Python SDK. Admins can also set up multiple private hubs with different lists of models discoverable for different groups of users. Users are then only able to discover and use models within the private hubs they have access to through Amazon SageMaker Studio and the SDK. This level of control empowers enterprises to consume the latest in open weight generative artificial intelligence (AI) development while enforcing governance guardrails. Finally, admins can share access to private hubs across multiple AWS accounts, enabling collaborative model management while maintaining centralized control. SageMaker JumpStart uses AWS Resource Access Manager (AWS RAM) to securely share private hubs with other accounts in the same organization. The new feature is available in the us-east-2 AWS Region as of writing, and will be available to more Regions soon.

The following diagram shows an example architecture of SageMaker JumpStart with its public and private hub features. The diagram illustrates how SageMaker JumpStart provides access to different model repositories, with some users accessing the public SageMaker JumpStart hub and others using private curated hubs.

In the following section, we demonstrate how admins can configure granular access control of models in SageMaker JumpStart using a private hub. Then we show how users can access and consume allowlisted models in the private hub using SageMaker Studio and the SageMaker Python SDK. Finally, we look at how an admin user can share the private hub with users in another account.

Prerequisites

To use the SageMaker Python SDK and run the code associated with this post, you need the following prerequisites:

An AWS account that contains all your AWS resources
An AWS Identity and Access Management (IAM) role with access to SageMaker Studio notebooks
SageMaker JumpStart enabled in a SageMaker Studio domain

Create a private hub, curate models, and configure access control (admins)

This section provides a step-by-step guide for administrators to create a private hub, curate models, and configure access control for your organization’s users.

Because the feature has been integrated in the latest SageMaker Python SDK, to use the model granular access control feature with a private hub, let’s first update the SageMaker Python SDK:
```
!pip3 install sagemaker —force-reinstall —quiet
```

Next, import the SageMaker and Boto3 libraries:

import boto3
from sagemaker import Session
from sagemaker.jumpstart.hub.hub import Hub

Configure your private hub:
```
HUB_NAME="CompanyHub"
HUB_DISPLAY_NAME="Allowlisted Models"
HUB_DESCRIPTION="These are allowlisted models taken from the JumpStart Public Hub."
REGION="<your_region_name>" # for example, "us-west-2"
```
In the preceding code, HUB_NAME specifies the name of your Hub. HUB_DISPLAY_NAME is the display name for your hub that will be shown to users in UI experiences. HUB_DESCRIPTION is the description for your hub that will be shown to users.

Set up a Boto3 client for SageMaker:

sm_client = boto3.client('sagemaker')
session = Session(sagemaker_client=sm_client)
session.get_caller_identity_arn()

Check if the following policies have been already added to your admin IAM role; if not, you can add them as inline policies:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "s3:ListBucket",
                "s3:GetObject",
                "s3:GetObjectTagging"
            ],
            "Resource": [
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>",
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>/*"
            ],
            "Effect": "Allow"
        }
    ]
}

Replace the <REGION> placeholder using the configurations in Step 3.

In addition to setting up IAM permissions to the admin role, you need to scope down permissions for your users so they can’t access public contents.

Use the following policy to deny access to the public hub for your users. These can be added as inline policies in the user’s IAM role:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": "s3:*",
            "Effect": "Deny",
            "Resource": [
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>",
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>/*"
            ],
            "Condition": {
                "StringNotLike": {"s3:prefix": ["*.ipynb", "*/eula.txt"]}
            }
        },
        {
            "Action": "sagemaker:*",
            "Effect": "Deny",
            "Resource": [
                "arn:aws:sagemaker:<REGION>:aws:hub/SageMakerPublicHub",
                "arn:aws:sagemaker:<REGION>:aws:hub-content/SageMakerPublicHub/*/*"
            ]
        }
    ]
}

Replace the <REGION> placeholder in the policy using the configurations in Step 3.

After you have set up the private hub configuration and permissions, you’re ready to create the private hub.

Use the following code to create the private hub within your AWS account in the Region you specified earlier:

hub = Hub(hub_name=HUB_NAME, sagemaker_session=session)

try:
  hub.create(
      description=HUB_DESCRIPTION,
      display_name=HUB_DISPLAY_NAME
  )
  print(f"Successfully created Hub with name {HUB_NAME} in {REGION}")
except Exception as e:
  if "ResourceInUse" in str(e):
    print(f"A hub with the name {HUB_NAME} already exists in your account.")
  else:
    raise e

Use hub.describe() to verify the configuration of your hub.After your private hub is set up, you can add a reference to models from the SageMaker JumpStart public hub to your private hub. No model artifacts need to be managed by the customer. The SageMaker team will manage any version or security updates.For a list of available models, refer to Built-in Algorithms with pre-trained Model Table.

To search programmatically, run the command

filter_value = "framework == meta"
response = hub.list_sagemaker_public_hub_models(filter=filter_value)
models = response["hub_content_summaries"]
while response["next_token"]:
    response = hub.list_sagemaker_public_hub_models(filter=filter_value,
                                                    next_token=response["next_token"])
    models.extend(response["hub_content_summaries"])

print(models)

The filter argument is optional. For a list of filters you can apply, refer to SageMaker Python SDK.

Use the retrieved models from the preceding command to create model references for your private hub:
```
for model in models:
    print(f"Adding {model.get('hub_content_name')} to Hub")
    hub.create_model_reference(model_arn=model.get("hub_content_arn"), 
                               model_name=model.get("hub_content_name"))
```
The SageMaker JumpStart private hub offers other useful features for managing and interacting with the curated models. Administrators can check the metadata of a specific model using the hub.describe_model(model_name=<model_name>) command. To list all available models in the private hub, you can use a simple loop:
```
response = hub.list_models()
models = response["hub_content_summaries"]
while response["next_token"]:
    response = hub.list_models(next_token=response["next_token"])
    models.extend(response["hub_content_summaries"])

for model in models:
    print(model.get('HubContentArn'))
```
If you need to remove a specific model reference from the private hub, use the following command:
```
hub.delete_model_reference("<model_name>")
```
If you want to delete the private hub from your account and Region, you’ll need to delete all the HubContents first, then delete the private hub. Use the following code:
```
for model in models:
    hub.delete_model_reference(model_name=model.get('HubContentName')) 

hub.delete()
```

Interact with allowlisted models (users)

This section offers a step-by-step guide for users to interact with allowlisted models in SageMaker JumpStart. We demonstrate how to list available models, identify a model from the public hub, and deploy the model to endpoints from SageMaker Studio as well as the SageMaker Python SDK.

User experience in SageMaker Studio

Complete the following steps to interact with allowlisted models using SageMaker Studio:

On the SageMaker Studio console, choose JumpStart in the navigation pane or in the Prebuilt and automated solutions section.
Choose one of model hubs you have access to. If the user has access to multiple hubs, you’ll see a list of hubs, as shown in the following screenshot.
If the user has access to only one hub, you’ll go straight to the model list.
You can view the model details and supported actions like train, deploy, and evaluate.
To deploy a model, choose Deploy.
Modify your model configurations like instances and deployment parameters, and choose Deploy.

User experience using the SageMaker Python SDK

To interact with your models using the SageMaker Python SDK, complete the following steps:

Just like the admin process, the first step is to force reinstall the SageMaker Python SDK:
```
!pip3 install sagemaker —force-reinstall —quiet
```

Import the SageMaker and Boto3 libraries:

import boto3
from sagemaker import Session
from sagemaker.jumpstart.hub.hub import Hub
from sagemaker.jumpstart.model import JumpStartModel
from sagemaker.jumpstart.estimator import JumpStartEstimator

To access the models in your private hub, you need the Region and the name of the hub on your account. Fill out the HUB_NAME and REGION fields with the information provided by your administrator:

HUB_NAME="CompanyHub" 
REGION="<your_region_name>" # for example, "us-west-2"
sm_client = boto3.client('sagemaker') 
sm_runtime_client = boto3.client('sagemaker-runtime') 
session = Session(sagemaker_client=sm_client, 
                    sagemaker_runtime_client=sm_runtime_client)
hub = Hub(hub_name=HUB_NAME, sagemaker_session=session)

List the models available in your private hub using the following command:

response = hub.list_models()
models = response["hub_content_summaries"]
while response["next_token"]:
    response = hub.list_models(next_token=response["next_token"])
    models.extend(response["hub_content_summaries"])

print(models)

To get more information about a particular model, use the describe_model method:

model_name = "huggingface-llm-phi-2"
response = hub.describe_model(model_name=model_name) 
print(response)

You can deploy models in a hub with the Python SDK by using JumpStartModel. To deploy a model from the hub to an endpoint and invoke the endpoint with the default payloads, run the following code. To select which model from your hub you want to use, pass in a model_id and version. If you pass in * for the version, it will take the latest version available for that model_id in the hub. If you’re using a model gated behind a EULA agreement, pass in accept_eula=True.
```
model_id, version = "huggingface-llm-phi-2", "1.0.0"
model = JumpStartModel(model_id, version, hub_name=HUB_NAME, 
                            region=REGION, sagemaker_session=session)
predictor = model.deploy(accept_eula=False)
```

To invoke your deployed model with the default payloads, use the following code:

example_payloads = model.retrieve_all_examples()
for payload in example_payloads:
    response = predictor.predict(payload.body)
    print("nInputn", payload.body, "nnOutputn", 
                response[0]["generated_text"], "nn===============")

To delete the model endpoints that you created, use the following code:
```
predictor.delete_model()
predictor.delete_endpoint()
```

Cross-account sharing of private hubs

SageMaker JumpStart private hubs support cross-account sharing, allowing you to extend the benefits of your curated model repository beyond your own AWS account. This feature enables collaboration across different teams or departments within your organization, even when they operate in separate AWS accounts. By using AWS RAM, you can securely share your private hubs while maintaining control over access.

To share your private hub across accounts, complete the following steps:

On the AWS RAM console, choose Create resource share.
When specifying resource share details, choose the SageMaker hub resource type and select one or more private hubs that you want to share. When you share a hub with any other account, all of its contents are also shared implicitly.
Associate permissions with your resource share.
Use AWS account IDs to specify the accounts to which you want to grant access to your shared resources.
Review your resource share configuration and choose Create resource share.

It may take a few minutes for the resource share and principal associations to complete.

Admins that want to perform the preceding steps programmatically can enter the following command to initiate the sharing:

# create a resource share using the private hub
aws ram create-resource-share 
    --name test-share 
    --resource-arns arn:aws:sagemaker:<region>:<resource_owner_account_id>:hub/<hub_name> 
    --principals <consumer_account_id>  
    --region <region>

Replace the <resource_owner_account_id>, <consumer_account_id>, <hub_name>, and <region> placeholders with the appropriate values for the resource owner account ID, consumer account ID, name of the hub, and Region to use.

After you set up the resource share, the specified AWS account will receive an invitation to join. They must accept this invitation through AWS RAM to gain access to the shared private hub. This process makes sure access is granted only with explicit consent from both the hub owner and the recipient account. For more information, refer to Using shared AWS resources.

You can also perform this step programmatically:

# list resource shares
aws ram get-resource-share-invitations 
    --region <region>

# accept resource share
# using the arn from the previous response 
aws ram accept-resource-share-invitation  
  --resource-share-invitation-arn <arn_from_ previous_request> 
  --region <region>

For detailed instructions on creating resource shares and accepting invitations, refer to Creating a resource share in AWS RAM. By extending your private hub across accounts, you can foster collaboration and maintain consistent model governance across your entire organization.

Conclusion

SageMaker JumpStart allows enterprises to adopt FMs while maintaining granular control over model access and usage. By creating a curated repository of approved models in private hubs, organizations can align their AI initiatives with corporate policies and regulatory requirements. The private hub decouples model curation from model consumption, enabling administrators to manage the model inventory while data scientists focus on developing AI solutions.

This post explained the private hub feature in SageMaker JumpStart and provided steps to set up and use a private hub, with minimal additional configuration required. Administrators can select models from the public SageMaker JumpStart hub, add them to the private hub, and manage user access through IAM policies. Users can then deploy these preapproved models, fine-tune them on custom datasets, and integrate them into their applications using familiar SageMaker interfaces. The private hub uses the SageMaker underlying infrastructure, allowing it to scale with enterprise-level ML demands.

For more information about SageMaker JumpStart, refer to SageMaker JumpStart. To get started using SageMaker JumpStart, access it through SageMaker Studio.

About the Authors

Raju Rangan is a Senior Solutions Architect at AWS. He works with government-sponsored entities, helping them build AI/ML solutions using AWS. When not tinkering with cloud solutions, you’ll catch him hanging out with family or smashing birdies in a lively game of badminton with friends.

Sherry Ding is a senior AI/ML specialist solutions architect at AWS. She has extensive experience in machine learning with a PhD in computer science. She mainly works with public sector customers on various AI/ML-related business challenges, helping them accelerate their machine learning journey on the AWS Cloud. When not helping customers, she enjoys outdoor activities.

June Won is a product manager with Amazon SageMaker JumpStart. He focuses on making foundation models easily discoverable and usable to help customers build generative AI applications. His experience at Amazon also includes mobile shopping applications and last mile delivery.

Bhaskar Pratap is a Senior Software Engineer with the Amazon SageMaker team. He is passionate about designing and building elegant systems that bring machine learning to people’s fingertips. Additionally, he has extensive experience with building scalable cloud storage services.

AI Generated Robotic Content