2.SageMakerVPC.yaml

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

AWSTemplateFormatVersion: '2010-09-09'
Description: >
  This CloudFormation stack creates all the necessary pre-requisites for Amazon SageMaker Hyperpod,
  these include VPC and two Subnets, a security group, a FSx Lustre Filesystem, S3 Bucket and IAM Role.
  A public subnet and a private subnet are created in an Availability Zone that you provide as a parameter.
  As part of the template you'll deploy an Internet Gateway and NAT Gateway in
  the public subnet. In addition you deploy endpoints for Amazon S3. The VPC contains 2 CIDR blocks with 10.0.0.0/16 and 10.1.0.0/16
  The first CIDR is used for the public subnet, the second is used for the private.
  The template creates an fsx lustre volume in the specified AZ with a default of
  1.2 TB storage which can be overridden by parameter. A role is also created which
  helps to execute HyperPod cluster operations.


####################
## Stack Metadata ##
####################

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: General configuration
        Parameters:
          - VPCName
          - SSMDocumentName
      - Label:
          default: Availability Zone configuration for the subnets
        Parameters:
          - PrimarySubnetAZ
          - BackupSubnetAZ
      - Label:
          default: Fsx Lustre storage size
        Parameters:
          - Capacity
      - Label:
          default: Network and S3 endpoints configuration
        Parameters:
          - CreateS3Endpoint
      - Label:
          default: S3 bucket name
        Parameters:
          - S3Bucket
    ParameterLabels:
      VPCName:
        default: Name of your VPC
      SSMDocumentName:
        default: True/False; Create SSM Session Manager Document. Only set to False if SSM-SessionManagerRunShellAsUbuntu document exists in your account.
      PrimarySubnetAZ:
        default: Availability zone id to deploy the primary subnets
      BackupSubnetAZ:
        default: (Optional) Availability zone id to deploy the backup private subnet
      CreateS3Endpoint:
        default: Create an S3 endpoint

######################
## Stack Parameters ##
######################

Parameters:
  VPCName:
    Description: Name of your VPC
    Default: 'SageMaker HyperPod VPC'
    Type: String

  PrimarySubnetAZ:
    Description: Availability zone id in which the public subnet and primary private subnet will be created.
    Type: String
    Default: usw2-az4

  BackupSubnetAZ:
    Description: Availability zone id in which the backup private subnet will be created. Specify it when you need multiple AZs for other AWS services (e.g. AWS Directory Service). Leave empty if you don't need it.
    Type: String
    Default: ''

  CreateS3Endpoint:
    AllowedValues:
      - 'true'
      - 'false'
    Default: 'true'
    Description:
      Set to false if to avoid creating an S3 endpoint on your VPC.
    Type: String

  Capacity:
    Description: Storage capacity in GiB (1200 or increments of 2400)
    Type: Number
    Default: 1200

  S3Bucket:
    Description: S3 Bucket to save lifecycle configuration file
    Type: String
    Default: "sagemaker-lifecycle"


  PerUnitStorageThroughput:
    Description: Provisioned Read/Write (MB/s/TiB)
    Type: Number
    Default: 250
    AllowedValues:
      - 125
      - 250
      - 500
      - 1000

  Compression:
    Description: Data compression type
    Type: String
    AllowedValues:
      - "LZ4"
      - "NONE"
    Default: "LZ4"

  LustreVersion:
    Description: Lustre software version
    Type: String
    AllowedValues:
      - "2.15"
      - "2.12"
    Default: "2.15"

  SSMDocumentName:
    Type: String
    Default: 'true'
    AllowedValues:
      - 'true'
      - 'false'
    Description: Boolean on whether or not to create the SSM Session Manager Document (choose false only if it already exists in your account).

  CreateOpenZFS:
    AllowedValues:
      - 'true'
      - 'false'
    Default: 'false'
    Description:
      Creates an FSx OpenZFS file system in addition to the FSxL file system. Default is false
    Type: String
 
  OpenZFSThroughput:
    AllowedValues:
      - 160
      - 320
      - 640
      - 1280
      - 2560
      - 3840
      - 5120
      - 7680
      - 10240
    Default: 160
    Description: OpenZFS throughput in MB/s. Allowed values are 160, 320, 640, 1280, 2560, 3840, 5120, 7680, 10240
    Type: Number

  OpenZFSStorageCapacity:
    Description: Storage capacity in GiB for OpenZFS (64 GiB to 524,288 GiB)
    Type: Number
    Default: 64
    MinValue: 64
    MaxValue: 524288

###############################
## Conditions for Parameters ##
###############################

