Deploy Jupyter notebook container with Amazon ECS

AWSTemplateFormatVersion: '2010-09-09'
Description: This stack deploys a large AWS VPC with internet access
Mappings:
  # Hard values for the subnet masks. These masks define
  # the range of internal IP addresses that can be assigned.
  # The VPC can have all IP's from 10.0.0.0 to 10.0.255.255
  # There are four subnets which cover the ranges:
  #
  # 10.0.0.0 - 10.0.63.255 (16384 IP addresses)
  # 10.0.64.0 - 10.0.127.255 (16384 IP addresses)
  # 10.0.128.0 - 10.0.191.255 (16384 IP addresses)
  # 10.0.192.0 - 10.0.255.0 (16384 IP addresses)
  #
  SubnetConfig:
    VPC:
      CIDR: '10.0.0.0/16'
    PublicOne:
      CIDR: '10.0.0.0/18'
    PublicTwo:
      CIDR: '10.0.64.0/18'
    PrivateOne:
      CIDR: '10.0.128.0/18'
    PrivateTwo:
      CIDR: '10.0.192.0/18'
Resources:
  # VPC in which containers will be networked.
  # It has two public subnets, and two private subnets.
  # We distribute the subnets across the first two available subnets
  # for the region, for high availability.
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      EnableDnsSupport: true
      EnableDnsHostnames: true
      CidrBlock: !FindInMap ['SubnetConfig', 'VPC', 'CIDR']

  # Two public subnets, where containers can have public IP addresses
  PublicSubnetOne:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 0
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref 'VPC'
      CidrBlock: !FindInMap ['SubnetConfig', 'PublicOne', 'CIDR']
      MapPublicIpOnLaunch: true
  PublicSubnetTwo:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 1
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref 'VPC'
      CidrBlock: !FindInMap ['SubnetConfig', 'PublicTwo', 'CIDR']
      MapPublicIpOnLaunch: true

  # Two private subnets where containers will only have private
  # IP addresses, and will only be reachable by other members of the
  # VPC
  PrivateSubnetOne:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 0
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref 'VPC'
      CidrBlock: !FindInMap ['SubnetConfig', 'PrivateOne', 'CIDR']
  PrivateSubnetTwo:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 1
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref 'VPC'
      CidrBlock: !FindInMap ['SubnetConfig', 'PrivateTwo', 'CIDR']

  # Setup networking resources for the public subnets. Containers
  # in the public subnets have public IP addresses and the routing table
  # sends network traffic via the internet gateway.
  InternetGateway:
    Type: AWS::EC2::InternetGateway
  GatewayAttachement:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref 'VPC'
      InternetGatewayId: !Ref 'InternetGateway'
  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref 'VPC'
  PublicRoute:
    Type: AWS::EC2::Route
    DependsOn: GatewayAttachement
    Properties:
      RouteTableId: !Ref 'PublicRouteTable'
      DestinationCidrBlock: '0.0.0.0/0'
      GatewayId: !Ref 'InternetGateway'
  PublicSubnetOneRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnetOne
      RouteTableId: !Ref PublicRouteTable
  PublicSubnetTwoRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnetTwo
      RouteTableId: !Ref PublicRouteTable

  # Setup networking resources for the private subnets. Containers
  # in these subnets have only private IP addresses, and must use a NAT
  # gateway to talk to the internet. We launch two NAT gateways, one for
  # each private subnet.
  NatGatewayOneAttachment:
    Type: AWS::EC2::EIP
    DependsOn: GatewayAttachement
    Properties:
        Domain: vpc
  NatGatewayTwoAttachment:
    Type: AWS::EC2::EIP
    DependsOn: GatewayAttachement
    Properties:
        Domain: vpc
  NatGatewayOne:
    Type: AWS::EC2::NatGateway
    Properties:
      AllocationId: !GetAtt NatGatewayOneAttachment.AllocationId
      SubnetId: !Ref PublicSubnetOne
  NatGatewayTwo:
    Type: AWS::EC2::NatGateway
    Properties:
      AllocationId: !GetAtt NatGatewayTwoAttachment.AllocationId
      SubnetId: !Ref PublicSubnetTwo
  PrivateRouteTableOne:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref 'VPC'
  PrivateRouteOne:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref PrivateRouteTableOne
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId: !Ref NatGatewayOne
  PrivateRouteTableOneAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTableOne
      SubnetId: !Ref PrivateSubnetOne
  PrivateRouteTableTwo:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref 'VPC'
  PrivateRouteTwo:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref PrivateRouteTableTwo
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId: !Ref NatGatewayTwo
  PrivateRouteTableTwoAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTableTwo
      SubnetId: !Ref PrivateSubnetTwo

