AWSTemplateFormatVersion: "2010-09-09"
Description:

  "Builds a VPC with an Elastic Inference endpoint; creates an EC2 container instance with Elastic Inference accelerator attached; creates a client EC2 instance to run inference calls from"


Mappings:

  ECSOptimizedAmazonLinux2AMIs:
    us-east-1:
      amiId: "ami-0b16d80945b1a9c7d"
    us-east-2:
      amiId: "ami-0e7c12c1bedd6bf21"
    us-west-2:
      amiId: "ami-0e434a58221275ed4"
    eu-west-1:
      amiId: "ami-02bf9e90a6e30dc74"
    ap-northeast-1:
      amiId: "ami-0e37e42dff65024ae"
    ap-northeast-2:
      amiId: "ami-000fbda700ba8fe9d"

  DeepLearningAmazonLinuxAMIs:
    us-east-1:
      amiId: "ami-01ac4e28da63bac3c"
    us-east-2:
      amiId: "ami-0a13ab27aa8e0e92a"
    us-west-2:
      amiId: "ami-0283132d7b60d70b9"
    eu-west-1:
      amiId: "ami-02273e0d16172dbd1"
    ap-northeast-1:
      amiId: "ami-0e102072c84e36eed"
    ap-northeast-2:
      amiId: "ami-078afb8923fdf1d94"

Parameters:

  AvailabilityZone:
    Type: "String"
    Default: "us-east-1a"
    Description: "Availability zone must be picked from one of the following regions: us-east-1, us-east-2, us-west-2, ap-northeast-1, ap-northeast-2, eu-west-1. Note that only certain AZs support Elastic Inference - which zones exactly varies on the account."

  KeyPairName:
    Type: "String"
    Default: "replace-with-your-regional-keypair-name"
    Description: "Name of an EC2 Key Pair in the selected region (used for SSH): https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-key-pairs.html"
    AllowedPattern: "[-_\\w]+"

  ECSContainerInstanceType:
    Type: "String"
    Default: "c5.large"
    Description: "Type of the EC2 container instance - the instance running TensorFlow Model Server containers"

  ClientInstanceType:
    Type: "String"
    Default: "c5.large"
    Description: "Type of the EC2 instance running client application"