Conditions:
  S3EndpointCondition: !Equals [!Ref 'CreateS3Endpoint', 'true']
  BackupSubnetCondition: !Not [ !Equals [!Ref 'BackupSubnetAZ', ''] ]
  CreateSSMDocument: !Equals [!Ref 'SSMDocumentName', 'true']
  CreateOpenZFSCondition: !Equals [!Ref 'CreateOpenZFS', 'true']


##########################
## Rules for Parameters ##
##########################

Rules:
  AZsRule:
    Assertions:
      - Assert: !Not
        - !Equals
          - !Ref PrimarySubnetAZ
          - !Ref BackupSubnetAZ
        AssertDescription: Backup AZ has to be different from the primary AZ.


#########################
## VPC & Network Setup ##
#########################

Mappings:
  Networking:
    VPC:
      CIDR0: 10.0.0.0/16
      CIDR1: 10.1.0.0/16

Resources:
  # Create a VPC
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      EnableDnsSupport: true
      EnableDnsHostnames: true
      CidrBlock: !FindInMap [Networking, VPC, CIDR0]
      Tags:
        - Key: Name
          Value: SageMaker HyperPod VPC

  VpcCidrBlock:
    Type: AWS::EC2::VPCCidrBlock
    DependsOn: VPC
    Properties:
      VpcId: !Ref VPC
      CidrBlock: !FindInMap [Networking, VPC, CIDR1]

  FlowLogsRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service: vpc-flow-logs.amazonaws.com
          Action: sts:AssumeRole
      Policies:
      - PolicyName: flowlogs-policy
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Action:
            - logs:CreateLogStream
            - logs:PutLogEvents
            - logs:DescribeLogGroups
            - logs:DescribeLogStreams
            Resource: !GetAtt FlowLogsGroup.Arn
  FlowLogsGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      RetentionInDays: 7

  FlowLogVPC:
    Type: AWS::EC2::FlowLog
    Properties:
      DeliverLogsPermissionArn: !GetAtt FlowLogsRole.Arn
      LogGroupName: FlowLogsGroup
      ResourceId: !Ref VPC
      ResourceType: VPC
      TrafficType: ALL

  # Create an IGW and add it to the VPC
  InternetGateway:
    Type: AWS::EC2::InternetGateway

  GatewayToInternet:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref VPC
      InternetGatewayId: !Ref InternetGateway

  # Create a NAT GW then add it to the public subnet
  NATGateway:
    Type: AWS::EC2::NatGateway
    Properties:
      AllocationId: !GetAtt ElasticIP.AllocationId
      SubnetId: !Ref PublicSubnet

  ElasticIP:
    Type: AWS::EC2::EIP
    Properties:
      Domain: vpc

  # NOTE: when you create additional security groups, you must ensure that every
  # security group has ingress/egress from/to its own security group id. Failure
  # to do so may cause trn1/p4d/p4de/p5 SMHP cluster creation to fail:
  #
  #     Instance i-aaaabbbbccccddddf failed to provision with the following
  #     error: "EFA health checks did not run successfully. Ensure that your
  #     VPC and security groups are properly configured before attempting to
  #     create a new cluster." Note that multiple instances may be impacted."
  SecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allow EFA communication for Multi-Node Parallel Batch jobs
      VpcId: !Ref VPC
  EFASecurityGroupIngress:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      Description: All to all communication for EFA Ingress within Security Group
      IpProtocol: -1
      FromPort: -1
      ToPort: -1
      GroupId: !Ref SecurityGroup
      SourceSecurityGroupId: !Ref SecurityGroup
  EFASecurityGroupEgress:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      Description: All to all communication for EFA Egress  within Security Group
      IpProtocol: -1
      FromPort: -1
      ToPort: -1
      GroupId: !Ref SecurityGroup
      DestinationSecurityGroupId: !Ref SecurityGroup
  EFASecurityGroupEgressECS:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      Description: All to all communication for Egress to all
      IpProtocol: -1
      FromPort: -1
      ToPort: -1
      GroupId: !Ref SecurityGroup
      CidrIp: 0.0.0.0/0

  # Build the public subnet
  PublicSubnet:
    Type: AWS::EC2::Subnet
    DependsOn: VPC
    Properties:
      MapPublicIpOnLaunch: true
      VpcId: !Ref VPC
      CidrBlock: !Select [ 0, !Cidr [ !GetAtt VPC.CidrBlock, 2, 15 ]]
      AvailabilityZoneId: !Ref PrimarySubnetAZ
      Tags:
        - Key: Name
          Value: !Join [ ' ', [ !Ref VPCName, 'Public Subnet -', !Ref PrimarySubnetAZ ] ]

  # Create the primary private subnet
  PrimaryPrivateSubnet:
    Type: AWS::EC2::Subnet
    DependsOn: [VpcCidrBlock]
    Properties:
      VpcId: !Ref VPC
      CidrBlock: !Select [ 0, !Cidr [ !FindInMap [Networking, VPC, CIDR1], 2, 15 ]]
      AvailabilityZoneId: !Ref PrimarySubnetAZ
      Tags:
        - Key: Name
          Value: !Join [ ' ', [ !Ref VPCName, 'Private Subnet -', !Ref PrimarySubnetAZ ] ]

  # Create the backup private subnet
  BackupPrivateSubnet:
    Condition: BackupSubnetCondition
    Type: AWS::EC2::Subnet
    DependsOn: [VpcCidrBlock]
    Properties:
      VpcId: !Ref VPC
      CidrBlock: !Select [ 1, !Cidr [ !FindInMap [Networking, VPC, CIDR1], 2, 15 ]]
      AvailabilityZoneId: !Ref BackupSubnetAZ
      Tags:
        - Key: Name
          Value: !Join [ ' ', [ !Ref VPCName, 'Private Subnet -', !Ref BackupSubnetAZ ] ]

  # Create and set the public route table
  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC

  PublicRoute:
    Type: AWS::EC2::Route
    DependsOn: GatewayToInternet
    Properties:
      RouteTableId: !Ref PublicRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway

  # Then the private route table
  PrivateRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC

  PrivateRouteToInternet:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref PrivateRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId: !Ref NATGateway

  # Associate the public route table to the public subnet
  PublicSubnetRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnet
      RouteTableId: !Ref PublicRouteTable

  # and the primary private subnet to the private route table
  PrimaryPrivateSubnetRTAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PrimaryPrivateSubnet
      RouteTableId: !Ref PrivateRouteTable

  # and the backup private subnet to the private route table
  BackupPrivateSubnetRTAssociation:
    Condition: BackupSubnetCondition
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref BackupPrivateSubnet
      RouteTableId: !Ref PrivateRouteTable

  # S3 endpoint
  S3Endpoint:
    Condition: S3EndpointCondition
    Type: AWS::EC2::VPCEndpoint
    Properties:
      PolicyDocument:
        Version: 2012-10-17
        Statement:
            - Effect: Allow
              Principal: '*'
              Action:
                - '*'
              Resource:
                - '*'
      RouteTableIds:
        - !Ref PublicRouteTable
        - !Ref PrivateRouteTable
      ServiceName: !Join
        - ''
        - - com.amazonaws.
          - !Ref AWS::Region
          - .s3
      VpcId: !Ref VPC

  FSxLFilesystem:
    Type: AWS::FSx::FileSystem
    DeletionPolicy: Delete
    UpdateReplacePolicy: Delete
    Properties:
      FileSystemType: LUSTRE
      StorageType: SSD
      FileSystemTypeVersion: !Ref LustreVersion
      StorageCapacity: !Ref Capacity
      SecurityGroupIds:
        - !Ref SecurityGroup
      SubnetIds:
        - !Ref PrimaryPrivateSubnet
      LustreConfiguration:
        DataCompressionType: !Ref Compression
        DeploymentType: PERSISTENT_2
        PerUnitStorageThroughput: !Ref PerUnitStorageThroughput
        MetadataConfiguration:
          Mode: AUTOMATIC

  FSxOpenZFSFileSystem:
    Type: AWS::FSx::FileSystem
    Condition: CreateOpenZFSCondition
    DeletionPolicy: Delete
    UpdateReplacePolicy: Delete
    Properties:
      FileSystemType: OPENZFS
      OpenZFSConfiguration:
        AutomaticBackupRetentionDays: 30
        CopyTagsToBackups: Yes
        CopyTagsToVolumes: Yes
        DailyAutomaticBackupStartTime: '19:00'
        DeploymentType: SINGLE_AZ_HA_2
        Options:
          - DELETE_CHILD_VOLUMES_AND_SNAPSHOTS
        RootVolumeConfiguration:
          DataCompressionType: NONE
          NfsExports: 
            - ClientConfigurations: 
              -  Clients: '*'
                 Options: 
                    - rw
                    - no_root_squash
                    - crossmnt
        ThroughputCapacity: !Ref OpenZFSThroughput
        WeeklyMaintenanceStartTime: '1:04:00'
      SecurityGroupIds:  
        - !Ref SecurityGroup
      StorageType: SSD
      StorageCapacity: !Ref OpenZFSStorageCapacity
      SubnetIds:
        - !Ref PrimaryPrivateSubnet

  AmazonSagemakerClusterExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/AmazonSageMakerClusterInstanceRolePolicy"
      Policies:
        - PolicyName: AmazonSagemakerClusterVPCPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                    [
                    "ec2:CreateNetworkInterface",
                    "ec2:CreateNetworkInterfacePermission",
                    "ec2:DeleteNetworkInterface",
                    "ec2:DeleteNetworkInterfacePermission",
                    "ec2:DescribeNetworkInterfaces",
                    "ec2:DescribeVpcs",
                    "ec2:DescribeDhcpOptions",
                    "ec2:DescribeSubnets",
                    "ec2:DescribeSecurityGroups",
                    "ec2:DetachNetworkInterface",
                    "ec2:CreateTags"
                ]
                Resource: "*"
  LCScriptsBucket:
    Type: 'AWS::S3::Bucket'
    DeletionPolicy: Retain
    Properties:
      BucketName:
       !Sub
          - '${S3Bucket}-${RandomGUID}'
          - { RandomGUID: !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref AWS::StackId ]]]] }

  SSMSessionManagerRunShellAsUbuntu:
    Type: AWS::SSM::Document
    Condition: CreateSSMDocument
    Properties:
      Name: "SSM-SessionManagerRunShellAsUbuntu"
      DocumentType: "Session"
      Content:
        schemaVersion: "1.0"
        description: "Document to hold regional settings for Session Manager"
        sessionType: "Standard_Stream"
        inputs:
          s3BucketName: ""
          s3KeyPrefix: ""
          s3EncryptionEnabled: true
          cloudWatchLogGroupName: ""
          cloudWatchEncryptionEnabled: true
          cloudWatchStreamingEnabled: true
          idleSessionTimeout: "20"
          maxSessionDuration: ""
          kmsKeyId: ""
          runAsEnabled: true
          runAsDefaultUser: "ubuntu"
          shellProfile:
            windows: ""
            linux: "/bin/bash -c 'sudo su - ubuntu'"