Outputs:
  VpcId:
    Description: The ID of the VPC that this stack is deployed in
    Value: !Ref 'VPC'
  PublicSubnetIds:
    Description: Comma seperated list of public facing subnets that have
                 a direct internet connection as long as you assign a public IP
    Value: !Sub '${PublicSubnetOne},${PublicSubnetTwo}'
  PrivateSubnetIds:
    Description: Comma seperated list of private subnets that use a NAT
                 gateway for internet access.
    Value: !Sub '${PrivateSubnetOne},${PrivateSubnetTwo}'

AWSTemplateFormatVersion: '2010-09-09'
Description: ECS Cluster optimized for machine learning inference workloads
Parameters:
  InstanceType:
    Type: String
    Default: inf2.8xlarge
    Description: Class of AWS Inferentia instance used to run containers.
    AllowedValues: [inf1.xlarge, inf1.2xlarge, inf1.6xlarge, inf1.24xlarge,
      inf2.xlarge, inf2.8xlarge, inf2.24xlarge, inf2.48xlarge]
    ConstraintDescription: Please choose a valid instance type.
  DesiredCapacity:
    Type: Number
    Default: '0'
    Description: Number of instances to initially launch in your ECS cluster.
  MaxSize:
    Type: Number
    Default: '2'
    Description: Maximum number of instances that can be launched in your ECS cluster.
  ECSAMI:
    Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
    Default: /aws/service/ecs/optimized-ami/amazon-linux-2/inf/recommended/image_id
    Description: The Amazon Machine Image ID used for the cluster, leave it as the default value to get the latest Inferentia optimized AMI
  VpcId:
    Type: AWS::EC2::VPC::Id
    Description: VPC ID where the ECS cluster is launched
  SubnetIds:
    Type: List<AWS::EC2::Subnet::Id>
    Description: List of subnet IDs where the EC2 instances will be launched

