Connect to Amazon services using AWS PrivateLink in Amazon SageMaker

AWS customers deploying secure development environments often need to restrict incoming and outgoing Internet traffic. This becomes increasingly important with the development of artificial intelligence (ai) due to the data assets that must be protected. Data transmission over 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 to protect 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 building architectures for production in which many AWS services must work together.

In this post, we present a solution to configure SageMaker notebook instances to connect to amazon Bedrock and other AWS services using AWS PrivateLink and amazon Elastic Compute Cloud (amazon EC2) security groups.

Solution Overview

The following example architecture shows a SageMaker instance connecting to several services. The SageMaker instance is isolated from the Internet, but can still access AWS services through PrivateLink. You will notice that the connection to amazon S3 is made through a VPC Gateway endpoint. You can learn more about Gateway VPC endpoints here.

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

Create security groups for outbound and inbound endpoint access

First, you must create the security groups that will be attached to the VPC endpoints and the SageMaker instance. Security groups are created before creating a SageMaker instance because once the instance is created, the security group settings cannot be changed.

Create two groups, one for outgoing and one for incoming. Complete the following steps:

1. In the amazon EC2 console, choose security groups in the navigation panel.

2. Choose Create security group.

3. To Security group nameEnter a name (for example, inbound-sagemaker).

4. To Descriptionenter a description.

5. To VPCchoose your VPC.

6. Note the security group ID to be used in the following steps.

7. Create a new outbound rule.

8. To Security group nameenter a name (for example, outbound-sagemaker).

9. To Descriptionenter description.

10. To VPCchoose the same VPC as the inbound rule.

11. In the Exit 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 outgoing security group ID to be used in the next step.

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

Create a SageMaker instance with the outgoing security group

Now create a SageMaker instance with the network configuration shown in the following screenshot. It is important to choose the same VPC that you used to create the inbound and outbound security groups. Then choose the outgoing security group you created earlier.

Create a VPC interface endpoint

In this step, you will 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. In the amazon VPC console, choose End points in the navigation panel.

2. Choose Create endpoint.

3. To name tagenter a name (for example, rock link).

4. To Service Categoryselect AWS services.

5. To ServicesSearch and choose com.amazonaws..bedrock-run.

6. Configure the VPC to the same one you have been working with.

7. Specify subnets.

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

8. Configure the security group to the incoming security group you created earlier.

After creating the endpoint, it should take some time for it to become available.

Repeat these steps for each service you need for your workflow. The following screenshots show examples of services for which you can create front-end VPC endpoints, such as amazon Simple Storage Service (amazon S3), amazon Kendra, and AWS Lambda. AWS PrivateLink allows you to connect privately to various AWS services; For an updated list, 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 ran this on a Jupyter notebook cell, it would give you an error because you didn't point the invocation to use the VPC endpoint. To do this, add 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 you created in the previous step and find the DNS names, as illustrated in the following screenshot. Private DNS is the best option as it is the same as public, meaning you don't have to change anything to use the private connection. The next best option is to use regional DNS, which is the first option under “DNS Names.” Both options allow your traffic to fail over to other healthy availability zones (AZs), in case the current AZ is corrupted.

Clean

To clean up your resources, complete the following steps:

1. In the SageMaker console, navigate to the notebook settings page.

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

3. In the amazon EC2 console, navigate to the incoming security group details page.

4. About him Behavior menu, choose Delete security groups.

5. Repeat these steps for the outgoing security group.

6. In the amazon VPC console, navigate to the VPC endpoint details page.

7. About him Behavior menu, choose Delete.

8. Repeat this step for each endpoint you created as part of this publication.

Conclusion

In this post, we show how to configure VPC endpoints and security groups to allow SageMaker to connect to amazon Bedrock. When a SageMaker instance has restricted access to the Internet, you can still develop and connect to other AWS services by using 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 start 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 free time, Francisco likes to play music and guitar, play soccer with his daughters, and enjoy time with his family.

Sung Min Hong is an applied scientist in the AWS Generative ai Innovation Center, where he helps accelerate the breadth of use cases for AWS customers. Before joining amazon, Sungmin was a postdoctoral researcher at Harvard Medical School. He has a Ph.D. in Computer Science from New York University. Outside of work, Sungmin enjoys hiking, traveling, and reading.

yash shah He is a scientific manager at the AWS Generative ai Innovation Center. He and his team of applied scientists and machine learning engineers work on a variety of machine learning use cases in healthcare, sports, automotive, and manufacturing.

Anila Joshi has over a decade of experience building ai solutions. As an applied sciences manager in the AWS Generative ai Innovation Center, Anila pioneers innovative ai applications that push the boundaries of possibilities and guides customers to strategically chart a course for the future of ai.