AWS Lambda function to shutdown AWS SageMaker Notebooks cross account

Hello!

I see that you are trying to set up a Lambda function in one account, (account A) so that it is able to shutdown SageMaker Notebooks in a separate account, (account B).

In this case, the Lambda function in account A would need to assume a role in account B to be able to shutdown the SageMaker Notebook instances in account B.

1. In account B, you would create a policy and then assign said policy to a role. This would be a policy which gives permission to Lambda to shutdown a SageMaker Notebook instance. In account B’s IAM, create a policy, choose SageMaker and give StopNotebookInstance permission. See below for example policy, where <account-B-ID> should populate with the AWS account ID of account B, and you can replace ‘name-of-instance’ with the name of the SageMaker instance. Then, create a role (AWS service, Lambda) by using the policy we made in step 1.

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AllowLambdaToStopSageMakerInstance",
            "Effect": "Allow",
            "Action": "sagemaker:StopNotebookInstance",
            "Resource": "arn:aws:sagemaker:*<account-B-ID>:notebook-instance/name-of-instance"
        }
    ]
}

2. Next, you would give the Lambda function in account A permissions to assume the role in account B. In account A’s IAM, create a policy, choose STS, and choose ‘AssumeRole’. The policy will populate similar to below but replace <account-B-ID> with the AWS account ID of account B, and replace ‘role-on-source-account’ with the name of the assumed role. Go to Lambda > Configuration > Permissions and open the default execution role. Add the policy we just made to that role.

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AllowLambdaToAssumeRoleInAccountB",
            "Effect": "Allow",
            "Action": [
                "sts:AssumeRole"
            ],
            "Resource": ["arn:aws:iam::<account-B-ID>:role/role-on-source-account"]
        }
    ]
}

3. Then, you must also modify your cross-account IAM role’s trust policy to allow the Lambda function to assume the role. Go back to account B’s IAM and go to the role we made in the beginning. Go to Trust Relationships > Edit Trust Policy and change it to look like below, but this time, replace <account-A-ID> with the AWS account ID of account A, and replace ‘my-lambda-execution-role’ with the name of the Lambda execution role we made in the step before.

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Principal": {
                "AWS": "arn:aws:iam::<account-A-ID>:role/my-lambda-execution-role"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}

4. Finally, you have to include the AWS STS AssumeRole API call in your Lambda function. This is an example code snippet in python, which includes the API call and what I used to test that it successfully shuts down the instance. You would replace <account-B-ID> with the the AWS account ID of account B, replace ‘role-on-source-account’ with the name of the assumed role, and replace name-of-instance with the name of your SageMaker Notebook instance.

import boto3

# IAM Role ARN in Account B with permissions to stop the SageMaker notebook instance
target_account_role_arn = 'arn:aws:iam::<account-B-ID>:role/SageMakerStopRole2'

# SageMaker notebook instance name to stop
notebook_instance_name = 'mySageMakerNotebook'

# Assume the IAM Role in Account B
sts_client = boto3.client('sts')
assumed_role = sts_client.assume_role(
    RoleArn=target_account_role_arn,
    RoleSessionName='AssumedRoleSession'
)

# Use the temporary credentials to create a SageMaker client in Account B
sagemaker_client = boto3.client(
    'sagemaker',
    aws_access_key_id=assumed_role['Credentials']['AccessKeyId'],
    aws_secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
    aws_session_token=assumed_role['Credentials']['SessionToken']
)

def lambda_handler(event, context):
    try:
        # Stop the SageMaker notebook instance in Account B
        sagemaker_client.stop_notebook_instance(NotebookInstanceName=notebook_instance_name)
        return {
            'statusCode': 200,
            'body': f'Successfully stopped {notebook_instance_name} in Account B.'
        }
    except Exception as e:
        return {
            'statusCode': 500,
            'body': f'Error: {str(e)}'
        }

Now, your Lambda function in account A should have the necessary permissions to shutdown the AWS SageMaker Notebooks in account B.

Note:

• The Lambda function has to be in the same region as the SageMaker Notebook instance.
• You may have to make slight adjustments with the IAM permissions

Here is the documentation as reference:

https://repost.aws/knowledge-center/lambda-function-assume-iam-role

https://repost.aws/knowledge-center/start-stop-lambda-eventbridge

https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-run-and-manage-shut-down.html

Feel free to reach out if you run into any issues!