base-resources.yaml

AWSTemplateFormatVersion: "2010-09-09"
Description: Setup an EC2 instance for deep racer

Parameters:
  MyIPAddress:
    Type: String
  RuleNumber:
    Type: String
  NetworkAclId:
    Type: String
  VPC:
    Type: String
  Subnet0:
    Type: String
    Default: ''
  Subnet1:
    Type: String
    Default: ''
  Subnet2:
    Type: String
    Default: ''
  Subnet3:
    Type: String
    Default: ''
  Subnet4:
    Type: String
    Default: ''
  Subnet5:
    Type: String
    Default: ''

Outputs:
  InstanceProfile:
    Description: Instance Profile
    Value: !Ref InstanceProfile
    Export:
      Name:
        !Sub '${AWS::StackName}-InstanceProfile'

  Region:
    Description: Region in use
    Value: !Ref AWS::Region
    Export:
      Name:
        !Sub '${AWS::StackName}-Region'

  Bucket:
    Description: S3 Bucket
    Value: !Ref Bucket
    Export:
      Name:
        !Sub '${AWS::StackName}-Bucket'

  SecurityGroup:
    Description: Security Group
    Value: !GetAtt SecurityGroup.GroupId
    Export:
      Name:
        !Sub '${AWS::StackName}-SecurityGroup'

  EFS:
    Description: Elastic FileSystem
    Value: !Ref EFS
    Export:
      Name:
        !Sub '${AWS::StackName}-EFS'

  EFSAccessPoint:
    Description: Elastic FileSystem Access Point
    Value: !Ref EFSAccessPoint
    Export:
      Name:
        !Sub '${AWS::StackName}-EFSAccessPoint'

  AMIParameter:
    Description: SSM Parameter path containing latest created Deepracer AMI
    Value: !Sub /DeepRacer/Images/${AWS::StackName}
    Export:
      Name:
        !Sub '${AWS::StackName}-AMIParameter'

  NetworkAclId:
    Description: Network ACL Id
    Value: !Ref NetworkAclId
    Export:
      Name:
        !Sub '${AWS::StackName}-NetworkAclId'

  InterruptionNotification:
    Description: Interruption Notification
    Value: !Ref InterruptionNotification
    Export:
      Name:
        !Sub '${AWS::StackName}-InterruptionNotification'

  FutureTimeCronExpressionLambdaArn:
    Description: Lambda arn of a function that generates a cloudwatch schedule compatible cron expression for a future time, measured in 'minutes in the future'
    Value: !GetAtt FutureTimeCronExpressionLambda.Arn
    Export:
      Name:
        !Sub '${AWS::StackName}-FutureTimeCronExpressionLambdaArn'

  TerminationLambdaArn:
    Description: Lambda arn of a function that safely terminates an instance
    Value: !GetAtt TerminationLambda.Arn
    Export:
      Name:
        !Sub '${AWS::StackName}-TerminationLambdaArn'

Conditions:
  HasSubnet0: !Not
    - !Equals [!Ref Subnet0, '']

  HasSubnet1: !Not
    - !Equals [!Ref Subnet1, '']

  HasSubnet2: !Not
    - !Equals [!Ref Subnet2, '']

  HasSubnet3: !Not
    - !Equals [!Ref Subnet3, '']

  HasSubnet4: !Not
    - !Equals [!Ref Subnet4, '']

  HasSubnet5: !Not
    - !Equals [!Ref Subnet5, '']

