Amazon SageMaker private workforce in CDK Python

This project demonstrates how to create an Amazon SageMaker private workforce, paired with an Amazon Cognito user pool using the AWS Cloud Development Kit (AWS CDK).

Please note that this sample solution and code are only for testing and non-production purposes.

Amazon SageMaker private workforces can be created in a fast and intuitive way using the AWS Management Console, but many organizations manage infrastructure deployment through infrastructure as code (IaC). However, creating a private workforce with IaC is not a straightforward task because of some complex technical dependencies between services during the initial creation, as detailed further below. Such dependencies are addressed by this solution.

Please note that only one private workforce can be created for each region in a given account. Furthermore, you can later use work teams to assign workers in a private workforce to labeling jobs, but work teams are not included in this sample.

Deployment

Prerequisites

• An AWS account, already bootstrapped for the AWS CDK.
• Verify that the target deployment region doesn't have an existing private workforce, since only one for each region in a given account is supported.
• AWS credentials with sufficient permissions to deploy the solution.
• AWS CDK installed (version 2.178.1 or later).
• Python (version 3.13 or later).
• AWS CLI.
• A mobile device with an authenticator app installed.

Local setup

Recommended: use a local Python environment to install the dependencies:

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Deploy the stack

Verify that you have available credentials with sufficient permissions to deploy in the expected AWS account:

aws sts get-caller-identity

Deploy with the AWS CDK, providing target account ID and region for deployment:

cdk synth -c account=111122223333 -c region=aa-example-1

cdk deploy -c account=111122223333 -c region=aa-example-1

replacing 111122223333 and aa-example-1 with the desired Account ID and deployment region, respectively. Account ID and region are required parameters.

Optionally, you can also include the prefix of the Amazon Cognito domain (must be globally unique) or a comma-separated list of email addresses that will be invited to join the private workforce:

cdk synth -c account=111122223333 -c region=aa-example-1 -c domainName=example -c emails="<email-1>,<email-2>"

cdk deploy -c account=111122223333 -c region=aa-example-1 -c domainName=example -c emails="<email-1>,<email-2>"

replacing 111122223333, aa-example-1, example, and "<email-1>,<email-2>" with the desired Account ID, the deployment region, the prefix that Amazon Cognito will prepend to the URL of the managed login, and the list of comma-separated emails to invite to the private workforce, respectively.

The domainName argument is a prefix that Amazon Cognito will use to create the managed login with the format example.auth.aa-example-1.amazoncognito.com. Note that the domain prefix is immutable after first deployment and persisted in AWS Systems Manager Parameter Store for consistency across stack updates. If you don't specify it, the solution will create a randomized prefix. On subsequent deployments, the name will be read from Parameter Store.

Open the AWS CloudFormation console to follow the deployment and find additional information.

Test

If you invited yourself from the CLI using the emails argument to join the private workforce, follow the instructions in the email that you received to register and access the labeling portal. Otherwise, follow the steps below to invite yourself and others to join the private workforce. For more information, see the “Create a user” section of the Amazon Cognito documentation.

1. On the Amazon Cognito console, choose User pools in the navigation pane.
2. Choose the existing user pool, MyWorkforceUserPool.
3. Choose Users, then choose Create a user.
4. Choose Email as the alias attribute to sign in.
5. Choose Send an email invitation as the invitation message.
6. For User name, enter a name for the new user. Make sure not to use the email address.
7. For Email address, enter the email address of the worker to be invited.
8. For simplicity, choose Generate a password for the user.
9. Choose Create.

Then follow the instructions in the invitation email to register and access the labeling portal.

Troubleshooting

For a successful deployment, verify the following:

• Confirm you have sufficient permissions to create all of the stack resources.
• Note that only one private workforce for each region in a given account is supported. Verify you don't have an existing private workforce in the chosen account and region.
• Note that the domain of the managed login must be globally unique, so make sure to enter an unique prefix when using domainName argument in the CLI.
• For subsequent deployments, maintain consistency with private workforce name, domain name, and provided email addresses. Those resources can cause failures when changed on redeployment.
• Please create users directly from the Amazon Cognito console as described above.

