Federated Access to AWS for Distributed Machine Learning Teams

Question

Pulumi · Accepted Answer

Federated access on AWS allows users from different domains to access AWS resources using their existing authentication mechanisms. Distributed Machine Learning (ML) teams may require federated access to collaborate on projects and use AWS resources like Sagemaker, S3 buckets, etc., without creating separate AWS IAM users for each member.

In AWS, federated access is typically implemented using AWS Identity and Access Management (IAM) roles that trust an external identity provider, such as an SAML provider or an OpenID Connect provider. These roles have policies that grant the necessary permissions to AWS resources and are assumed by users authenticated with the trusted identity provider.

These are the Pulumi AWS components you would use for federated access:

- `aws.iam.Role`: Create IAM roles for your federated users, which dictate what AWS resources they can access and actions they can perform.
- `aws.iam.RolePolicyAttachment`: Attach managed IAM policies to the IAM roles for fine-grained permissions.
- `aws.iam.SAMLProvider`: (if using SAML) Define a SAML provider in IAM, which allows federated users authenticated by the trusted identity provider to assume the roles.

In the context of a distributed Machine Learning team, federated access can be paired with specific AWS services such as:

- `aws.sagemaker`: Provide services to build, train, and deploy machine learning models at scale.
- `aws.s3`: Create and manage storage buckets to store datasets, models, etc.

I'm going to write a program that creates:

1. A SAML-based identity provider.
2. An IAM role with a trust relationship for the SAML provider.
3. A sample policy attachment providing access to Sagemaker services.

You'll need to adapt the details of the trust policy and permissions to fit the specific identity provider and permissions your ML team needs.

```python
import pulumi
import pulumi_aws as aws

# The ARN of the SAML provider in your Identity Provider (IdP), for example, your on-premises Active Directory Federation Services.
# This will be provided by your IdP admin after configuring the trust relationship with AWS.
saml_provider_arn = 'arn:aws:iam::123456789012:saml-provider/IDP_NAME'

# Create a SAML provider in AWS IAM. This allows federated users authenticated by the external IdP to assume roles in AWS.
saml_provider = aws.iam.SamlProvider("mlTeamSamlProvider",
    saml_metadata_document=pulumi.Output.secret("YOUR_SAML_METADATA_DOCUMENT"))

# The IAM role that the SAML-authenticated users will assume when accessing AWS resources.
# The trust relationship policy establishes trust between the SAML provider and AWS.
ml_team_role = aws.iam.Role("mlTeamRole",
    assume_role_policy=pulumi.Output.all(saml_provider.arn, saml_provider_arn).apply(lambda args: f'''
        {{
            "Version": "2012-10-17",
            "Statement": [
                {{
                    "Effect": "Allow",
                    "Principal": {{
                        "Federated": "{args[1]}"
                    }},
                    "Action": "sts:AssumeRoleWithSAML",
                    "Condition": {{
                        "StringEquals": {{
                            "SAML:aud": "https://signin.aws.amazon.com/saml"
                        }}
                    }}
                }}
            ]
        }}
    '''))

# Attach a managed policy to the IAM role to grant Sagemaker access to the federated users.
# Ensure that this policy grants the minimum necessary privileges.
ml_sagemaker_policy_attachment = aws.iam.RolePolicyAttachment("mlSagemakerPolicyAttachment",
    role=ml_team_role.name,
    policy_arn="arn:aws:iam::aws:policy/AmazonSageMakerFullAccess")

# Export the role ARN to be used with your IdP configuration.
pulumi.export('ml_team_role_arn', ml_team_role.arn)
```

This program:

1. Defines an IAM SAML provider in AWS IAM using Pulumi's `aws.iam.SamlProvider`.
2. Creates an IAM role (`aws.iam.Role`) with a trust relationship policy to the SAML provider, allowing federated users to assume the role.
3. Attaches the Amazon SageMaker full access managed policy to the role, allowing federated users to fully manage SageMaker resources.

You'll need to replace `YOUR_SAML_METADATA_DOCUMENT` with the actual SAML metadata XML document provided by your organization's identity provider. The `saml_provider_arn` should also be updated with the ARN of the SAML provider your IdP returns when configuring the trust relationship with AWS.

Ensure that all IAM policies and the trust relationship policy you configure meet your organization's security requirements and grant the least privilege necessary for your ML team to accomplish their tasks.

After deploying this Pulumi program, you'll need to configure your Identity Provider to use the created role's ARN (`ml_team_role_arn`) in its SAML assertions. This allows your federated users to assume the role upon authentication via the IdP.