dcv-start-stop.yaml

AWSTemplateFormatVersion: '2010-09-09'

Description: Simple solution to start/stop DCV instances based an AWS SSM Maintenance Window workflow

Parameters: 
  EC2KeyPair: 
    Type: 'AWS::EC2::KeyPair::KeyName'
  SecurityGroup: 
    Type: 'AWS::EC2::SecurityGroup::Id'
  Instance:
    Type: String
    Default: g4dn.2xlarge
    AllowedValues:
      - g3s.xlarge
      - g3.4xlarge
      - g3.8xlarge
      - g3.16xlarge
      - g3s.xlarge
      - g4dn.xlarge
      - g4dn.2xlarge
      - g4dn.4xlarge
      - g4dn.8xlarge
      - g4dn.12xlarge
      - g4dn.16xlarge
      - p3.2xlarge
      - p3.8xlarge
      - p3.16xlarge
      - p3dn.24xlarge
  Minutes:
    Type: Number
    Default: 10
    AllowedValues:
      - 10
      - 15
      - 20
      - 25
      - 30
      - 60
  Tag:
    Type: String
    Default: "DCV"
    AllowedPattern: "^[a-zA-Z0-9]*$"
  TagValue: 
    Type: String 
    Default: "Yes"
    AllowedPattern: "^[a-zA-Z0-9]*$"
  Action:
    Type: String
    Default: stopInstance
    AllowedValues: 
      - stopInstance
      - terminateInstance

Metadata: 
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label: 
          default: 'Network information'
        Parameters:
          - SecurityGroup
      - Label: 
          default: 'Amazon EC2 information' 
        Parameters:
          - EC2KeyPair
          - Instance
      - Label: 
          default: 'AWS System Manager information'
        Parameters:
          - Tag 
          - TagValue
          - Minutes
      - Label:
          default: 'AWS Lambda Environment Variable'
        Parameters:
          - Action
    ParameterLabels:
      EC2KeyPair:
        default: EC2 Key Pair to decode the Windows password
      SecurityGroup:
        default: Security Group to use with the DCV instance - must allow ingress on port 8443
      Instance: 
        default: Instance Type - must have GPU if your workloads requires intense graphic processing
      Minutes:
        default: How frequently the Maintenance Window should run its tasks to identify idle DCV instances 
      Tag:
        default: Name of the Tag that will be monitored by the solution 
      TagValue:
        default: Value of the Tag that will be monitored by the solution 
      Action:
        default: Should Stop or Terminate instances when idle?

Mappings:
  DCVServer: ## last update 2021-10-26
    us-east-1:
      "HVM64": ami-045d4d13f489439a2
    us-east-2:
      "HVM64": ami-00bf94f338d214470
    us-west-1:
      "HVM64": ami-00493e347ecc5291c
    us-west-2:
      "HVM64": ami-053f7f184071920af
    eu-north-1:
      "HVM64": ami-075318ab63fd9ea12
    ap-south-1:
      "HVM64": ami-0dac00757d3028cfc
    eu-west-3:
      "HVM64": ami-05a753079a82dff60
    eu-west-2:
      "HVM64": ami-092acaf972e3d5c50
    eu-west-1:
      "HVM64": ami-00f427ee60aff2cef
    ap-northeast-2:
      "HVM64": ami-013ab9c597a4f1b98
    ap-northeast-1:
      "HVM64": ami-058767e4ed89104d4
    ap-southeast-1:
      "HVM64": ami-065a0888182bf6243
    ap-southeast-2:
      "HVM64": ami-0b29148c6c117c6e9
    eu-central-1:
      "HVM64": ami-0593516fc586441b6
    sa-east-1:
      "HVM64": ami-020b5ef016db35c9a
    ca-central-1:
      "HVM64": ami-077008f20bbef0eb4

