AWSTemplateFormatVersion: 2010-09-09
Parameters:
  VpcId:
    Type: 'AWS::EC2::VPC::Id'
    Description: VpcId of your existing Virtual Private Cloud (VPC)
    ConstraintDescription: must be the VPC Id of an existing Virtual Private Cloud
  Subnets:
    Type: 'List<AWS::EC2::Subnet::Id>'
    Description: The list of SubnetIds in your Virtual Private Cloud (VPC)
    ConstraintDescription: >-
      must be a list of at least two existing subnets associated with at least
      two different availability zones. They should be residing in the selected
      Virtual Private Cloud
  sgdask:
    Type: 'AWS::EC2::SecurityGroup::Id'
    Description: security group of dask fargate cluster
    ConstraintDescription:  private security group of dask fargate cluster
Resources:
  logsGroupScheduler:
    Type: 'AWS::Logs::LogGroup'
    Properties:
      LogGroupName: /ecs/dask-scheduler
    Metadata:
      'AWS::CloudFormation::Designer':
        id: 527588db-c546-4f16-a30b-faa587c9f2bc
  logsGroupWorkers:
    Type: 'AWS::Logs::LogGroup'
    Properties:
      LogGroupName: /ecs/dask-workers

  PrivateNameSpace:
    Type: 'AWS::ServiceDiscovery::PrivateDnsNamespace'
    Properties:
      Vpc: !Ref VpcId
      Name: local-dask

  ECSExecutionRoleV1:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ecs-tasks.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      RoleName: ECSExecutionRoleV1

  RolePolicies:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: ECSExecutionRoleV1
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action:
              - 'ecr:BatchCheckLayerAvailability'
              - 'ecr:GetDownloadUrlForLayer'
              - 'ecr:BatchGetImage'
              - 'ecr:GetAuthorizationToken'
              - 'logs:CreateLogStream'
              - 'logs:PutLogEvents'
            Effect: Allow
            Resource: '*'
      Roles:
        - Ref: ECSExecutionRoleV1
  
  DaskCluster:
    Type: 'AWS::ECS::Cluster'
    Properties:
      ClusterName: Fargate-Dask-Cluster
      ClusterSettings:
        - Name: containerInsights
          Value: enabled
      CapacityProviders:
        - FARGATE
        - FARGATE_SPOT
      DefaultCapacityProviderStrategy:
        - CapacityProvider: FARGATE
          Weight: 1
        - CapacityProvider: FARGATE_SPOT
          Weight: 1


  serviceRegistryScheduler:
    Type: 'AWS::ServiceDiscovery::Service'
    Properties:
      Name: Dask-Scheduler
      HealthCheckCustomConfig:
        FailureThreshold: 10
      DnsConfig:
        DnsRecords:
          - Type: A
            TTL: '60'
        NamespaceId: !Ref PrivateNameSpace
  serviceRegistryWorker:
    Type: 'AWS::ServiceDiscovery::Service'
    Properties:
      Name: Dask-Worker
      HealthCheckCustomConfig:
        FailureThreshold: 10
      DnsConfig:
        DnsRecords:
          - Type: A
            TTL: '60'
        NamespaceId: !Ref PrivateNameSpace
  taskdefinition:
    Type: 'AWS::ECS::TaskDefinition'
    DependsOn: serviceRegistryScheduler
    Properties:
      Family: !Join 
        - ''
        - - cloudformation
          - '-dask-scheduler-v1'
      Cpu: '2048'
      Memory: '4096'
      NetworkMode: awsvpc
      TaskRoleArn: !Ref ECSExecutionRoleV1
      ExecutionRoleArn: !Ref ECSExecutionRoleV1
      RequiresCompatibilities:
        - FARGATE
      ContainerDefinitions:
        - Name: dask-scheduler-v1
          Image: 'ramgvittal/dask:latest'
          Cpu: '2048'
          Memory: '4096'
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: /ecs/dask-scheduler
              awslogs-region: !Ref 'AWS::Region'
              awslogs-stream-prefix: ecs
          Essential: 'true'
          Command:
            - dask-scheduler

  taskdefinitionWorker:
    Type: 'AWS::ECS::TaskDefinition'
    DependsOn: DaskSchedulerService
    Properties:
      Family: !Join 
        - ''
        - - cloudformation
          - '-dask-workers-v1'
      Cpu: '2048'
      Memory: '4096'
      NetworkMode: awsvpc
      TaskRoleArn: !Ref ECSExecutionRoleV1
      ExecutionRoleArn: !Ref ECSExecutionRoleV1
      RequiresCompatibilities:
        - FARGATE
      ContainerDefinitions:
        - Name: dask-workers-v1
          Image: 'ramgvittal/dask:latest'
          Cpu: '2048'
          Memory: '4096'
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: /ecs/dask-workers
              awslogs-region: !Ref 'AWS::Region'
              awslogs-stream-prefix: ecs
          Essential: 'true'
          Command:
            - dask-worker
            - 'dask-scheduler.local-dask:8786'
            - '--memory-limit 4096MB'
            - '--worker-port 9000'
            - '--nanny-port 9001'
            - '--bokeh-port 9002'
  DaskSchedulerService:
    Type: 'AWS::ECS::Service'
    DependsOn:
      - serviceRegistryScheduler
      - logsGroupScheduler
    Properties:
      Cluster: !Ref DaskCluster
      DeploymentConfiguration:
        MaximumPercent: 200
        MinimumHealthyPercent: 100
      DesiredCount: 1
      LaunchType: FARGATE
      ServiceName: Dask-Scheduler
      ServiceRegistries:
        - RegistryArn: !GetAtt serviceRegistryScheduler.Arn
      TaskDefinition: !Ref taskdefinition
      NetworkConfiguration:
        AwsvpcConfiguration:
          AssignPublicIp: ENABLED
          SecurityGroups:
            - !Ref sgdask
          Subnets: !Ref Subnets

  DaskWorkerService:
    Type: 'AWS::ECS::Service'
    DependsOn:
      - serviceRegistryWorker
      - logsGroupWorkers
    Properties:
      Cluster: !Ref DaskCluster
      DeploymentConfiguration:
        MaximumPercent: 200
        MinimumHealthyPercent: 100
      DesiredCount: 1
      LaunchType: FARGATE
      ServiceName: Dask-Workers
      ServiceRegistries:
        - RegistryArn: !GetAtt serviceRegistryWorker.Arn
      TaskDefinition: !Ref taskdefinitionWorker
      NetworkConfiguration:
        AwsvpcConfiguration:
          AssignPublicIp: ENABLED
          SecurityGroups:
            - !Ref sgdask
          Subnets: !Ref Subnets