Enhancing AI Workloads with Amazon SageMaker HyperPod: New Features for Security and Storage Management

Introduction to Amazon SageMaker HyperPod

Amazon SageMaker HyperPod is a purpose-built infrastructure for optimizing foundation model training and inference at scale, eliminating the complexities involved in building and optimizing machine learning infrastructure.

Addressing the Need for Compliance and Security

As AI applications expand across diverse domains, the integration of robust security measures and storage options is increasingly vital for large enterprises. SageMaker HyperPod introduces critical features that bolster control and flexibility in production deployments.

Key Features

Customer Managed Keys (CMK) Support

• Encryption Control: Customers can encrypt EBS volumes and custom AMIs using their own keys.
• Compliance Assurance: Helps meet organizational security requirements such as HIPAA and FIPS compliance.

Amazon EBS CSI Driver Support

• Storage Lifecycle Management: Integration with Amazon EBS facilitates efficient dynamic storage management for large-scale AI workloads.

Configuration and Implementation

Setting Up Customer Managed Keys (CMK)

• IAM Role Configuration: Ensure proper permissions are set up for KMS key interactions.
• Using KMS Keys: Detailed instructions on specifying customer-managed keys when creating or updating clusters.

Working with Amazon EBS

• Dynamic and Ephemeral Volumes: Utilize EBS volumes effectively for both persistent and ephemeral storage needs.

Demonstrating New Features

Volume Resizing Demo with EBS CSI Driver

Step-by-step process to resize a volume using the EBS CSI driver, showcasing the added flexibility and efficiency for storage management.

Conclusion and Future Directions

The enhancements made in SageMaker HyperPod through CMK support and EBS CSI driver integration not only elevate security and compliance standards but also enrich storage management capabilities. This positions organizations to effectively handle large-scale AI workloads while maintaining organizational policies and regulations.

About the Authors

A diverse team of experts from AWS share insights on their experiences and contributions to the SageMaker HyperPod service, highlighting the commitment to scalable and efficient ML solutions.

Optimizing ML Workloads at Scale with Amazon SageMaker HyperPod

As the demand for artificial intelligence (AI) grows across various industries, organizations face the challenge of efficiently training and deploying foundation models (FMs) at scale. Amazon SageMaker HyperPod offers a solution: a purpose-built infrastructure designed specifically for optimizing the training and inference of these large-scale models. By simplifying the complexities of building and managing machine learning (ML) infrastructure, SageMaker HyperPod empowers data scientists and engineers to focus on innovation rather than infrastructure maintenance.

The Importance of Security and Compliance

For large enterprises, security and compliance are paramount. As companies deploy AI solutions across multiple domains, ensuring that GPU clusters adhere to organizational security policies becomes crucial. The latest updates to SageMaker HyperPod enhance its capabilities with two key features: continuous scaling and integration with customer-managed keys (CMK).

Customer Managed Keys (CMK) Support

One standout feature of HyperPod EKS is its support for CMK, enabling customers to encrypt EBS volumes with their own encryption keys. This control is critical for organizations needing to meet regulatory requirements, such as HIPAA and FIPS compliance.

Key Points About CMK Configuration

• Optional Encryption: CMK is optional for EBS volumes. If a customer does not specify a key, AWS-managed keys are used.
• Immutable Key: Once set, the CMK for existing volumes cannot be changed.
• Multiple Configurations: Each instance group can configure one root volume with a CMK and one secondary volume.
• Custom AMIs: Customers can also independently encrypt custom AMIs with CMK.

Implementing CMK ensures comprehensive data-at-rest protection, maintaining tight control over sensitive information throughout the lifecycle of the instances.

Amazon EBS CSI Driver Support

SageMaker HyperPod now supports the Amazon Elastic Block Store (Amazon EBS) Container Storage Interface (CSI) driver. This development allows for dynamic storage management within Kubernetes, making it easier to handle large datasets and model artifacts necessary for training foundation models.

With CSI, users can:

• Provision and Mount Volumes Dynamically: Choose between cluster-level provisioning or Pod-level management, allowing flexibility based on workload requirements.
• Manage Lifecycles of EBS Volumes: Handle both ephemeral and persistent storage more effectively.

Getting Started

To leverage these new features, ensure you meet the following prerequisites:

1. AWS IAM Configuration: Your IAM execution role should have the required permissions for AWS KMS to allow operations such as scaling instance counts and updating cluster software.
2. KMS Key Policy: Properly configure your KMS key policy to authorize the necessary operations with your HyperPod execution role.

Setting Up CMK on HyperPod

1. Create a KMS Key: Identify the KMS key that you’ll use for encrypting EBS volumes and custom AMIs.
2. Update Cluster Configuration: When creating or updating your HyperPod cluster, specify the customer-managed keys in your API calls.

Example: Create a Cluster with CMK Support

aws sagemaker create-cluster \
  --cluster-name my-cluster \
  --instance-groups '[{
       "ExecutionRole": "arn:aws:iam:::role/YourRole",
       "InstanceCount": 2,
       "InstanceGroupName": "YourGroupName",
       "InstanceStorageConfigs": [
             {
                 "EbsVolumeConfig": {
                     "RootVolume": True,
                     "VolumeKmsKeyId": "arn:aws:kms:region:key/key-id"
                 }
             },
             {
                 "EbsVolumeConfig": {
                     "VolumeSizeInGB": 100,
                     "VolumeKmsKeyId": "arn:aws:kms:region:key/key-id"
                 }
             }
        ],
       "InstanceType": "YourInstanceType"
  }]' \
  --vpc-config '{
      "SecurityGroupIds": ["YourSecurityGroupId"],
      "Subnets": ["YourSubnetId"]
  }'

Using Amazon EBS CSI Driver

To implement Amazon EBS CSI, you will need to navigate the following steps:

1. IAM Setup: Create an IAM Service Account with an appropriate policy for interacting with EBS.
2. Deploy the EBS CSI Add-On: Use eksctl to create and install the EBS CSI driver add-on within your EKS cluster.

Example: Create IAM Service Account

eksctl create iamserviceaccount \
  --name ebs-csi-controller-sa \
  --namespace kube-system \
  --cluster YourClusterName \
  --role-name DemoRole \
  --attach-policy-arn arn:aws:iam::aws:policy/service-role/AmazonEBSCSIDriverPolicy \
  --approve

Conclusion

With the introduction of customer-managed key support and Amazon EBS CSI driver integration, Amazon SageMaker HyperPod stands as a robust solution for enterprises looking to optimize their machine learning workflows. These enhancements ensure adherence to security policies while providing flexible storage management options, ultimately enabling organizations to efficiently train and deploy foundational models at scale.

For those interested in diving deeper into these features, check out our comprehensive documentation and explore how to best utilize SageMaker HyperPod for your projects.

About the Authors

This blog post was contributed by a team of AWS solutions architects and engineers with extensive experience in AI and ML workloads. Learn more about our contributors and their backgrounds:

• Mark Vinciguerra: Associate Specialist Solutions Architect focusing on generative AI.
• Rostislav Povelikin: Senior Specialist Solutions Architect with expertise in distributed training.
• Kunal Jha: Principal Product Manager for Amazon SageMaker HyperPod.
• Takuma Yoshitani: Senior Software Development Engineer focusing on SageMaker HyperPod improvements.
• Vivek Koppuru: Engineering leader on the HyperPod team.
• Ajay Mahendru: Engineering leader with a focus on scalable solutions.
• Siddharth Senger: Senior Software Development Engineer with a rich background in AWS services.

Feel free to connect with us on LinkedIn, and let’s take your AI initiatives to new heights!