Resources:

  VPC:
    Type: "AWS::EC2::VPC"
    Properties:
      CidrBlock: "10.0.0.0/16"
      EnableDnsSupport: true
      EnableDnsHostnames: true

  Subnet:
    Type: "AWS::EC2::Subnet"
    Properties:
      VpcId:
        "Ref": "VPC"
      CidrBlock: "10.0.0.0/24"
      AvailabilityZone:
        "Ref": "AvailabilityZone"
      MapPublicIpOnLaunch: "true"

  EIEndpointSecurityGroup:
    Type: "AWS::EC2::SecurityGroup"
    Properties:
      GroupName: "ElasticInferenceEndpointSG"
      GroupDescription: "Security Group managing Elastic Inference service endpoint"
      VpcId:
        "Ref": "VPC"
      SecurityGroupIngress:
        - SourceSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          SourceSecurityGroupOwnerId:
            "Ref": "AWS::AccountId"
          IpProtocol: "tcp"
          FromPort: 443
          ToPort: 443
      SecurityGroupEgress:
        - DestinationSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          IpProtocol: "tcp"
          FromPort: 443
          ToPort: 443

  Allow443FromEIIngressRule:
    Type : "AWS::EC2::SecurityGroupIngress"
    Properties:
      IpProtocol: "tcp"
      FromPort: 443
      ToPort: 443
      SourceSecurityGroupId:
        "Ref": "EIEndpointSecurityGroup"
      GroupId:
        Fn::GetAtt:
          - "VPC"
          - "DefaultSecurityGroup"

  Allow22FromAnywhereIngressRule:
    Type: "AWS::EC2::SecurityGroupIngress"
    Properties:
      IpProtocol: "tcp"
      FromPort: 22
      ToPort: 22
      CidrIp: "0.0.0.0/0"
      GroupId:
        Fn::GetAtt:
          - "VPC"
          - "DefaultSecurityGroup"

  ElasticInferenceEndpoint:
    Type: "AWS::EC2::VPCEndpoint"
    Properties:
      ServiceName:
        Fn::Sub: "com.amazonaws.${AWS::Region}.elastic-inference.runtime"
      VpcEndpointType: "Interface"
      VpcId:
        "Ref": "VPC"
      SubnetIds:
        - "Ref": "Subnet"
      SecurityGroupIds:
        - "Ref": "EIEndpointSecurityGroup"
      PrivateDnsEnabled: true

  ECSContainerInstanceRole:
    Type: "AWS::IAM::Role"
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Statement:
          Effect: "Allow"
          Action: "sts:AssumeRole"
          Principal:
            Service:
              - "ec2.amazonaws.com"
      RoleName: "ECSContainerInstanceRole"
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role"

  ECSContainerInstanceProfile:
    Type: "AWS::IAM::InstanceProfile"
    Properties:
      InstanceProfileName: "BlogECSContainerInstanceProfile"
      Path: "/"
      Roles:
        - Ref: "ECSContainerInstanceRole"

  ElasticInferenceConnectPolicy:
    Type: "AWS::IAM::Policy"
    Properties:
      PolicyName: "ElasticInferenceConnectPolicy"
      PolicyDocument:
        Statement:
          - Effect: "Allow"
            Resource: "*"
            Action:
              - "elastic-inference:Connect"
              - "iam:List*"
              - "iam:Get*"
              - "ec2:Describe*"
              - "ec2:Get*"
      Roles:
        - "Ref": "ECSEITaskRole"

  ClientInstanceRole:
    Type: "AWS::IAM::Role"
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Statement:
          Effect: "Allow"
          Action: "sts:AssumeRole"
          Principal:
            Service:
              - "ec2.amazonaws.com"
      RoleName: "ClientInstanceRole"

  ClientInstanceProfile:
    Type: "AWS::IAM::InstanceProfile"
    Properties:
      InstanceProfileName: "BlogClientInstanceProfile"
      Path: "/"
      Roles:
        - Ref: "ClientInstanceRole"

  ECSRegisterAndRunTaskPolicy:
    Type: "AWS::IAM::Policy"
    Properties:
      PolicyName: "ECSRegisterAndRunTaskPolicy"
      PolicyDocument:
        Statement:
          - Effect: "Allow"
            Resource: "*"
            Action:
              - "ecs:RegisterTaskDefinition"
              - "ecs:RunTask"
          - Effect: "Allow"
            Resource:
              - Fn::Sub: "arn:aws:iam::${AWS::AccountId}:role/ecs-ei-task-role"
              - Fn::Sub: "arn:aws:iam::${AWS::AccountId}:role/ecsTaskExecutionRole"
            Action:
              - "iam:PassRole"
      Roles:
        - "Ref": "ClientInstanceRole"

  ContainerInstanceSecurityGroup:
    Type: "AWS::EC2::SecurityGroup"
    Properties:
      GroupName: "ContainerInstanceSG"
      GroupDescription: "Security Group managing EC2 container instance with Elastic Inference service endpoint"
      VpcId:
        "Ref": "VPC"
      SecurityGroupIngress:
        - SourceSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          SourceSecurityGroupOwnerId:
            "Ref": "AWS::AccountId"
          IpProtocol: "tcp"
          FromPort: 8500
          ToPort: 8501
        - SourceSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          SourceSecurityGroupOwnerId:
            "Ref": "AWS::AccountId"
          IpProtocol: "tcp"
          FromPort: 9000
          ToPort: 9001
      SecurityGroupEgress:
        - DestinationSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          IpProtocol: "tcp"
          FromPort: 443
          ToPort: 443

  ClientInstanceSecurityGroup:
    Type: "AWS::EC2::SecurityGroup"
    Properties:
      GroupName: "ClientInstanceSG"
      GroupDescription: "Security Group managing client EC2 instance"
      VpcId:
        "Ref": "VPC"
      SecurityGroupEgress:
        - DestinationSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          IpProtocol: "tcp"
          FromPort: 8500
          ToPort: 8501
        - DestinationSecurityGroupId:
            Fn::GetAtt:
              - "VPC"
              - "DefaultSecurityGroup"
          IpProtocol: "tcp"
          FromPort: 9000
          ToPort: 9001

  ECSContainerInstance:
    Type: "AWS::EC2::Instance"
    Properties:
      ImageId:
        Fn::FindInMap:
          - "ECSOptimizedAmazonLinux2AMIs"
          - "Ref": "AWS::Region"
          - "amiId"
      InstanceType:
        "Ref": "ECSContainerInstanceType"
      KeyName:
        "Ref": "KeyPairName"
      SubnetId:
        "Ref": "Subnet"
      IamInstanceProfile:
        "Ref": "ECSContainerInstanceProfile"
      PrivateIpAddress: "10.0.0.42"
      Tags:
        - Key: "Name"
          Value: "EI/ECS blog - EC2 container instance"
  ClientInstance:
    Type: "AWS::EC2::Instance"
    DependsOn: "ECSContainerInstance"
    Properties:
      ImageId:
        Fn::FindInMap:
          - "DeepLearningAmazonLinuxAMIs"
          - "Ref": "AWS::Region"
          - "amiId"
      IamInstanceProfile:
        "Ref": "ClientInstanceProfile"
      InstanceType:
        "Ref": "ClientInstanceType"
      KeyName:
        "Ref": "KeyPairName"
      SubnetId:
        "Ref": "Subnet"
      Tags:
        - Key: "Name"
          Value: "EI/ECS blog - Client instance"
      UserData:
        Fn::Base64:
          Fn::Sub: |
            #!/bin/bash

            set -e

            /usr/bin/pip3 install awscli --upgrade

            cat <<TASK_DEF > /tmp/task_def.json
            {
              "executionRoleArn": "arn:aws:iam::${AWS::AccountId}:role/ecsTaskExecutionRole",
              "inferenceAccelerators": [
              {
                "deviceName": "device1",
                "deviceType": "eia1.medium"
              }
              ],
              "containerDefinitions": [
              {
                "entryPoint": [
                  "bash",
                  "-c",
                  "apt-get update; apt-get --assume-yes install unzip; wget https://s3-us-west-2.amazonaws.com/aws-tf-serving-ei-example/ssd_resnet.zip -P \${!MODEL_BASE_PATH}/\${!MODEL_NAME}/; unzip -j \${!MODEL_BASE_PATH}/\${!MODEL_NAME}/ssd_resnet.zip -d \${!MODEL_BASE_PATH}/\${!MODEL_NAME}/1/; /usr/bin/tf_serving_entrypoint.sh"
                ],
                "portMappings": [
                {
                  "hostPort": 8500,
                  "protocol": "tcp",
                  "containerPort": 8500
                },
                {
                  "hostPort": 8501,
                  "protocol": "tcp",
                  "containerPort": 8501
                }
                ],
                "cpu": 0,
                "environment": [
                {
                  "name": "KMP_SETTINGS",
                  "value": "0"
                },
                {
                  "name": "TENSORFLOW_INTRA_OP_PARALLELISM",
                  "value": "2"
                },
                {
                  "name": "MODEL_NAME",
                  "value": "ssdresnet"
                },
                {
                  "name": "KMP_AFFINITY",
                  "value": "granularity=fine,compact,1,0"
                },
                {
                  "name": "MODEL_BASE_PATH",
                  "value": "/models"
                },
                {
                  "name": "KMP_BLOCKTIME",
                  "value": "0"
                },
                {
                  "name": "TENSORFLOW_INTER_OP_PARALLELISM",
                  "value": "2"
                },
                {
                  "name": "OMP_NUM_THREADS",
                  "value": "1"
                }
                ],
                "mountPoints": [],
                "volumesFrom": [],
                "image": "763104351884.dkr.ecr.${AWS::Region}.amazonaws.com/tensorflow-inference:1.14.0-cpu-py27-ubuntu18.04",
                "essential": true,
                "name": "ubuntu-tfs"
              },
              {
                "entryPoint": [
                  "bash",
                  "-c",
                  "apt-get update; apt-get --assume-yes install unzip; wget https://s3-us-west-2.amazonaws.com/aws-tf-serving-ei-example/ssd_resnet.zip -P /models/ssdresnet/; unzip -j /models/ssdresnet/ssd_resnet.zip -d /models/ssdresnet/1; /usr/bin/tensorflow_model_server --port=\${!GRPC_PORT} --rest_api_port=\${!REST_PORT} --model_name=\${!MODEL_NAME} --model_base_path=\${!MODEL_BASE_PATH}/\${!MODEL_NAME}"
                ],
                "resourceRequirements": [
                {
                  "type": "InferenceAccelerator",
                  "value": "device1"
                }
                ],
                "healthCheck": {
                  "retries": 2,
                  "command": [
                    "CMD-SHELL",
                    "LD_LIBRARY_PATH=/opt/ei_health_check/lib /opt/ei_health_check/bin/health_check"
                  ],
                  "timeout": 15,
                  "interval": 30,
                  "startPeriod": 60
                },
                "portMappings": [
                {
                  "hostPort": 9000,
                  "protocol": "tcp",
                  "containerPort": 9000
                },
                {
                  "hostPort": 9001,
                  "protocol": "tcp",
                  "containerPort": 9001
                }
                ],
                "cpu": 0,
                "environment": [
                {
                  "name": "GRPC_PORT",
                  "value": "9000"
                },
                {
                  "name": "REST_PORT",
                  "value": "9001"
                },
                {
                  "name": "MODEL_NAME",
                  "value": "ssdresnet"
                },
                {
                  "name": "MODEL_BASE_PATH",
                  "value": "/models"
                }
                ],
                "mountPoints": [],
                "volumesFrom": [],
                "image": "763104351884.dkr.ecr.${AWS::Region}.amazonaws.com/tensorflow-inference-eia:1.14.0-cpu-py27-ubuntu16.04",
                "essential": true,
                "name": "ubuntu-tfs-ei"
              }
              ],
              "memory": "3072",
              "taskRoleArn": "arn:aws:iam::${AWS::AccountId}:role/ecs-ei-task-role",
              "family": "ei-ecs-ubuntu-tfs-bridge",
              "requiresCompatibilities": [
                "EC2"
              ],
              "networkMode": "bridge",
              "volumes": [],
              "placementConstraints": []
            }
            TASK_DEF

            /usr/local/bin/aws ecs register-task-definition --cli-input-json file:///tmp/task_def.json --region ${AWS::Region}
            /usr/local/bin/aws ecs run-task --cluster default --task-definition ei-ecs-ubuntu-tfs-bridge --region ${AWS::Region}

            cat <<'OUTER' > /tmp/run_ei_example.sh
            #!/bin/bash


            curl -O https://raw.githubusercontent.com/amikelive/coco-labels/master/coco-labels-paper.txt
            curl -O https://s3.amazonaws.com/amazonei/media/3giraffes.jpg

            cat <<'INNER' > ssd_resnet_client.py

            from __future__ import print_function
            import grpc
            import tensorflow as tf
            from PIL import Image
            import numpy as np
            import time
            import os
            from tensorflow_serving.apis import predict_pb2
            from tensorflow_serving.apis import prediction_service_pb2_grpc

            tf.app.flags.DEFINE_string('server', 'localhost:8500',
                                       'PredictionService host:port')
            tf.app.flags.DEFINE_string('image', '', 'path to image in JPEG format')
            FLAGS = tf.app.flags.FLAGS

            if(FLAGS.image == ''):
              print("Supply an Image using '--image [path/to/image]'")
              exit(1)

            local_coco_classes_txt = "coco-labels-paper.txt"

            # Setting default number of predictions
            NUM_PREDICTIONS = 20

            # Reading coco labels to a list
            with open(local_coco_classes_txt) as f:
              classes = ["No Class"] + [line.strip() for line in f.readlines()]

            def main(_):

              channel = grpc.insecure_channel(FLAGS.server)
              stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

              with Image.open(FLAGS.image) as f:
                f.load()

                # Reading the test image given by the user
                data = np.asarray(f)

                # Setting batch size to 1
                data = np.expand_dims(data, axis=0)

                # Creating a prediction request
                request = predict_pb2.PredictRequest()

                # Setting the model spec name
                request.model_spec.name = 'ssdresnet'

                # Setting up the inputs and tensors from image data
                request.inputs['inputs'].CopyFrom(
                    tf.contrib.util.make_tensor_proto(data, shape=data.shape))

                # Iterating over the predictions. The first inference request can take saveral seconds to complete

                durations = []

                for curpred in range(NUM_PREDICTIONS):
                  if(curpred == 0):
                    print("The first inference request loads the model into the accelerator and can take several seconds to complete. Please standby!")

                  # Start the timer
                  start = time.time()

                  # This is where the inference actually happens
                  result = stub.Predict(request, 60.0)  # 60 secs timeout
                  duration = time.time() - start
                  durations.append(duration)
                  print("Inference %d took %f seconds" % (curpred, duration))

                # Extracting results from output
                outputs = result.outputs
                detection_classes = outputs["detection_classes"]

                # Creating an ndarray from the output TensorProto
                detection_classes = tf.make_ndarray(detection_classes)

                # Creating an ndarray from the detection_scores
                detection_scores = tf.make_ndarray(outputs['detection_scores'])

                # Getting the number of objects detected in the input image from the output of the predictor
                num_detections = int(tf.make_ndarray(outputs["num_detections"])[0])
                print("%d detection[s]" % (num_detections))

                # Getting the class ids from the output and mapping the class ids to class names from the coco labels with associated detection score
                class_label_score = ["%s: %.2f" % (classes[int(detection_classes[0][index])], detection_scores[0][index])
                               for index in range(num_detections)]
                print("SSD Prediction is (label, probability): ", class_label_score)
                print("Latency:")
                for percentile in [95, 50]:
                  print("p%d: %.2f seconds" % (percentile, np.percentile(durations, percentile, interpolation='lower')))

            if __name__ == '__main__':
              tf.app.run()
            INNER

            source activate amazonei_tensorflow_p27
            export SERVER_IP=10.0.0.42

            is_model_available() {

              local -r port=$1
              curl -s http://${!SERVER_IP}:${!port}/v1/models/ssdresnet | jq -r '.model_version_status|.[]|.state'
            }

            while :;
            do
              [ "$(is_model_available 8501)" == "AVAILABLE" ] && [ "$(is_model_available 9001)" == "AVAILABLE" ] && break
              sleep 1
            done

            for PORT in 8500 9000; do
            	python ssd_resnet_client.py --server=${!SERVER_IP}:${!PORT} --image 3giraffes.jpg;
            done

            OUTER

            chown ec2-user:ec2-user /tmp/run_ei_example.sh
            chmod u+x /tmp/run_ei_example.sh
            su -l ec2-user /tmp/run_ei_example.sh


  InternetGateway:
    Type: "AWS::EC2::InternetGateway"

  VPCGatewayAttachment:
    Type: "AWS::EC2::VPCGatewayAttachment"
    Properties:
      InternetGatewayId:
        "Ref": "InternetGateway"
      VpcId:
        "Ref": "VPC"

  MainRouteTable:
    Type: "AWS::EC2::RouteTable"
    Properties:
      VpcId:
        "Ref": "VPC"

  RouteToIG:
    Type: "AWS::EC2::Route"
    DependsOn: "InternetGateway"
    Properties:
      RouteTableId:
        "Ref": "MainRouteTable"
      DestinationCidrBlock: "0.0.0.0/0"
      GatewayId:
        "Ref": "InternetGateway"


  SubnetRouteTableAssociation:
    Type: "AWS::EC2::SubnetRouteTableAssociation"
    Properties:
      SubnetId:
        "Ref": "Subnet"
      RouteTableId:
        "Ref": "MainRouteTable"

  ECSEITaskRole:
    Type: "AWS::IAM::Role"
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Statement:
          Effect: "Allow"
          Action: "sts:AssumeRole"
          Principal:
            Service:
              - "ecs-tasks.amazonaws.com"
      RoleName: "ecs-ei-task-role"

  ECSTaskExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Statement:
          Effect: "Allow"
          Action: "sts:AssumeRole"
          Principal:
            Service:
              - "ecs-tasks.amazonaws.com"
      RoleName: "ecsTaskExecutionRole"
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy"
