Patch Management for AI Training Clusters using AWS SSM

Question

Pulumi · Accepted Answer

Patch management is an essential part of maintaining the security and performance of any cloud infrastructure. It involves regularly updating software on your systems to patch vulnerabilities, fix bugs, and improve functionality. AWS Systems Manager (SSM) is a powerful service that provides visibility and control over your AWS infrastructure and can automate the patching process.

In this context, managing patches for AI training clusters likely involves a fleet of EC2 instances that run compute-intensive workloads. AWS SSM includes a feature called Patch Manager, which automates the process of patching managed instances with both security-related and other types of updates.

Below is a Pulumi program written in Python that sets up AWS SSM to manage patches for an AI training cluster. This program will create the following resources:

1. An SSM Patch Baseline: Defines which patches are approved for your instances.
2. An SSM Patch Group: Groups together instances that should have the same patch baseline applied.
3. An SSM Association: Associates the patch baseline with the patch group and specifies a schedule for patching.
4. An SSM Maintenance Window (optional): Defines a specific time window during which patching should occur to avoid disrupting operations.

Here's the program:

```python
import pulumi
import pulumi_aws as aws

# AWS Systems Manager Patch Manager requires SSM Agent to be installed and running
# on each EC2 instance. Ensure your AI Training cluster instances have the agent
# installed before using the following setup.

# Create a Patch Baseline for your instances
patch_baseline = aws.ssm.PatchBaseline("ai-training-patch-baseline",
    operating_system="AMAZON_LINUX_2", # Change as needed for your operating system
    approved_patches=["arn:aws:ssm:us-west-2:123456789012:patchbaseline/SOME-PATCH-BASELINE"],
    approval_rules=[
        aws.ssm.PatchBaselineApprovalRuleArgs(
            approve_after_days=7, # Automatically approve patches 7 days after release
            compliance_level="CRITICAL", # Set the compliance severity level
            patch_filters=[
                aws.ssm.PatchBaselineApprovalRulePatchFilterArgs(
                    key="CLASSIFICATION",
                    values=["CriticalUpdates"],
                ),
            ],
        ),
    ],
)

# Create a Patch Group and associate it with the Patch Baseline
patch_group = aws.ssm.PatchGroup("ai-training-patch-group",
    baseline_id=patch_baseline.id,  # The ID of the patch baseline resource we created
    patch_group="ai-training-cluster-group", # Name of the patch group
)

# Create an SSM Association for patching
ssm_association = aws.ssm.Association("ai-training-patch-association",
    name="AWS-RunPatchBaseline", # Predefined SSM Document for patching
    targets=[
        aws.ssm.AssociationTargetArgs(
            key="tag:PatchGroup", # Use tags to target the instances in the patch group
            values=["ai-training-cluster-group"], # Tag value matches the patch group name
        ),
    ],
)

# Optionally, define a Maintenance Window if needed for controlled patch timings
maintenance_window = aws.ssm.MaintenanceWindow("ai-training-maintenance-window",
    schedule="cron(0 4 ? * SUN *)", # Example cron expression for 4 AM every Sunday
    duration=3, # Maintenance window duration in hours
    cutoff=1, # Time before the end of the window to stop scheduling new tasks
)

# Now, create a target mapping for Maintenance Window and Patch Group
maintenance_window_target = aws.ssm.MaintenanceWindowTarget("ai-training-maintenance-target",
    window_id=maintenance_window.id, # The ID of the maintenance window resource
    resource_type="INSTANCE",
    targets=[
        aws.ssm.MaintenanceWindowTargetTargetArgs(
            key="tag:PatchGroup",
            values=["ai-training-cluster-group"],
        ),
    ],
)

# Link the Maintenance Window with the Patching Association
maintenance_window_task = aws.ssm.MaintenanceWindowTask("ai-training-maintenance-task",
    window_id=maintenance_window.id,
    targets=[
        aws.ssm.MaintenanceWindowTaskTargetArgs(
            key="WindowTargetIds",
            values=[maintenance_window_target.id],
        ),
    ],
    task_type="RUN_COMMAND",
    task_arn=ssm_association.name, # ARN of the patching association
    service_role_arn="arn:aws:iam::123456789012:role/MaintenanceWindowRole", # IAM Role with permissions
    max_concurrency="2", # How many targets to patch at the same time
    max_errors="1", # Stop patching if an error occurs on a target
)

# Export the ARN of the created resources for reference
pulumi.export('patch_baseline_arn', patch_baseline.arn)
pulumi.export('patch_group_name', patch_group.patch_group)
pulumi.export('ssm_association_id', ssm_association.id)
pulumi.export('maintenance_window_id', maintenance_window.id)
```

This program is intended for use as part of an existing Pulumi project configured with the proper AWS credentials. It assumes your compute instances are tagged correctly to be associated with the SSM patch group, and it uses a simple cron expression to define a weekly patching schedule.

Remember to adjust the `operating_system`, AWS region, patch classification, IAM role ARNs, and other parameters as needed to fit your deployment. Review the AWS SSM documentation on Patch Manager and be mindful of the cron expression to ensure patching happens during a time that's appropriate for your workload and business requirements.