Resources:
  Bucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      LifecycleConfiguration:
        Rules:
          - Id: Delete
            Prefix: imagebuilder
            Status: Enabled
            ExpirationInDays: 1

  BucketAccessPolicy:
    Type: 'AWS::S3::BucketPolicy'
    Properties:
      Bucket: !Ref Bucket
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: 'Allow'
            Principal:
              AWS:
                - !GetAtt InstanceRole.Arn
              Service:
                - deepracer.amazonaws.com
            Action:
              - 's3:GetObject'
              - 's3:ListBucket'
              - 's3:PutObject'
              - 's3:PutObjectAcl'
            Resource:
              - !Join
                - ''
                - - 'arn:aws:s3:::'
                  - !Ref Bucket
              - !Join
                - ''
                - - 'arn:aws:s3:::'
                  - !Ref Bucket
                  - '/*'

  SecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allow access from My IP Address

  SecurityGroupIngress:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      GroupName: !Ref SecurityGroup
      IpProtocol: -1
      CidrIp: !Sub '${MyIPAddress}/32'


  EFSSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      VpcId: !Ref VPC
      GroupDescription: EFS SG
      SecurityGroupIngress:
        -
          SourceSecurityGroupId: !GetAtt SecurityGroup.GroupId
          FromPort: 2049
          IpProtocol: tcp
          ToPort: 2049

  EFSMountTarget0:
    Type: AWS::EFS::MountTarget
    Condition: HasSubnet0
    Properties:
      FileSystemId: !Ref EFS
      SecurityGroups:
        - !Ref EFSSecurityGroup
      SubnetId: !Ref Subnet0

  EFSMountTarget1:
    Type: AWS::EFS::MountTarget
    Condition: HasSubnet1
    Properties:
      FileSystemId: !Ref EFS
      SecurityGroups:
        - !Ref EFSSecurityGroup
      SubnetId: !Ref Subnet1

  EFSMountTarget2:
    Type: AWS::EFS::MountTarget
    Condition: HasSubnet2
    Properties:
      FileSystemId: !Ref EFS
      SecurityGroups:
        - !Ref EFSSecurityGroup
      SubnetId: !Ref Subnet2

  EFSMountTarget3:
    Type: AWS::EFS::MountTarget
    Condition: HasSubnet3
    Properties:
      FileSystemId: !Ref EFS
      SecurityGroups:
        - !Ref EFSSecurityGroup
      SubnetId: !Ref Subnet3

  EFSMountTarget4:
    Type: AWS::EFS::MountTarget
    Condition: HasSubnet4
    Properties:
      FileSystemId: !Ref EFS
      SecurityGroups:
        - !Ref EFSSecurityGroup
      SubnetId: !Ref Subnet4

  EFSMountTarget5:
    Type: AWS::EFS::MountTarget
    Condition: HasSubnet5
    Properties:
      FileSystemId: !Ref EFS
      SecurityGroups:
        - !Ref EFSSecurityGroup
      SubnetId: !Ref Subnet5

  NACLEntry:
    Type: AWS::EC2::NetworkAclEntry
    Properties:
      NetworkAclId: !Ref NetworkAclId
      CidrBlock: !Sub '${MyIPAddress}/32'
      Protocol: '-1'
      RuleAction: allow
      RuleNumber: !Ref RuleNumber

  InstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonEC2FullAccess
        - arn:aws:iam::aws:policy/AmazonElasticFileSystemFullAccess
        - arn:aws:iam::aws:policy/AmazonSSMFullAccess
      Policies:
        - PolicyName: !Sub '${AWS::StackName}-inline-policy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ssmmessages:CreateControlChannel
                  - ssmmessages:CreateDataChannel
                  - ssmmessages:OpenControlChannel
                  - ssmmessages:OpenDataChannel
                Resource: '*'
              - Effect: Allow
                Action:
                  - kms:*
                  - s3:*
                Resource: '*'

  InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles:
      - !Ref InstanceRole

  EFS:
    Type: AWS::EFS::FileSystem
    Properties:
      Encrypted: true

  EFSAccessPoint:
    Type: AWS::EFS::AccessPoint
    Properties:
      FileSystemId: !Ref EFS
      PosixUser:
        Uid: '1000'
        Gid: '1000'
      RootDirectory:
        CreationInfo:
          OwnerGid: '1000'
          OwnerUid: '1000'
          Permissions: '0755'
        Path: '/home/ubuntu/efs'

  InterruptionNotification:
    Type: AWS::SNS::Topic
    Properties:
      DisplayName: InterruptionNotification

  TopicPolicy:
    Type: AWS::SNS::TopicPolicy
    Properties:
      PolicyDocument:
        Statement:
          - Sid: AllowServices
            Effect: Allow
            Principal: '*'
            Action:
              - sns:Publish
              - sns:Subscribe
            Resource: '*'
      Topics:
        - !Ref 'InterruptionNotification'

  FutureTimeCronExpressionLambda:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.handler
      Role: !GetAtt FutureTimeCronExpressionLambdaRole.Arn
      Runtime: python3.7
      Code:
        ZipFile: |
          import cfnresponse
          import datetime

          def handler(event, context):
              physical_resource_id = 'FutureTimeCronExpression'
              response_code = cfnresponse.SUCCESS
              response_data = {'cron_expression': ''}
              if event.get('RequestType', '') == 'Delete':
                  if 'ResponseURL' in event:
                      cfnresponse.send(event, context, response_code, response_data, physical_resource_id)
                  print(response_data)
                  return response_data

              try:
                  ttl = int(event['ResourceProperties']['ttl'])
              except:
                  ttl = 60 * 12 # 12 hours

              expires = datetime.datetime.now() + datetime.timedelta(minutes=ttl)
              cron_expression = f"cron({expires.minute} {expires.hour} {expires.day} {expires.month} ? {expires.year})"
              print(F"cron expression: {cron_expression}")
              response_data['cron_expression'] = cron_expression
              if 'ResponseURL' in event:
                  cfnresponse.send(event, context, response_code, response_data, physical_resource_id)
              print(response_data)
              return response_data

  FutureTimeCronExpressionLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  FutureTimeCronExpressionLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt FutureTimeCronExpressionLambda.Arn
      Action: lambda:InvokeFunction
      Principal: cloudformation.amazonaws.com
      SourceAccount: !Ref AWS::AccountId

  TerminationLambda:
    Type: AWS::Lambda::Function
    Properties:
      Role: !GetAtt TerminationLambdaRole.Arn
      Handler: index.handler
      Runtime: python3.7
      Timeout: 900
      Code:
        ZipFile: !Sub |
          import json
          import boto3
          import logging
          import os
          from botocore.exceptions import ClientError

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)

          def execute_safe_termination(INSTANCE):
              ssm_client = boto3.client('ssm')
              asg_client = boto3.client('autoscaling')
              instance=INSTANCE
              try:
                  asg_info = asg_client.describe_auto_scaling_groups(
                      AutoScalingGroupNames=[
                      instance,
                      ],
                  )
                  instance = asg_info['AutoScalingGroups'][0]['Instances'][0]['InstanceId']
              except:
                  print("Instance isn't part of an ASG, will execute on individual EC2 instance")
              try:
                  response = ssm_client.send_command(
                      InstanceIds=[instance],
                      DocumentName='AWS-RunShellScript',
                      Parameters={'commands': ['su - ubuntu bash -lc /home/ubuntu/bin/safe_termination.sh']},
                      CloudWatchOutputConfig={'CloudWatchOutputEnabled': True},
                      TimeoutSeconds=60)
                  command_id = response['Command']['CommandId']
                  print(f'Running commands on instance {instance}. Command id: {command_id}')
                  waiter = ssm_client.get_waiter('command_executed')
                  waiter.wait(CommandId=command_id, InstanceId=instance, WaiterConfig={'Delay': 30, 'MaxAttempts': 3})
              except ssm_client.exceptions.InvalidInstanceId:
                  print("SSM agent not running.")
              except ClientError as e:
                  print(e.response['Error']['Message'])

          def initiate_stack_deletion(stack):
              cfn_client = boto3.client('cloudformation')
              cfn_client.delete_stack(StackName=stack)

          def notify(stack, instance):
              sns_client = boto3.client('sns')
              topic = '${InterruptionNotification}'
              sns_client.publish(TopicArn=topic, Message=f'Training has completed for instance: {instance} stack: {stack}')

          def handler(event, context):
              instance = event['instance']
              stack = event['stack']
              execute_safe_termination(instance)
              initiate_stack_deletion(stack)
              notify(stack, instance)
              logger.info('Printing event: {}'.format(event))
              return None

              notification_topic = os.environ['INTERRUPTION_NOTIFICATION']

  TerminationLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonEC2FullAccess
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: !Sub '${AWS::StackName}-termination-lambda-policy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ssmmessages:CreateControlChannel
                  - ssmmessages:CreateDataChannel
                  - ssmmessages:OpenControlChannel
                  - ssmmessages:OpenDataChannel
                  - iam:*
                  - sns:*
                  - ssm:*
                  - events:*
                  - lambda:*
                  - cloudformation:DeleteStack
                Resource: '*'