#############
## Outputs ##
#############
Outputs:
  VPC:
    Value: !Ref VPC
    Description: ID of the VPC
    Export:
      Name: !Sub ${AWS::StackName}-VPC
  PublicSubnet:
    Value: !Ref PublicSubnet
    Description: ID of the public subnet
    Export:
      Name: !Sub ${AWS::StackName}-PublicSubnet
  PrimaryPrivateSubnet:
    Value: !Ref PrimaryPrivateSubnet
    Description: ID of the primary private subnet
    Export:
      Name: !Sub ${AWS::StackName}-PrimaryPrivateSubnet
  BackupPrivateSubnet:
    Value: !Ref BackupPrivateSubnet
    Description: ID of the backup private subnet
    Export:
      Name: !Sub ${AWS::StackName}-BackupPrivateSubnet
    Condition: BackupSubnetCondition
  SecurityGroup:
    Value: !Ref SecurityGroup
    Description: SecurityGroup for Batch
    Export:
      Name: !Sub ${AWS::StackName}-SecurityGroup
  FSxLustreFilesystemMountname:
    Description: The ID of the FSxL filesystem that has been created
    Value: !GetAtt FSxLFilesystem.LustreMountName
    Export:
      Name: !Sub ${AWS::StackName}-FSxLustreFilesystemMountname
  FSxLustreFilesystemDNSname:
    Description: The DNS of the FSxL filesystem that has been created
    Value: !GetAtt FSxLFilesystem.DNSName
    Export:
      Name: !Sub ${AWS::StackName}-FSxLustreFilesystemDNSname
  FSxLustreFilesystemId:
    Description: The ID of the FSxL filesystem that has been created
    Value: !Ref FSxLFilesystem
    Export:
      Name: !Sub ${AWS::StackName}-FSxLustreFilesystemId
  AmazonSagemakerClusterExecutionRoleArn:
    Description: The Role ARN used for cluster creation
    Value: !GetAtt AmazonSagemakerClusterExecutionRole.Arn
    Export:
      Name: !Sub ${AWS::StackName}-AmazonSagemakerClusterExecutionRoleArn
  AmazonS3BucketName:
    Description: The S3 bucket where Lifecycle scripts are stored
    Value: !Ref LCScriptsBucket
    Export:
      Name: !Sub ${AWS::StackName}-AmazonSagemakerLCScriptsBucketName
  FSxOpenZFSFileSystemId:
    Condition: CreateOpenZFSCondition
    Description: The ID of the FSxOpenZFS filesystem that has been created
    Value: !Ref FSxOpenZFSFileSystem
  FSxOpenZFSFileSystemDNSname:
    Condition: CreateOpenZFSCondition
    Description: The DNS of the FSxOpenZFS filesystem that has been created
    Value: !GetAtt FSxOpenZFSFileSystem.DNSName