Resources:
  # Cluster that keeps track of container deployments
  ECSCluster:
    Type: AWS::ECS::Cluster
    Properties:
      ClusterSettings:
        - Name: containerInsights
          Value: enabled

  # Autoscaling group. This launches the actual EC2 instances that will register
  # themselves as members of the cluster, and run the docker containers.
  ECSAutoScalingGroup:
    Type: AWS::AutoScaling::AutoScalingGroup
    DependsOn:
      # This is to ensure that the ASG gets deleted first before these
      # resources, when it comes to stack teardown.
      - ECSCluster
      - EC2Role
    Properties:
      VPCZoneIdentifier:
        - !Select [ 0, !Ref SubnetIds ]
        - !Select [ 1, !Ref SubnetIds ]
      LaunchTemplate:
        LaunchTemplateId: !Ref ContainerInstances
        Version: !GetAtt ContainerInstances.LatestVersionNumber
      MinSize: 0
      MaxSize: !Ref MaxSize
      DesiredCapacity: !Ref DesiredCapacity
      NewInstancesProtectedFromScaleIn: true
    UpdatePolicy:
      AutoScalingReplacingUpdate:
        WillReplace: 'true'

  # The config for each instance that is added to the cluster
  ContainerInstances:
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateData:
        ImageId: !Ref ECSAMI
        InstanceType: !Ref InstanceType
        IamInstanceProfile:
          Name: !Ref EC2InstanceProfile
        SecurityGroupIds:
          - !Ref ContainerHostSecurityGroup
        UserData:
          # This injected configuration file is how the EC2 instance
          # knows which ECS cluster on your AWS account it should be joining
          Fn::Base64: !Sub |
            #!/bin/bash
            echo ECS_CLUSTER=${ECSCluster} >> /etc/ecs/ecs.config
        BlockDeviceMappings:
          - DeviceName: "/dev/xvda"
            Ebs:
              VolumeSize: 50
              VolumeType: gp3
        # Disable IMDSv1, and require IMDSv2
        MetadataOptions:
          HttpEndpoint: enabled
          HttpTokens: required
  EC2InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: /
      Roles:
        - !Ref EC2Role

  # Custom resource that force destroys the ASG. This cleans up EC2 instances that had
  # managed termination protection enabled, but which are not yet released.
  # This is necessary because ECS does not immediately release an EC2 instance from termination
  # protection as soon as the instance is no longer running tasks. There is a cooldown delay.
  # In the case of tearing down the CloudFormation stack, CloudFormation will delete the
  # AWS::ECS::Service and immediately move on to tearing down the AWS::ECS::Cluster, disconnecting
  # the AWS::AutoScaling::AutoScalingGroup from ECS management too fast, before ECS has a chance
  # to asynchronously turn off managed instance protection on the EC2 instances.
  # This will leave some EC2 instances stranded in a state where they are protected from scale-in forever.
  # This then blocks the AWS::AutoScaling::AutoScalingGroup from cleaning itself up.
  # The custom resource function force destroys the autoscaling group when tearing down the stack,
  # avoiding the issue of protected EC2 instances that can never be cleaned up.
  CustomAsgDestroyerFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: !Sub |
          const { AutoScalingClient, DeleteAutoScalingGroupCommand } = require("@aws-sdk/client-auto-scaling");
          const autoscaling = new AutoScalingClient({ region: '${AWS::Region}' });
          const response = require('cfn-response');

          exports.handler = async function(event, context) {
            console.log(event);

            if (event.RequestType !== "Delete") {
              await response.send(event, context, response.SUCCESS);
              return;
            }

            const input = {
              AutoScalingGroupName: '${ECSAutoScalingGroup}',
              ForceDelete: true
            };
            const command = new DeleteAutoScalingGroupCommand(input);
            const deleteResponse = await autoscaling.send(command);
            console.log(deleteResponse);

            await response.send(event, context, response.SUCCESS);
          };
      Handler: index.handler
      Runtime: nodejs20.x
      Timeout: 30
      Role: !GetAtt CustomAsgDestroyerRole.Arn

  # The role used by the ASG destroyer
  CustomAsgDestroyerRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      ManagedPolicyArns:
        # https://docs.aws.amazon.com/aws-managed-policy/latest/reference/AWSLambdaBasicExecutionRole.html
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: allow-to-delete-autoscaling-group
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: autoscaling:DeleteAutoScalingGroup
                Resource: !Sub arn:aws:autoscaling:${AWS::Region}:${AWS::AccountId}:autoScalingGroup:*:autoScalingGroupName/${ECSAutoScalingGroup}

  CustomAsgDestroyer:
    Type: Custom::AsgDestroyer
    DependsOn:
      - CapacityProviderAssociation
    Properties:
      ServiceToken: !GetAtt CustomAsgDestroyerFunction.Arn
      Region: !Ref "AWS::Region"

  # Create an ECS capacity provider to attach the ASG to the ECS cluster
  # so that it autoscales as we launch more containers
  CapacityProvider:
    Type: AWS::ECS::CapacityProvider
    Properties:
      AutoScalingGroupProvider:
        AutoScalingGroupArn: !Ref ECSAutoScalingGroup
        ManagedScaling:
          InstanceWarmupPeriod: 60
          MinimumScalingStepSize: 1
          MaximumScalingStepSize: 100
          Status: ENABLED
          # Percentage of cluster reservation to try to maintain
          TargetCapacity: 100
        ManagedTerminationProtection: ENABLED
        ManagedDraining: ENABLED

  # Create a cluster capacity provider assocation so that the cluster
  # will use the capacity provider
  CapacityProviderAssociation:
    Type: AWS::ECS::ClusterCapacityProviderAssociations
    Properties:
      CapacityProviders:
        - !Ref CapacityProvider
      Cluster: !Ref ECSCluster
      DefaultCapacityProviderStrategy:
        - Base: 0
          CapacityProvider: !Ref CapacityProvider
          Weight: 1

  # A security group for the EC2 hosts that will run the containers.
  # This can be used to limit incoming traffic to or outgoing traffic
  # from the container's host EC2 instance.
  ContainerHostSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Access to the EC2 hosts that run containers
      VpcId: !Ref VpcId

  # Role for the EC2 hosts. This allows the ECS agent on the EC2 hosts
  # to communciate with the ECS control plane, as well as download the docker
  # images from ECR to run on your host.
  EC2Role:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: [ec2.amazonaws.com]
            Action: ['sts:AssumeRole']
      Path: /

      # See reference: https://docs.aws.amazon.com/AmazonECS/latest/developerguide/security-iam-awsmanpol.html#security-iam-awsmanpol-AmazonEC2ContainerServiceforEC2Role
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role

  # This is a role which is used within Fargate to allow the Fargate agent
  # to download images, and upload logs.
  ECSTaskExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: [ecs-tasks.amazonaws.com]
            Action: ['sts:AssumeRole']
            Condition:
              ArnLike:
                aws:SourceArn: !Sub arn:aws:ecs:${AWS::Region}:${AWS::AccountId}:*
              StringEquals:
                aws:SourceAccount: !Ref AWS::AccountId
      Path: /

      # This role enables basic features of ECS. See reference:
      # https://docs.aws.amazon.com/AmazonECS/latest/developerguide/security-iam-awsmanpol.html#security-iam-awsmanpol-AmazonECSTaskExecutionRolePolicy
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy

Outputs:
  ClusterName:
    Description: The ECS cluster into which to launch resources
    Value: !Ref ECSCluster
  ECSTaskExecutionRole:
    Description: The role used to start up a task
    Value: !Ref ECSTaskExecutionRole
  CapacityProvider:
    Description: The cluster capacity provider that the service should use
                 to request capacity when it wants to start up a task
    Value: !Ref CapacityProvider

AWSTemplateFormatVersion: '2010-09-09'
Description: An example service that deploys an Jupyter Lab notebook
             with AWS Neuron support for machine learning.

Parameters:
  VpcId:
    Type: String
    Description: The VPC that the service is running inside of
  PublicSubnetIds:
    Type: List<AWS::EC2::Subnet::Id>
    Description: List of public facing subnets
  PrivateSubnetIds:
    Type: List<AWS::EC2::Subnet::Id>
    Description: List of private subnets
  ClusterName:
    Type: String
    Description: The name of the ECS cluster into which to launch capacity.
  ECSTaskExecutionRole:
    Type: String
    Description: The role used to start up an ECS task
  CapacityProvider:
    Type: String
    Description: The cluster capacity provider that the service should use
                 to request capacity when it wants to start up a task
  ServiceName:
    Type: String
    Default: jupyter
    Description: A name for the service
  ImageUrl:
    Type: String
    Description: The URL of a Juypter notebook container to run
  ContainerCpu:
    Type: Number
    Default: 10240
    Description: How much CPU to give the container. 1024 is 1 CPU
  ContainerMemory:
    Type: Number
    Default: 32768
    Description: How much memory in megabytes to give the container
  MyIp:
    Type: String
    Default: 0.0.0.0/0
    Description: The IP addresses that you want to accept traffic from.
                 Default accepts traffic from anywhere on the internet.

