base-resources.yaml

AWSTemplateFormatVersion: "2010-09-09"
Description: Setup base requirements for DeepRacer on the Spot

Parameters:
  MyIPAddress:
    Type: String
  RuleNumber:
    Type: String
  NetworkAclId:
    Type: String

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

  DeepRacerServiceRole:
    Description: DeepRacer Service Role
    Value: !GetAtt DeepRacerServiceRole.Arn
    Export:
      Name:
        !Sub '${AWS::StackName}-DeepRacerServiceRole'

  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'

  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'

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'

  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/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:*
                  - sns:*
                  - ssm:*
                  - iam:DeleteRolePolicy
                  - iam:DetachRolePolicy
                  - iam:DeleteRole
                  - events:DeleteRule
                  - events:RemoveTargets
                  - events:PutRule
                  - events:DescribeRule
                  - lambda:RemovePermission
                  - lambda:DeleteFunction
                  - lambda:InvokeFunction
                Resource: '*'
              - Effect: Allow
                Action:
                  - cloudformation:DeleteStack
                  - cloudformation:UpdateStack
                Resource: 
                  - !Sub "arn:${AWS::Partition}:cloudformation:${AWS::Region}:${AWS::AccountId}:stack/*"
              - Effect: Allow
                Action:
                  - iam:passRole
                Resource: 
                  - !GetAtt DeepRacerServiceRole.Arn

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

  DeepRacerServiceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - deepracer.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: !Sub '${AWS::StackName}-inline-policy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - s3:ListAllMyBuckets
                Resource: '*'
              - Effect: Allow
                Action:
                  - s3:GetBucketPolicy
                  - s3:PutBucketPolicy
                  - s3:ListBucket
                  - s3:GetBucketAcl
                  - s3:GetObject
                  - s3:GetObjectVersion
                  - s3:GetObjectAcl
                  - s3:GetBucketLocation
                Resource:
                  - !Join
                    - ''
                    - - 'arn:aws:s3:::'
                      - !Ref Bucket
                  - !Join
                    - ''
                    - - 'arn:aws:s3:::'
                      - !Ref Bucket
                      - '/*'

  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.11
      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

  FutureTimeCronExpressionLambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      RetentionInDays: 7
      LogGroupName: !Sub '/aws/lambda/${FutureTimeCronExpressionLambda}'

  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.11
      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']

  TerminationLambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      RetentionInDays: 7
      LogGroupName: !Sub '/aws/lambda/${TerminationLambda}'

  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
                  - logs:DeleteLogGroup
                Resource: '*'

  QuotaLambda:
    Type: AWS::Lambda::Function
    Properties:
      Role: !GetAtt QuotaLambdaRole.Arn
      Handler: index.handler
      Runtime: python3.11
      Timeout: 30
      Code:
        ZipFile: !Sub |
          import boto3
          import logging
          import cfnresponse
          from botocore.exceptions import ClientError

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          client = boto3.client('service-quotas')

          def handler(event, context):
              if event['RequestType'] == 'Create':
                try:
                  existing_spot_quota = client.get_service_quota(
                      ServiceCode='ec2',
                      QuotaCode='L-3819A6DF'
                  )
                  if existing_spot_quota['Quota']['Value'] < 16:
                    spot_response = client.request_service_quota_increase(
                        ServiceCode='ec2',
                        QuotaCode='L-3819A6DF',
                        DesiredValue=16
                    )
                  else:
                    spot_response="Existing spot quota is already sufficient"
                  existing_ondemand_quota = client.get_service_quota(
                      ServiceCode='ec2',
                      QuotaCode='L-DB2E81BA'
                  )
                  if existing_ondemand_quota['Quota']['Value'] < 16:
                    ondemand_response = client.request_service_quota_increase(
                        ServiceCode='ec2',
                        QuotaCode='L-DB2E81BA',
                        DesiredValue=16
                    )
                  else:
                    ondemand_response="Existing on-demand quota is already sufficient"
                  logger.info('Printing event: {}'.format(spot_response))
                  logger.info('Printing event: {}'.format(ondemand_response))
                  responseData = {}
                  responseData['Data'] = str(spot_response) + ", " +str(ondemand_response)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)
                except ClientError as e:
                  responseData = {}
                  responseData['Data'] = e.response['Error']['Message'] + ".  Likely you have an existing quota request in process, check AWS Support console."
                  print(e.response['Error']['Message'])
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)
              else:
                responseData = {}
                responseData['Data'] = 'Event not a create event, nothing to do'
                cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)

  QuotaLambdaRole:
    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
        - arn:aws:iam::aws:policy/IAMFullAccess
      Policies:
        - PolicyName: !Sub '${AWS::StackName}-quota-lambda-policy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - servicequotas:*
                Resource: '*'
              - Effect: Deny
                Action: 
                  - logs:CreateLogGroup
                Resource: '*'

  QuotaLambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      RetentionInDays: 7
      LogGroupName: !Sub '/aws/lambda/${QuotaLambda}'

  QuotaInvoke:
      Type: AWS::CloudFormation::CustomResource
      Version: "1.0"
      Properties:
        ServiceToken: !GetAtt QuotaLambda.Arn
        FunctionName: !Ref QuotaLambda

  SSMShellScriptLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      RetentionInDays: 7
      LogGroupName: !Sub '/aws/ssm/AWS-RunShellScript'