Resources:
  # SSM Document to start the DCV Monitor inside the instance - must be run every time an instance start
  StartDCVMonitorForConnections:
    Type: AWS::SSM::Document
    Properties: 
      Content: 
        schemaVersion: '2.2'
        description: 'Start a script that will monitor if there is an active session'
        parameters: 
          dcvtopic:
            type: String
            description: "SNS topic ARN"
            default: !Ref DCVTopic
          commands: 
            type: String 
            description: 'monitor for onDisconnect and wait 30min before sending a SNS message to stop the instance'
            default: | 
              $stopWatch = New-Object -TypeName System.Diagnostics.StopWatch;
              $timer = New-TimeSpan -Minutes 1;
              $stopWatch.Start();
              
              $conn = Invoke-Expression -Command 'C:\Program` Files\NICE\DCV\Server\bin\dcv list-connections -j console |ConvertFrom-Json';
              if ($conn) {
                  Write-Host "There are active connections. Exiting.";
                  exit
              } else {
                  Start-Sleep -Seconds 10;
                  while ($stopWatch.Elapsed -le $timer) {
                      $conn = Invoke-Expression -Command 'C:\Program` Files\NICE\DCV\Server\bin\dcv list-connections -j console |ConvertFrom-Json';
                      if (!$conn) {
                          Start-Sleep -Seconds 60;
                      } else {
                          Write-Host "There are active connections. Exiting."
                          exit
                      }
                  }
                  if (!$conn) {
                      Write-Host "No connections for 30min. Sending message to topic.";
                      $instance_id = Get-EC2InstanceMetadata -Category InstanceId;
                      Publish-SNSMessage -TopicArn ${dcvtopic} -Message "{""default"":""${instance_id}""}" -MessageStructure 'json' -Subject 'DCV user disconnected';
                  } else { 
                      Write-Host "Some error occurred and the message was not posted."
                      exit
                  }
              }
        mainSteps:
        - action: aws:runPowerShellScript
          name: 'runPowerShellScript'
          inputs: 
            timeoutSeconds: 2100
            runCommand: 
              - "$dcvtopic=\"{{ dcvtopic }}\""
              - "{{ commands }}"
      DocumentType: Command 
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  SSMMaintenanceWindow:
    Type: AWS::SSM::MaintenanceWindow
    Properties: 
      AllowUnassociatedTargets: False
      Cutoff: 0
      Description: Will verify if there are active connections
      Duration: 1
      Name: DCV-Monitor-Idle-Instances
      Schedule: !Sub cron(0 */${Minutes} * * * ? *)
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  SSMMaintenanceWindowTargets:
    Type: AWS::SSM::MaintenanceWindowTarget
    Properties:
      Description: DCV instances that have been instantiated with the correct tag 
      ResourceType: INSTANCE 
      Targets: 
        - Key: !Sub tag:${Tag} 
          Values: 
            - !Ref TagValue 
      WindowId: !Ref SSMMaintenanceWindow
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  SSMMaintenanceWindowTask:
    Type: AWS::SSM::MaintenanceWindowTask
    Properties:
      Description: using the SSM Document that will verify if there are idle connections in the DCV instance 
      MaxConcurrency: '7' 
      MaxErrors: '7'
      Priority: 0
      Targets: 
        - Key: WindowTargetIds
          Values: 
            - !Ref SSMMaintenanceWindowTargets
      TaskArn: !Ref StartDCVMonitorForConnections
      TaskType: RUN_COMMAND
      WindowId: !Ref SSMMaintenanceWindow
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  # Lambda function to stop a DCV instance based on the instanceId
  DCVStopInstanceLambda:
    Type: AWS::Lambda::Function
    Properties:
      Environment:
        Variables:
          ACTION: !Ref Action
      Code:
        ZipFile: |
          import json
          import boto3
          import os

          ec2 = boto3.client('ec2')
          ddb = boto3.client('dynamodb')

          def stopDcvInstance(e):
              try:
                  stopResponse = ec2.stop_instances(
                      InstanceIds=[
                          e
                      ]
                  )
                  print(stopResponse)
              except: 
                  print('Could not stop the instance %', e)
                  return {
                      'statusCode': '200',
                      'body': 'Could not stop the instance'
                  }

          def terminateDcvInstance(e):
              try:
                  terminateResponse = ec2.terminate_instances(
                      InstanceIds=[
                          e
                      ]
                  )
                  print(terminateResponse)
              except: 
                  print('Could not terminate the instance %', e)
                  return {
                      'statusCode': '200',
                      'body': 'Could not terminate the instance'
                  }                  

          def lambda_handler(event, context):
              action = os.environ['ACTION']
              
              if ('instanceId' in event):
                  # message from SSM is actually a string within a JSON key, need to transform into JSON
                  instance_id = json.loads(event['Records'][0]['Sns']['Message'])['instanceId']        
                  print(instance_id)
              else:
                  instance_id = event['Records'][0]['Sns']['Message']
                  print(instance_id)
              
              # choose between terminating or stopping the instances
              try:
                if (action == 'stopInstance'):
                 stopDcvInstance(instance_id)
                elif (action == 'terminateInstance'):
                  terminateDcvInstance(instance_id)
                else:
                  print('Instance action not known')
              except:
                print('could not call EC2 action. Exiting')
              else: 
                print('action taken: %', action)

            return {
              'statusCode': '200',
              'body': 'Instance stopped succesfully'
            }              

      FunctionName: 'manage-dcv-instance'
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !GetAtt DCVLambdaExecutionRole.Arn
      Runtime: python3.7
      Timeout: 30
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  # IAM Role to attach to the Lambda Function
  # must allow: write to log group and EC2 StopInstance
  DCVLambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Path: /
      AssumeRolePolicyDocument: 
        Statement:
          - Action: 'sts:AssumeRole'
            Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
      Policies:
        - PolicyName: "AllowInvokeFunction"
          PolicyDocument:
            Version: "2012-10-17"
            Statement: 
              - Effect: "Allow" 
                Action: "lambda:InvokeFunction"
                Resource: "*"
        - PolicyName: 'AllowStopInstance'
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: 
                  - "ec2:StopInstances"
                  - "ec2:TerminateInstances"
                Resource: !Sub arn:${AWS::Partition}:ec2:${AWS::Region}:${AWS::AccountId}:instance/* # shouldn't use * anywhere
                #Condition:
                #  StringEquals:
                #    - !Sub '"aws:ResourceTag/"${Tag}: ${TagValue} ' ##TODO:  Still trying to add tags as a condition and implement least privilege
        - PolicyName: 'LambdaExecutionPolicy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement: 
              - Effect: Allow
                Action: 
                  - 'logs:CreateLogGroup'
                Resource:
                  - !Sub arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:*
              - Effect: Allow
                Action: 
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource:
                  - !Sub arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/*:*
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  # Allow SNS to call the DCVStopInstance function
  DCVLambdaPermissions:
    Type: AWS::Lambda::Permission
    Properties:
      Action: "lambda:InvokeFunction"
      FunctionName: !GetAtt DCVStopInstanceLambda.Arn
      Principal: "sns.amazonaws.com"
      SourceArn: !Ref DCVTopic
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  # SNS Topic to orchestrate the integration
  DCVTopic:
    Type: AWS::SNS::Topic
    Properties:
      DisplayName: ManageDCVInstances
      TopicName: ManageDCVInstances
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  # SNS Subscription for the DCVStopInstance lambda
  DCVLambdaSubscription:
    Type: AWS::SNS::Subscription
    Properties:
      TopicArn: !Ref DCVTopic
      Protocol: lambda
      Endpoint: !GetAtt DCVStopInstanceLambda.Arn
      Region: !Ref 'AWS::Region'
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  DCVInstanceRole:
    Type: AWS::IAM::Role 
    Properties:
      Description: allows an EC2 Instance to be managed by SSM and send messages to an Amazon SNS topic 
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement: 
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
                - ssm.amazonaws.com
            Action: 
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: AllowPublishMessage
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: 
                  - 'sns:Publish' 
                Resource: !Ref DCVTopic
        - PolicyName: AccessDCVLicenseBucket
          PolicyDocument:
            Version: '2012-10-17'
            Statement: 
              - Effect: Allow
                Action: 
                  - 's3:GetObject' 
                Resource: !Sub arn:${AWS::Partition}:s3:::dcv-license.${AWS::Region}/*
        - PolicyName: SSMManageInstance
          PolicyDocument: 
            Version: '2012-10-17'
            Statement: 
              - Effect: Allow
                Action: 
                  - 'ssm:GetDeployablePatchSnapshotForInstance'
                  - 'ssm:GetManifest'
                  - 'ssm:ListAssociations'
                  - 'ssm:PutInventory'
                  - 'ssm:PutConfigurePackageResult'
                  - 'ssm:UpdateInstanceInformation'
                Resource: "*" # those actions only support all resources. TODO: scope down to the account itself.
                # Add Condition to scope down
              - Effect: Allow
                Action: 
                  - 'ssm:GetDocument'
                  - 'ssm:DescribeDocument'
                  - 'ssm:GetParameter'
                  - 'ssm:GetParameters'
                Resource: 
                  - !Sub arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:document/*
                  - !Sub arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:parameter/*
              - Effect: Allow
                Action: 
                  - 'ssm:DescribeAssociation'
                  - 'ssm:ListInstanceAssociations'
                  - 'ssm:PutComplianceItems'
                  - 'ssm:UpdateAssociationStatus'
                  - 'ssm:UpdateInstanceAssociationStatus'
                Resource: 
                  - !Sub arn:${AWS::Partition}:ec2:${AWS::Region}:${AWS::AccountId}:instance/*
                  - !Sub arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:managed-instance/*
                  - !Sub arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:document/DCV-Identify-Active-Connections
                  - !Sub arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:association/*
                  - !Sub arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:parameter/*
              - Effect: Allow
                Action: 
                  - 'ssmmessages:CreateControlChannel'
                  - 'ssmmessages:CreateDataChannel'
                  - 'ssmmessages:OpenControlChannel'
                  - 'ssmmessages:OpenDataChannel'
                Resource: "*" # those actions only support all resources. TODO: scope down to the account itself.
              - Effect: Allow
                Action:
                  - 'ec2messages:AcknowledgeMessage'
                  - 'ec2messages:DeleteMessage'
                  - 'ec2messages:FailMessage'
                  - 'ec2messages:GetEndpoint'
                  - 'ec2messages:GetMessages'
                  - 'ec2messages:SendReply'
                Resource: "*" # those actions only support all resources. TODO: scope down to the account itself.
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  DCVInstanceProfile:
    Type: AWS::IAM::InstanceProfile 
    Properties:
      Path: / 
      Roles:
        - !Ref DCVInstanceRole
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  DCVLaunchTemplate:
    Type: AWS::EC2::LaunchTemplate 
    Properties:
      LaunchTemplateData:
        IamInstanceProfile:
          Arn: !GetAtt DCVInstanceProfile.Arn 
        ImageId: !FindInMap [DCVServer, !Ref 'AWS::Region', HVM64] 
        InstanceType: !Ref Instance 
        KeyName: !Ref EC2KeyPair
        SecurityGroupIds: 
          - !Ref SecurityGroup 
        TagSpecifications: 
          - ResourceType: instance
            Tags:
              - Key: !Ref Tag
                Value: !Ref TagValue
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

Outputs:
  SNSTopic:
    Description: SNS Topic used to orchestrate the solution
    Value: !GetAtt DCVTopic.TopicName
  LambdaFunction: 
    Description: Lambda function that manage stopping the instances
    Value: !Ref DCVStopInstanceLambda
  SSMDocument: 
    Description: the SSM Document that verifies if there are active connections in the DCV instance
    Value: !Ref StartDCVMonitorForConnections
  LaunchTemplate:
    Description: the EC2 Launch Template created for the solution
    Value: !Ref DCVLaunchTemplate
  SSMDocumentName:
    Description: the AWS System Manager Document created 
    Value: !Ref StartDCVMonitorForConnections