Connect to Amazon services using AWS PrivateLink in Amazon SageMaker

AWS customers that implement secure development environments often have to restrict outbound and inbound internet traffic. This becomes increasingly important with artificial intelligence (AI) development because of the data assets that need to be protected. Transmitting data across the internet is not secure enough for highly sensitive data. Therefore, accessing AWS services without leaving the AWS network can be a secure workflow.

One of the ways you can secure AI development is by creating Amazon SageMaker instances within a virtual private cloud (VPC) with direct internet access disabled. This isolates the instance from the internet and makes API calls to other AWS services not possible. This presents a challenge for developers that are building architectures for production in which many AWS services need to function together.

In this post, we present a solution for configuring SageMaker notebook instances to connect to Amazon Bedrock and other AWS services with the use of AWS PrivateLink and Amazon Elastic Compute Cloud (Amazon EC2) security groups.

Solution overview

The following example architecture shows a SageMaker instance connecting to various services. The SageMaker instance is isolated from the internet but is still able to access AWS services through PrivateLink. One will notice that the connection to Amazon S3 is through a Gateway VPC endpoint. You can learn more about Gateway VPC endpoints here.

In the following sections, we show how to configure this on the AWS Management Console.

Create security groups for outbound and inbound endpoint access

First, you have to create the security groups that will be attached to the VPC endpoints and the SageMaker instance. You create the security groups before creating a SageMaker instance because after the instance has been created, the security group configuration can’t be changed.

You create two groups, one for outbound and another for inbound. Complete the following steps:

1. On the Amazon EC2 console, choose Security Groups in the navigation pane.

2. Choose Create security group.

3. For Security group name, enter a name (for example, inbound-sagemaker).

4. For Description, enter a description.

5. For VPC, choose your VPC.

6. Note the security group ID to use in the next steps.

7. Create a new outbound rule.

8. For Security group name, enter a name (for example, outbound-sagemaker).

9. For Description, enter description.

10. For VPC, choose the same VPC as the inbound rule.

11. In the Outbound rules section, choose Add rule.

12. Add an outbound rule with the inbound security group ID as the destination using HTTPS as the type.

13. Note the outbound security group ID to use in the next step.

14. Return to the inbound security group and add an inbound rule of HTTPS type with the destination set to the outbound security group ID.

Create a SageMaker instance with the outbound security group

You now create a SageMaker instance with the network configuration shown in the following screenshot. It’s important to choose the same VPC that you used to create the inbound and outbound security groups. You then choose the outbound security group you created earlier.

Create an Interface VPC endpoint

In this step, you create an Interface VPC endpoint using Amazon Virtual Private Cloud (Amazon VPC) that automatically uses PrivateLink, which allows calls from your SageMaker instance to AWS services.

1. On the Amazon VPC console, choose Endpoints in the navigation pane.

2. Choose Create endpoint.

3. For Name tag, enter a name (for example, bedrock-link).

4. For Service category, select AWS services.

5. For Services, search for and choose com.amazonaws..bedrock-runtime.

6. Set the VPC to the same one you’ve been working with.

7. Specify the subnet(s).

A subnet is a range of IP addresses within a VPC. If you don’t know what subnet to specify, any subnet will work. Otherwise, specify the subnet that is required by any security requirements from your cloud security team.

8. Set the security group to the inbound security group you created earlier.

After you create the endpoint, it should take some time to become available.

Repeat these steps for every service that you need for your workflow. The following screenshots show examples of services that you can create interface VPC endpoints for, such as Amazon Simple Storage Service (Amazon S3), Amazon Kendra, and AWS Lambda. AWS PrivateLink enables you to connect privately to several AWS services, for a current list please see this page.

Test the connection

You can test the connection to Amazon Bedrock using a simple Python API call. The following is a code snippet that invokes the Amazon Bedrock model:

import boto3
import json

bedrock = boto3.client(service_name='bedrock-runtime')
prompt = """
Human: What type of sharks are there?

Assistant:"""

body = json.dumps({
"prompt": prompt,
"max_tokens_to_sample": 4000,
"temperature": 0.1,
"top_p": 0.9,
})

modelId = 'anthropic.claude-instant-v1'
accept = 'application/json'
contentType = 'application/json'

response = bedrock.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType)
response_body = json.loads(response.get('body').read())

print(response_body.get('completion'))

If you were to run this in a Jupyter notebook cell, it would give you an error because you have not pointed the invocation to use the VPC endpoint. You do this by adding an endpoint URL to the client instantiation:

bedrock = boto3.client(
    service_name='bedrock-runtime',
    endpoint_url = 'https://vpce-0e452bc86b1f87c50-5xltzdpo.bedrock-runtime.us-west-2.vpce.amazonaws.com'
)

To find the endpoint URL, go back to the VPC endpoint that you created in the previous step and look for DNS names, illustrated in the following screenshot. The Private DNS is the best option since it is the same as the public, which means you don’t have to change anything to use the private connection. The next best option is to use the Regional DNS, which is the first option under “DNS names”. Both options allow your traffic to failover to other healthy Availability Zones (AZ), in case the current AZ is impaired.

Clean up

To clean up your resources, complete the following steps:

1. On the SageMaker console, navigate to the notebook configuration page.

2. Stop the instance, then choose Delete to delete the instance.

3. On the Amazon EC2 console, navigate to the inbound security group’s detail page.

4. On the Actions menu, choose Delete security groups.

5. Repeat these steps for the outbound security group.

6. On the Amazon VPC console, navigate to the VPC endpoint’s details page.

7. On the Actions menu, choose Delete.

8. Repeat this is step for every endpoint you created as part of this post.

Conclusion

In this post, we showed how to set up VPC endpoints and security groups to allow SageMaker to connect to Amazon Bedrock. When a SageMaker instance has restricted internet access, you can still develop and connect to other AWS services through the use of AWS PrivateLink. This post showed how to connect to Amazon Bedrock from an isolated SageMaker instance, but you can replicate the steps for other services.

We encourage you to get started developing AI applications on AWS. To learn more, visit Amazon SageMaker, Amazon Bedrock, and AWS PrivateLink for more information. Happy coding!

About the Author

Francisco Calderon is a Data Scientist at the AWS Generative AI Innovation Center. As a member of the GenAI Innovation Center, he helps solve critical business problems for AWS customers using the latest technology in Generative AI. In his spare time, Francisco likes to play music and guitar, play soccer with his daughters, and enjoy time with his family.

Sungmin Hong is an Applied Scientist at AWS Generative AI Innovation Center where he helps expedite the variety of use cases of AWS customers. Before joining Amazon, Sungmin was a postdoctoral research fellow at Harvard Medical School. He holds Ph.D. in Computer Science from New York University. Outside of work, Sungmin enjoys hiking, traveling and reading.

Yash Shah is a Science Manager in the AWS Generative AI Innovation Center. He and his team of applied scientists and machine learning engineers work on a range of machine learning use cases from healthcare, sports, automotive and manufacturing.

Anila Joshi has more than a decade of experience building AI solutions. As an Applied Science Manager at AWS Generative AI Innovation Center, Anila pioneers innovative applications of AI that push the boundaries of possibility and guides customers to strategically chart a course into the future of AI.