Check the messages from the AWS CloudFormation console for more information in case of deployment errors.

Clean up

To clean up your resources, open the AWS CloudFormation console and delete the stack “Workforce”. Alternatively, if you deployed using the AWS CDK CLI, you can also run cdk destroy from the same terminal where you ran cdk deploy and using the same AWS CDK CLI arguments as during deployment.

Stack Overview

This CDK stack creates an Amazon SageMaker private workforce with a dedicated Amazon Cognito user pool. The main components of the stack include:

1. Amazon Cognito user pool, configured with specific password policies and multi-factor authentication (MFA) settings, and related resources including:
   1. Amazon Cognito user pool app client.
   2. Amazon Cognito user pool domain. It is defined via a prefix (see documentation). The prefix is stored in AWS Systems Manager Parameter Store to be persisted between stack updates. Stack deletion will also delete the stored prefix. You can specify a domain prefix of your choice as input.
   3. Associated log delivery to Amazon CloudWatch.
   4. Associated AWS WAF.
   5. Initial set of users if provided as input.
2. Custom AWS Identity and Access Management (IAM) roles that provide prmissions to AWS CloudFormation custom resources to interact with Amazon Cognito, Amazon SageMaker AI, and other services.
3. AWS CloudFormation custom resources that are used to create and manage various resources, including:
   1. SageMaker private workforce and other API calls.
   2. Custom AWS Lambda function to filter and update Amazon Cognito app client URLs.

Stack creation flow

The following diagram illustrates the resource creation sequence and dependencies:

Amazon Cognito      Amazon Cognito              AWS Systems Manager      Amazon Cognito
  user pool    -->  user pool app client  --->  Parameter Store    ---> user pool domain
    ^                |               ^           string parameter
    |                |               |
    |                V               |
    |        Custom Resource         |
    |             IAM Role           |
    |                |               |
    |                |               |
    |                V               |
    |          Amazon SageMaker      |
    |          private workforce     |
    |                |               |
Update email <-------|-------> Update callback URLs
invitation with
labeling portal
    |
    v
Amazon Cognito users
from CLI-provided emails

The stack creates the resources in a specific order with several dependencies.

• We create the Amazon Cognito user pool app client and domain after the user pool. A placeholder is used as callback URL.
• Then, using custom resources, we create the Amazon SageMaker private workforce using the various Amazon Cognito resources.
  • Note: aws_cdk.custom_resources.AwsCustomResource automatically adds permissions to the IAM role of the underlying Lambda function to allow the SDK calls.
  • Note: Creating the Amazon SageMaker private workforce via API automatically adds callback and logout URLs to the Amazon Cognito user pool app client.
• We update the email invitation for Amazon Cognito users to include the URL of the labeling portal.
• We remove the placeholder URL from the Amazon Cognito app client.
• We create Amazon Cognito users and invite them to join the private workforce uning the email.

Development

Run synth offline

To validate the infrastructure code, use cdk synth (which will also apply cdk-nag rules that are configured) to compile the AWS CDK code into a CloudFormation template. To avoid AWS API calls during the synth process (that check for the existence of a string parameter in AWS Systems Manager Parameter Store), explicitly provide a domainName argument:

pip install -r requirements.txt
cdk synth -c account=111122223333 -c region=aa-example-1 -c domainName=example

Pre-commit setup

Run pre-commit to validate the code of the solution. Its configuration includes (see .pre-commit-config.yaml for further details):

• Launching nox.
• General checks.
• Running bandit.
• Run cdk synth.

Install utilities for development:

pip install -r requirements-dev.txt -r requirements.txt

Run pre-commit:

pre-commit run -a

Additionally, you can also run only tests (pytest) and linters/formatters with nox.

nox

See noxfile.py for details.

Security

See CONTRIBUTING for more information.

License

This library is licensed under the MIT-0 License. See the LICENSE file.