Jan Meppe

Machine learning engineer in love with building digital products.

3 minute read

In this blog post I want to show you how to do cross-account Sagemaker endpoint requests in AWS.

We have two accounts, a parent account and a child account. The parent account hosts the Sagemaker endpoint and we want to call it from the child account.

Based on this blog post

High level solution

  • In parent account
    • Create a policy SageMakerExternalInvokeEndpoint that allows you to execute the Sagemaker endpoint
    • Create a role SagemakerExternal and add child account as trusted account
  • In child account
    • Create a lambda on the child account to invoke the API with
    • Change the execution role of the lambda to allow sts:AssumeRole on SagemakerExternal (which is possible because we added it as a trusted account)

In the parent account

In the parent account we are going to do 2 things:

  1. Create a policy SageMakerExternalInvokeEndpoint that will give us permission to call the endpoints on this account
  2. Create a role SagemakerExternal having this policy that can be assumed by someone else (the child account)

Log in to the parent account.

Here we configured the endpoint(s) we want to hit cross-account.

Copy the ARN of the endpoint we want to hit.

Make a new policy that allows someone to invoke Sagemaker endpoints in this account.

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": ["sagemaker:InvokeEndpoint"],
            "Effect": "Allow",
            "Resource": "arn:aws:sagemaker:eu-west-1:<parent account id>:endpoint/*"
        }
    ]
}

Log in to IAM.

Go to policies.

Click on create policy.

Paste the json above.

If you want to, make the policy more generic with wildcards *.

Name it SageMakerExternalInvokeEndpoint.

Now we want to make a new role SageMakerExternal

Click on Trusted Account and paste child id account here

Add the new policy you just made.

Give it a name.

Now the child account is a trusted entity in the parent account.

What did we just do?

We created a new role SagemakerExternal that allows us to invoke the specific sagemaker endpoints on this account with the child account as a trusted account.

Now all that is left to do is to assume this role from the child account.

In the child account: Grant assume role to child account

Now what we need to do is get the child account to allow to assume this role.

Then we need to make a lambda to test it and change the execution role of that lambda to allow sts:AssumeRole.

Log in to child account.

Create new policy AllowAssumeSagemakerExternalRole

This policy will allow the child account to assume this new role.

Give it a nice name

Now create a new lambda.

This is what we want to paste in.

import json
import boto3
import os

from data import body2

def lambda_handler(event, context):
    # Assume the cross-account role (allowed because of sts:AssumeRole)
    sts_client = boto3.client('sts')
    assumed_role = sts_client.assume_role(
        RoleArn='arn:aws:iam::<parent account id>:role/SageMakerExternal',  # add parent id here
        RoleSessionName='SageMakerInvokeSession',
        ExternalId='<parent account id>', # add parent id here
    )
    
    # Create a new session using the assumed role credentials
    credentials = assumed_role['Credentials']
    sagemaker_client = boto3.client(
        'sagemaker-runtime',
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken'],
    )

    # Add the endpoint you want to call cross-account here
    endpoint_name = "f73559-..." 

    # Invoke the SageMaker endpoint
    response = sagemaker_client.invoke_endpoint(
        EndpointName=endpoint_name,
        ContentType='application/json',
        Body=json.dumps({}),
    )

    # Parse the response
    result = json.loads(response['Body'].read().decode())

    return {
        'statusCode': 200,
        'body': json.dumps(result)
    }

If you try to run this now you get the following error

“errorMessage”: “An error occurred (AccessDenied) when calling the AssumeRole operation: User: arn:aws:sts:::assumed-role/ is not authorized to perform: sts:AssumeRole on resource: arn:aws:iam:::role/SageMakerExternal",

What we need to do is change the lambda execution role.

Remember we made this policy: AllowAssumeSagemakerExternalRole

We need to add this policy to the lambda execution role.

Go to the lambda and click on edit.

Go to the bottom and click on edit

Add this piece

{
    "Effect": "Allow",
    "Action": "sts:AssumeRole",
    "Resource": "arn:aws:iam::<parent account id>:role/SageMakerExternal"
},

That’s it now you can run cross account sagemaker endpoint requests on AWS!

Conclusion

In this blog posted I showed you how to do cross account Sagemaker requests.

Summary:

  1. In the parent account with create a policy for invocations and a role with the child account as a trusted account
  2. In the child account allow that account to assume the role from step 1.

Subscribe

Comments

You May Also Enjoy