Integrating Cognito with SageMaker for Managed AI Workflows

Question

Pulumi · Accepted Answer

To integrate AWS Cognito with AWS SageMaker for Managed AI Workflows, you would typically use AWS Cognito to manage user authentication and authorization for your machine learning application, while SageMaker provides a fully managed service to build, train, and deploy machine learning models.

The following Pulumi Python program illustrates how to create a Cognito User Pool, a User Pool Client, an Identity Pool, and connects them to a SageMaker domain for managing AI workflows. In this example, we will:

1. Create a Cognito User Pool, which is a user directory in Amazon Cognito that provides sign-up and sign-in options for app users.
2. Create a Cognito User Pool Client, which allows an application to use the user pool to authenticate users.
3. Create a Cognito Identity Pool, which enables you to create unique identities for your users and federate them with identity providers.
4. Lastly, we create a SageMaker domain which provides an integrated development environment (IDE) for machine learning and supports managing AI workflows. The SageMaker domain will be linked to the Cognito User Pool for authenticating users.

The SageMaker domain configuration includes the setting for our default execution role and user settings. Specifically, the `userSettings` parameter includes a `JupyterServerAppSettings` configuration to link with the Cognito resources.

Let's dive into the code:

```python
import pulumi
import pulumi_aws as aws

# Create a Cognito User Pool
cognito_user_pool = aws.cognito.UserPool("aiUserPool",
    password_policy=aws.cognito.UserPoolPasswordPolicyArgs(
        minimum_length=8,
        require_lowercase=True,
        require_numbers=True,
    ))

# Create a Cognito User Pool Client
cognito_user_pool_client = aws.cognito.UserPoolClient("aiUserPoolClient",
    user_pool_id=cognito_user_pool.id,
    # It's a good practice to set the token validity periods and allowed OAuth flows according to your security requirements.
    allowed_oauth_flows=["code"],
    allowed_oauth_flows_user_pool_client=True,
    allowed_oauth_scopes=["email", "openid"],
    callback_urls=["https://www.example.com/callback"],
    default_redirect_uri="https://www.example.com/callback",
    explicit_auth_flows=["ALLOW_REFRESH_TOKEN_AUTH", "ALLOW_USER_SRP_AUTH"],
    generate_secret=True,
    logout_urls=["https://www.example.com/logout"])

# Create a Cognito Identity Pool
cognito_identity_pool = aws.cognito.IdentityPool("aiIdentityPool",
    identity_pool_name="AIIdentityPool",
    allow_unauthenticated_identities=False,  # Set to True if you want to support unauthenticated logins
    cognito_identity_providers=[aws.cognito.IdentityPoolCognitoIdentityProviderArgs(
        client_id=cognito_user_pool_client.id,
        provider_name=cognito_user_pool.endpoint,
    )])

# Assume role for Sagemaker's authentication and authorization
assume_role_policy = aws.iam.get_policy_document(statements=[
    aws.iam.GetPolicyDocumentStatementArgs(
        actions=["sts:AssumeRoleWithWebIdentity"],
        effect="Allow",
        principals=[
            aws.iam.GetPolicyDocumentStatementPrincipalArgs(
                type="Federated",
                identifiers=[cognito_identity_pool.arn],
            ),
        ],
    ),
])

# Create an IAM role Sagemaker will assume when accessing resources
sagemaker_execution_role = aws.iam.Role("aiSagemakerExecutionRole",
    assume_role_policy=assume_role_policy.json)

# Attach the necessary AWS managed policies for Sagemaker
for policy_arn in ["arn:aws:iam::aws:policy/AmazonSageMakerFullAccess", "arn:aws:iam::aws:policy/AmazonS3FullAccess"]:
    aws.iam.RolePolicyAttachment(policy_arn.replace(":", "-"),
        role=sagemaker_execution_role.name,
        policy_arn=policy_arn)

# Create a SageMaker Domain
sagemaker_domain = aws.sagemaker.Domain("aiSagemakerDomain",
    auth_mode="SSO",
    default_user_settings=aws.sagemaker.DomainDefaultUserSettingsArgs(
        execution_role=sagemaker_execution_role.arn,
        jupyter_server_app_settings=aws.sagemaker.DomainDefaultUserSettingsJupyterServerAppSettingsArgs(
            default_resource_spec=aws.sagemaker.DomainDefaultUserSettingsJupyterServerAppSettingsDefaultResourceSpecArgs(
                instance_type="ml.t3.medium",  # Adjust the instance type according to your needs
            ),
        ),
        security_groups=[],  # Add any security groups if needed
    ),
    domain_name="ai-sagemaker-domain",
    subnet_ids=["subnet-xxxxxxxx"], # Replace with your actual subnet IDs
    vpc_id="vpc-xxxxxxxx",  # Replace with your actual VPC ID
    app_network_access_type="PublicInternetOnly")

# Output the necessary identifiers for later use in the application
pulumi.export("cognitoUserPoolId", cognito_user_pool.id)
pulumi.export("cognitoUserPoolClientId", cognito_user_pool_client.id)
pulumi.export("cognitoIdentityPoolId", cognito_identity_pool.id)
pulumi.export("sagemakerDomainId", sagemaker_domain.id)
```

In this program, the last section exports identifiers such as the `cognitoUserPoolId` and `sagemakerDomainId` so you can reference them in your application or in other parts of your infrastructure as required.

Here's an overview of why this program uses the resources it does:

- **Cognito User Pool**: Acts as a user directory that will manage all the users for our application.
- **Cognito User Pool Client**: Provides our application a way to interface with the Cognito User Pool to authenticate users.
- **Cognito Identity Pool**: Allows users to exchange their user pool tokens for AWS credentials so that they can access other AWS services. In this case, we are granting access to SageMaker.
- **SageMaker Domain**: Represents the main gateway for developers to manage their SageMaker resources and work through integrated Jupyter notebooks.

Configure the VPC and subnet identifiers according to your AWS environment where you plan to deploy the resources. Additionally, ensure that the allowed OAuth flows and scopes are appropriate for your app's authentication requirements. Make sure your Pulumi AWS provider is configured correctly with sufficient permissions.

Please note that this program assumes that you will handle the deployment and management of the actual machine learning models and notebooks within SageMaker separately. The focus here is on setting up the infrastructure for authentication and authorization to the SageMaker environment.