Resources:

  # The task definition. This is a simple metadata description of what
  # container to run, and what resource requirements it has.
  TaskDefinition:
    Type: AWS::ECS::TaskDefinition
    DependsOn:
      - TaskAccessToJupyterToken
    Properties:
      Family: !Ref ServiceName
      NetworkMode: awsvpc
      RequiresCompatibilities:
        - EC2
      ExecutionRoleArn: !Ref ECSTaskExecutionRole
      ContainerDefinitions:
        - Name: jupyter
          Cpu: !Ref ContainerCpu
          MemoryReservation: !Ref ContainerMemory
          Image: !Ref ImageUrl
          EntryPoint:
            - '/bin/sh'
            - '-c'
          Command:
            - '/opt/conda/bin/jupyter-lab --allow-root --ServerApp.token=${JUPYTER_TOKEN} --ServerApp.ip=*'
          Secrets:
            - Name: JUPYTER_TOKEN
              ValueFrom: !Ref JupyterToken
          PortMappings:
            - ContainerPort: 8888
              HostPort: 8888
          MountPoints:
            - SourceVolume: efs-volume
              ContainerPath: /home
          LogConfiguration:
            LogDriver: 'awslogs'
            Options:
              mode: non-blocking
              max-buffer-size: 25m
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: !Ref ServiceName
          LinuxParameters:
            Devices:
              # Ensure that the AWS Neuron SDK inside the container
              # can access the underlying host device provided by AWS Inferentia
              - ContainerPath: /dev/neuron0
                HostPath: /dev/neuron0
                Permissions:
                  - read
                  - write
            Capabilities:
              Add:
                - "IPC_LOCK"
      Volumes:
        - Name: efs-volume
          EFSVolumeConfiguration:
            FilesystemId: !Ref EFSFileSystem
            RootDirectory: /
            TransitEncryption: ENABLED

  # The secret token used to protect the Jupyter notebook
  JupyterToken:
    Type: AWS::SecretsManager::Secret
    Properties:
      GenerateSecretString:
        PasswordLength: 30
        ExcludePunctuation: true

  # Attach a policy to the task execution role, which grants
  # the ECS agent the ability to fetch the Jupyter notebook
  # secret token on behalf of the task.
  TaskAccessToJupyterToken:
    Type: AWS::IAM::Policy
    Properties:
      Roles:
        - !Ref ECSTaskExecutionRole
      PolicyName: AccessJupyterToken
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - secretsmanager:DescribeSecret
              - secretsmanager:GetSecretValue
            Resource: !Ref JupyterToken

  # Attach a policy to the task execution role which allows the
  # task to access the EFS filesystem that provides durable storage
  # for the task.
  TaskAccessToFilesystem:
    Type: AWS::IAM::Policy
    Properties:
      Roles:
        - !Ref ECSTaskExecutionRole
      PolicyName: EFSAccess
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - elasticfilesystem:ClientMount
              - elasticfilesystem:ClientWrite
              - elasticfilesystem:DescribeMountTargets
              - elasticfilesystem:DescribeFileSystems
            Resource: !GetAtt EFSFileSystem.Arn

  # The service. The service is a resource which allows you to run multiple
  # copies of a type of task, and gather up their logs and metrics, as well
  # as monitor the number of running tasks and replace any that have crashed
  Service:
    Type: AWS::ECS::Service
    DependsOn: PublicLoadBalancerListener
    Properties:
      ServiceName: !Ref ServiceName
      Cluster: !Ref ClusterName
      PlacementStrategies:
        - Field: attribute:ecs.availability-zone
          Type: spread
        - Field: cpu
          Type: binpack
      CapacityProviderStrategy:
        - Base: 0
          CapacityProvider: !Ref CapacityProvider
          Weight: 1
      DeploymentConfiguration:
        MaximumPercent: 200
        MinimumHealthyPercent: 75
      DesiredCount: 1
      NetworkConfiguration:
        AwsvpcConfiguration:
          SecurityGroups:
            - !Ref ServiceSecurityGroup
          Subnets:
            - !Select [ 0, !Ref PrivateSubnetIds ]
            - !Select [ 1, !Ref PrivateSubnetIds ]
      TaskDefinition: !Ref TaskDefinition
      LoadBalancers:
        - ContainerName: jupyter
          ContainerPort: 8888
          TargetGroupArn: !Ref JupyterTargetGroup

  # Because we are launching tasks in AWS VPC networking mode
  # the tasks themselves also have an extra security group that is unique
  # to them. This is a unique security group just for this service,
  # to control which things it can talk to, and who can talk to it
  ServiceSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: !Sub Access to service ${ServiceName}
      VpcId: !Ref VpcId

  # This log group stores the stdout logs from this service's containers
  LogGroup:
    Type: AWS::Logs::LogGroup

   # Keeps track of the list of tasks running on EC2 instances
  JupyterTargetGroup:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    Properties:
      HealthCheckIntervalSeconds: 6
      HealthCheckPath: /api
      HealthCheckProtocol: HTTP
      HealthCheckTimeoutSeconds: 5
      HealthyThresholdCount: 2
      TargetType: ip
      Port: 8888
      Protocol: HTTP
      UnhealthyThresholdCount: 10
      VpcId: !Ref VpcId
      TargetGroupAttributes:
        - Key: deregistration_delay.timeout_seconds
          Value: 0

  # A public facing load balancer, this is used as ingress for
  # public facing internet traffic. The traffic is forwarded
  # down to the Juypter notebook where ever it is currently hosted
  # on whichever machine Amazon ECS placed it on.
  PublicLoadBalancerSG:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Access to the public facing load balancer
      VpcId: !Ref VpcId
      SecurityGroupIngress:
        # Allow access to ALB from the specified IP address
        - CidrIp: !Ref MyIp
          IpProtocol: -1
  PublicLoadBalancer:
    Type: AWS::ElasticLoadBalancingV2::LoadBalancer
    Properties:
      Scheme: internet-facing
      LoadBalancerAttributes:
      - Key: idle_timeout.timeout_seconds
        Value: '30'
      Subnets:
        # The load balancer is placed into the public subnets, so that traffic
        # from the internet can reach the load balancer directly via the internet gateway
        - !Select [ 0, !Ref PublicSubnetIds ]
        - !Select [ 1, !Ref PublicSubnetIds ]
      SecurityGroups:
        - !Ref PublicLoadBalancerSG
  PublicLoadBalancerListener:
    Type: AWS::ElasticLoadBalancingV2::Listener
    Properties:
      DefaultActions:
        - Type: 'forward'
          ForwardConfig:
            TargetGroups:
              - TargetGroupArn: !Ref JupyterTargetGroup
                Weight: 100
      LoadBalancerArn: !Ref 'PublicLoadBalancer'
      Port: 80
      Protocol: HTTP

  # The Jupyter services' security group allows inbound
  # traffic from the public facing ALB
  JupyterIngressFromPublicALB:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      Description: Ingress from the public ALB
      GroupId: !Ref 'ServiceSecurityGroup'
      IpProtocol: -1
      SourceSecurityGroupId: !Ref 'PublicLoadBalancerSG'

  # Filesystem that provides durable storage for the notebook
  EFSFileSystem:
    Type: AWS::EFS::FileSystem
    Properties:
      Encrypted: true
      PerformanceMode: generalPurpose
      ThroughputMode: bursting

  # Mount target allows usage of the EFS inside of subnet one
  EFSMountTargetOne:
    Type: AWS::EFS::MountTarget
    Properties:
      FileSystemId: !Ref EFSFileSystem
      SubnetId: !Select [ 0, !Ref PrivateSubnetIds ]
      SecurityGroups:
        - !Ref EFSFileSystemSecurityGroup

  # Mount target allows usage of the EFS inside of subnet two
  EFSMountTargetTwo:
    Type: AWS::EFS::MountTarget
    Properties:
      FileSystemId: !Ref EFSFileSystem
      SubnetId: !Select [ 1, !Ref PrivateSubnetIds ]
      SecurityGroups:
        - !Ref EFSFileSystemSecurityGroup

  # This security group is used by the mount targets so
  # that they will allow inbound NFS connections from
  # the ECS tasks that we launch
  EFSFileSystemSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Security group for EFS file system
      VpcId: !Ref VpcId
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 2049
          ToPort: 2049
          SourceSecurityGroupId: !Ref ServiceSecurityGroup

