Tracking User Activity on AWS SageMaker with CloudTrail

Question

Pulumi · Accepted Answer

To track user activity on AWS SageMaker, you can enable AWS CloudTrail, which provides a history of AWS API calls for your account, including API calls made via the AWS Management Console, AWS SDKs, command-line tools, and other AWS services. CloudTrail captures all API calls for SageMaker as events, including calls from the SageMaker console and from code calls to the SageMaker APIs.

Here's how you can use Pulumi to create an AWS CloudTrail Trail that tracks SageMaker activity:

1. **Create a new S3 bucket** to store your CloudTrail logs.
2. **Set up a CloudTrail Trail** that logs all SageMaker actions to the S3 bucket you just created.
3. (Optional) You can also create an SNS topic and have CloudTrail send notifications to this topic on every log delivery.

Let's walk through each of these steps with Pulumi in Python.

First, you'll need to import the required Pulumi packages. We're using the `pulumi_aws` package here:

```python
import pulumi
import pulumi_aws as aws
```

Next, create a new S3 bucket:

```python
# Create an S3 bucket to store CloudTrail logs
cloudtrail_logs_bucket = aws.s3.Bucket("cloudtrail-logs-bucket")
```

After the S3 bucket is created, set up the CloudTrail Trail:

```python
# Create a CloudTrail that logs all events related to SageMaker
sagemaker_trail = aws.cloudtrail.Trail("sagemaker-trail",
    s3_bucket_name=cloudtrail_logs_bucket.id,
    enable_logging=True,
    event_selectors=[
        aws.cloudtrail.TrailEventSelectorArgs(
            read_write_type="All",
            include_management_events=True,
            data_resources=[
                aws.cloudtrail.TrailEventSelectorDataResourceArgs(
                    type="AWS::SageMaker::NotebookInstance",
                    values=["arn:aws:sagemaker:*"]
                ),
            ],
        ),
    ],
    include_global_service_events=True,
)
```

Here, we've set `read_write_type` to `"All"` to capture both read and write activity, which includes creation, modification, and deletion of resources. The `include_management_events` flag is set to `True` to ensure you’re logging management operations that are performed on resources in your AWS account.

(Optional) Create an SNS topic, which will receive simple notifications service messages for log file delivery from CloudTrail:

```python
# (Optional) Create an SNS topic for CloudTrail notifications
cloudtrail_logs_topic = aws.sns.Topic("cloudtrail-logs-topic")

# Configure the Trail to send logs to the new SNS topic
sagemaker_trail = aws.cloudtrail.Trail("sagemaker-trail",
    # ... existing configuration ...
    sns_topic_name=cloudtrail_logs_topic.id,
)
```

Finally, export the S3 bucket name and SNS topic ARN for easy access:

```python
# Export the names of the resources
pulumi.export("cloudtrail_logs_bucket", cloudtrail_logs_bucket.id)
pulumi.export("cloudtrail_logs_topic_arn", cloudtrail_logs_topic.arn)
```

Here's the entire program, combining all the steps:

```python
import pulumi
import pulumi_aws as aws

# Create an S3 bucket to store CloudTrail logs
cloudtrail_logs_bucket = aws.s3.Bucket("cloudtrail-logs-bucket")

# Create an SNS topic for CloudTrail notifications (Optional)
cloudtrail_logs_topic = aws.sns.Topic("cloudtrail-logs-topic")

# Create a CloudTrail that logs all events related to SageMaker
sagemaker_trail = aws.cloudtrail.Trail("sagemaker-trail",
    s3_bucket_name=cloudtrail_logs_bucket.id,
    enable_logging=True,
    sns_topic_name=cloudtrail_logs_topic.id,  # Omit this line if you don't want SNS notifications
    event_selectors=[
        aws.cloudtrail.TrailEventSelectorArgs(
            read_write_type="All",
            include_management_events=True,
            data_resources=[
                aws.cloudtrail.TrailEventSelectorDataResourceArgs(
                    type="AWS::SageMaker::NotebookInstance",
                    values=["arn:aws:sagemaker:*"]
                ),
            ],
        ),
    ],
    include_global_service_events=True,
)

# Export the names of the resources
pulumi.export("cloudtrail_logs_bucket", cloudtrail_logs_bucket.id)
pulumi.export("cloudtrail_logs_topic_arn", cloudtrail_logs_topic.arn)
```

To run this Pulumi program:
- Save the above code to a file named `__main__.py`.
- Run `pulumi up` in the terminal in the directory containing the file.

This will create the S3 Bucket, CloudTrail Trail, and optionally the SNS Topic. All SageMaker activities will now be tracked and the logs will be stored in the designated S3 bucket.