Outputs:
  LoadBalancerUrl:
    Description: The URL at which you can access the application
    Value: !GetAtt PublicLoadBalancer.DNSName
  Secret:
    Description: The ARN of the secret that was created to protect the Juypter Lab
    Value: !Ref JupyterToken

AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: Parent stack that deploys VPC, Amazon ECS cluster with AWS Inferentia capacity
             and then deploys a JupyterLab IDE (latest Jupyter Notebook) with AWS Neuron SDK for machine
             learning projects

Parameters:
  ImageUrl:
    Type: String
    Description: The URL of the Jupyter image that you built

Resources:

  # The networking configuration. This creates an isolated
  # network specific to this particular environment
  VpcStack:
    Type: AWS::Serverless::Application
    Properties:
      Location: vpc.yml

  # This stack contains cluster wide resources that will be shared
  # by all services that get launched in the stack
  BaseStack:
    Type: AWS::Serverless::Application
    Properties:
      Location: inferentia-cluster.yml
      Parameters:
        VpcId: !GetAtt VpcStack.Outputs.VpcId
        SubnetIds: !GetAtt VpcStack.Outputs.PrivateSubnetIds

  # Deploys the JupyterLab application into the cluster
  JupyterNotebookStack:
    Type: AWS::Serverless::Application
    Properties:
      Location: jupyter-notebook.yml
      Parameters:
        ImageUrl: !Ref ImageUrl
        VpcId: !GetAtt VpcStack.Outputs.VpcId
        PublicSubnetIds: !GetAtt VpcStack.Outputs.PublicSubnetIds
        PrivateSubnetIds: !GetAtt VpcStack.Outputs.PrivateSubnetIds
        ClusterName: !GetAtt BaseStack.Outputs.ClusterName
        ECSTaskExecutionRole: !GetAtt BaseStack.Outputs.ECSTaskExecutionRole
        CapacityProvider: !GetAtt BaseStack.Outputs.CapacityProvider

Outputs:
  JuypterLabUrl:
    Description: The URL at which you can find your Juypter Lab installation
    Value: !GetAtt JupyterNotebookStack.Outputs.LoadBalancerUrl
  Secret:
    Description: The ARN of the secret token that protects your JuypterLab
    Value: !GetAtt JupyterNotebookStack.Outputs.Secret