dask-environment.yaml

Parameters:
  SagemakerCodeRepo:
    Type: String
    Description: Github Repository for loading into the Sagemaker Jupyter environment
    Default: https://github.com/awslabs/amazon-asdi.git

  SagemakerNotebookInstance:
    Type: String
    Description: Sagemaker Notebook instance type
    AllowedValues:
      - ml.t3.medium
      - ml.t3.large
      - ml.t3.xlarge
      - ml.t3.2xlarge
      - ml.m5.large
      - ml.m5.2xlarge
      - ml.m5.4xlarge
    Default: ml.m5.2xlarge

  DaskImage:
    Type: String
    Default: daskdev/dask:latest
    Description: Container image to use for the Dask workers before our custom image is built

  DaskWorkerCPU:
    Type: String
    Default: '1024'
    Description: CPU units to assign to dask workers (1024 = 1 vCPU)

  DaskWorkerMemory:
    Type: String
    Default: '8192'
    Description: Memory in MiB for dask workers

  DaskWorkerSpotPricing:
    Type: String
    Default: 'NO'
    Description: Use Fargate spot pricing for dask workers, to save money (workers may be terminated)
    AllowedValues:
      - 'YES'
      - 'NO'

Conditions:
  DaskWorkerSpotPricingCondition: !Equals [!Ref DaskWorkerSpotPricing, 'YES']

Resources:
  DaskVpc:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: 10.10.0.0/16
      EnableDnsHostnames: true
      EnableDnsSupport: true
      InstanceTenancy: default
      Tags:
        - Key: Name
          Value: dask-fargate/MyVpc
  PublicSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      CidrBlock: 10.10.0.0/18
      VpcId:
        Ref: DaskVpc
      AvailabilityZone:
        Fn::Select:
          - 0
          - Fn::GetAZs: ""
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: dask-fargate/PublicSubnet1
  PublicSubnet1RouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId:
        Ref: DaskVpc
      Tags:
        - Key: Name
          Value: dask-fargate/MyVpc/PublicSubnet1
  PublicSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId:
        Ref: PublicSubnet1RouteTable
      SubnetId:
        Ref: PublicSubnet1
  PublicSubnet1DefaultRoute:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId:
        Ref: PublicSubnet1RouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId:
        Ref: InternetGateway
    DependsOn:
      - VPCGatewayA
  PublicSubnet1EIP:
    Type: AWS::EC2::EIP
    Properties:
      Domain: vpc
      Tags:
        - Key: Name
          Value: dask-fargate/MyVpc/PublicSubnet1
  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: dask-fargate/MyVpc
  VPCGatewayA:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId:
        Ref: DaskVpc
      InternetGatewayId:
        Ref: InternetGateway
  NATGateway:
    Type: AWS::EC2::NatGateway
    Properties:
      AllocationId: !GetAtt PublicSubnet1EIP.AllocationId
      SubnetId: !Ref PublicSubnet1
  PrivateSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      CidrBlock: 10.10.64.0/18
      VpcId:
        Ref: DaskVpc
      AvailabilityZone:
        Fn::Select:
          - 0
          - Fn::GetAZs: ""
      Tags:
        - Key: Name
          Value: dask-fargate/PrivateSubnet1
  PrivateSubnet1RouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId:
        Ref: DaskVpc
      Tags:
        - Key: Name
          Value: dask-fargate/MyVpc/PrivateSubnet1
  PrivateSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId:
        Ref: PrivateSubnet1RouteTable
      SubnetId:
        Ref: PrivateSubnet1
  PrivateSubnet1DefaultRoute:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId:
        Ref: PrivateSubnet1RouteTable
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId:
        Ref: NATGateway
  PrivateSubnetS3Endpoint:
    Type: AWS::EC2::VPCEndpoint
    Properties:
      RouteTableIds:
        - !Ref PrivateSubnet1RouteTable
      ServiceName: !Sub 'com.amazonaws.${AWS::Region}.s3'
      VpcId: !Ref DaskVpc

  LogGroupScheduler:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/${AWS::StackName}/dask/scheduler
      RetentionInDays: 3
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete
  LogGroupWorker:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/${AWS::StackName}/dask/worker
      RetentionInDays: 3
    UpdateReplacePolicy: Delete
    DeletionPolicy: Delete

  ECSExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: ecs-tasks.amazonaws.com
        Version: "2012-10-17"
  ECSExecutionRolePolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Statement:
          - Action:
              - ecr:BatchCheckLayerAvailability
              - ecr:GetDownloadUrlForLayer
              - ecr:BatchGetImage
              - ecr:GetAuthorizationToken
            Effect: Allow
            Resource: "*"
          - Action:
              - logs:CreateLogStream
              - logs:PutLogEvents
            Effect: Allow
            Resource:
              Fn::GetAtt:
                - LogGroupScheduler
                - Arn
          - Action:
              - logs:CreateLogStream
              - logs:PutLogEvents
            Effect: Allow
            Resource:
              Fn::GetAtt:
                - LogGroupWorker
                - Arn
        Version: "2012-10-17"
      PolicyName: ECSExecutionRolePolicy
      Roles:
        - Ref: ECSExecutionRole
  DaskCluster:
    Type: AWS::ECS::Cluster
    Properties:
      ClusterName: !Join ["-", [!Ref "AWS::StackName", "Fargate-Dask-Cluster"]]
      CapacityProviders:
        - FARGATE
        - FARGATE_SPOT
      DefaultCapacityProviderStrategy:
        - CapacityProvider: FARGATE
          Weight: 1
        - CapacityProvider: FARGATE_SPOT
          Weight: 0

  DaskClusterPrivateNS:
    Type: AWS::ServiceDiscovery::PrivateDnsNamespace
    Properties:
      Name: local-dask
      Vpc:
        Ref: DaskVpc

  SchedulerDefinition:
    Type: AWS::ECS::TaskDefinition
    Properties:
      ContainerDefinitions:
        - Command:
            - dask-scheduler
            - --dashboard
            - --dashboard-address
            - '8787'
          Cpu: 4096
          Essential: true
          Image: !Ref DaskImage
          LogConfiguration: 
            LogDriver: awslogs
            Options:
              awslogs-group:
                Ref: LogGroupScheduler
              awslogs-stream-prefix: ecs
              awslogs-region:
                Ref: AWS::Region
          Memory: 30720
          MemoryReservation: 30720
          Name: Dask
      Cpu: "4096"
      ExecutionRoleArn:
        Fn::GetAtt:
          - ECSExecutionRole
          - Arn
      Family: Dask-Scheduler
      Memory: "30720"
      NetworkMode: awsvpc
      RequiresCompatibilities:
        - FARGATE
      TaskRoleArn:
        Fn::GetAtt:
          - ECSExecutionRole
          - Arn
  WorkerDefinition:
    Type: AWS::ECS::TaskDefinition
    Properties:
      ContainerDefinitions:
        - Command:
            - dask-worker
            - dask-scheduler.local-dask:8786
            - --memory-limit
            - '7000MB'
            - --worker-port
            - '9000'
            - --no-nanny
            - --no-dashboard
            - --death-timeout
            - '30'
            - --nthreads
            - '2'
            - --nprocs
            - '1'
            - --reconnect
          Cpu: !Ref DaskWorkerCPU
          Essential: true
          Image: !Ref DaskImage
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group:
                Ref: LogGroupWorker
              awslogs-stream-prefix: ecs
              awslogs-region:
                Ref: AWS::Region
          Memory: !Ref DaskWorkerMemory
          MemoryReservation: !Ref DaskWorkerMemory
          Name: Dask
      Cpu: !Ref DaskWorkerCPU
      ExecutionRoleArn:
        Fn::GetAtt:
          - ECSExecutionRole
          - Arn
      Family: Dask-Worker
      Memory: !Ref DaskWorkerMemory
      NetworkMode: awsvpc
      RequiresCompatibilities:
        - FARGATE
      TaskRoleArn:
        Fn::GetAtt:
          - ECSExecutionRole
          - Arn

  DaskSchedulerSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Enable Scheduler ports access
      GroupName: DaskSchedulerSecurityGroup
      SecurityGroupEgress:
        - CidrIp: 0.0.0.0/0
          Description: Allow all outbound traffic by default
          IpProtocol: "-1"
      SecurityGroupIngress:
        - CidrIp: 0.0.0.0/0
          Description: Allow access to the dask dashboard
          FromPort: 8787
          IpProtocol: tcp
          ToPort: 8787
        - SourceSecurityGroupId: !GetAtt NotebookSecurityGroup.GroupId
          Description: Allow access from SageMaker notebook
          FromPort: 8786
          IpProtocol: tcp
          ToPort: 8789
      VpcId: !Ref DaskVpc

  NotebookSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allow access to SageMaker notebook instance
      GroupName: NotebookSecurityGroup
      SecurityGroupEgress:
        - CidrIp: 0.0.0.0/0
          Description: Allow all outbound traffic by default
          IpProtocol: '-1'
      VpcId: !Ref DaskVpc

  DaskWorkerSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Restrict connectivity to dask workers
      GroupName: DaskWorkerSecurityGroup
      SecurityGroupEgress:
        - CidrIp: 0.0.0.0/0
          Description: Allow all outbound traffic by default
          IpProtocol: '-1'
      SecurityGroupIngress:
        - SourceSecurityGroupId: !GetAtt DaskSchedulerSecurityGroup.GroupId
          Description: Allow full access from dask scheduler
          IpProtocol: '-1'
        - SourceSecurityGroupId: !GetAtt NotebookSecurityGroup.GroupId
          Description: Allow full access from the notebook
          IpProtocol: '-1'
      VpcId: !Ref DaskVpc

  DaskWorkerToWorkerIngress:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      Description: Allow dask workers to communicate freely with each other
      GroupId: !GetAtt DaskWorkerSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt DaskWorkerSecurityGroup.GroupId
      IpProtocol: '-1'

  DaskWorkerToSchedulerIngress:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      Description: Allow dask workers to communicate with the scheduler
      GroupId: !GetAtt DaskSchedulerSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt DaskWorkerSecurityGroup.GroupId
      IpProtocol: '-1'

  DaskSchedulerService:
    Type: AWS::ECS::Service
    Properties:
      Cluster:
        Ref: DaskCluster
      DeploymentConfiguration:
        MaximumPercent: 200
        MinimumHealthyPercent: 100
      DesiredCount: 1
      EnableECSManagedTags: false
      LaunchType: FARGATE
      NetworkConfiguration:
        AwsvpcConfiguration:
          AssignPublicIp: ENABLED
          SecurityGroups:
            - Fn::GetAtt:
                - DaskSchedulerSecurityGroup
                - GroupId
          Subnets:
            - Ref: PublicSubnet1
      ServiceName: Dask-Scheduler
      ServiceRegistries:
        - RegistryArn:
            Fn::GetAtt:
              - DaskSchedulerServiceDiscovery
              - Arn
      TaskDefinition:
        Ref: SchedulerDefinition

  DaskSchedulerServiceDiscovery:
    Type: AWS::ServiceDiscovery::Service
    Properties:
      DnsConfig:
        DnsRecords:
          - TTL: 5
            Type: A
        NamespaceId:
          Fn::GetAtt:
            - DaskClusterPrivateNS
            - Id
        RoutingPolicy: MULTIVALUE
      HealthCheckCustomConfig:
        FailureThreshold: 1
      Name: Dask-Scheduler
      NamespaceId:
        Fn::GetAtt:
          - DaskClusterPrivateNS
          - Id

  DaskWorkerService:
    Type: AWS::ECS::Service
    Properties:
      CapacityProviderStrategy:
        - CapacityProvider:
            Fn::If:
              - DaskWorkerSpotPricingCondition
              - FARGATE_SPOT
              - FARGATE
          Weight: 1
      Cluster:
        Ref: DaskCluster
      DeploymentConfiguration:
        MaximumPercent: 200
        MinimumHealthyPercent: 100
      DesiredCount: 0
      EnableECSManagedTags: false
      NetworkConfiguration:
        AwsvpcConfiguration:
          SecurityGroups:
            - Fn::GetAtt:
                - DaskWorkerSecurityGroup
                - GroupId
          Subnets:
            - Ref: PrivateSubnet1
      ServiceName: Dask-Worker
      TaskDefinition:
        Ref: WorkerDefinition

  SagemakerRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action: sts:AssumeRole
            Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
        Version: "2012-10-17"
  SagemakerPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument:
        Statement:
          - Action:
              - s3:*
              - ecs:*
              - logs:*
              - ec2:DescribeNetworkInterfaces
            Effect: Allow
            Resource: "*"
        Version: "2012-10-17"
      PolicyName: notebookAccessPolicy
      Roles:
        - Ref: SagemakerRole

  NotebookLifecycleConfig:
    Type: AWS::SageMaker::NotebookInstanceLifecycleConfig
    Properties:
        OnStart:
        - Content:
            Fn::Base64: |
                #!/bin/sh

                set -e
                cat > /home/ec2-user/setup.sh << EOF
                #!/bin/bash
                sleep 10
                echo "\$(date) Creating daskpy3 conda environment"
                echo "\$(date) Running in directory ${PWD}"
                conda create --name daskpy3 python="3.7.10" -y
                source activate daskpy3
                conda install -c conda micromamba -y 
                micromamba install -c conda basemap proj4 nodejs ipywidgets -y 
                pip install botocore==1.20.106 aiobotocore==1.4.1 zarr==2.9.5 rechunker==0.4.2 ipykernel==6.3.1 boto3==1.17.106 dask==2021.8.1 distributed==2021.8.1 tornado==6.1 cloudpickle==1.6.0 msgpack==1.0.2 blosc==1.10.2 numpy==1.21.2 pandas==1.3.2 lz4==3.1.3 netcdf4==1.5.7 xarray==0.19.0 bokeh==2.2.3 s3fs==2021.8.1 fsspec==2021.8.1 h5netcdf==0.11.0 h5py==3.4.0 intake-esm==2021.8.17 intake==0.6.3
                python -m ipykernel install --user --name daskpy3 --display-name "conda_daskpy3"
                echo "\$(date) Updating Jupyter config options"
                cp ~/.jupyter/jupyter_notebook_config.py ~/.jupyter/jupyter_notebook_config.py.bak
                sed -i 's/^#c.NotebookApp.iopub_data_rate_limit.*$/c.NotebookApp.iopub_data_rate_limit = 1.0e10/' ~/.jupyter/jupyter_notebook_config.py
                sed -i 's/^#c.NotebookApp.iopub_msg_rate_limit.*$/c.NotebookApp.iopub_msg_rate_limit = 1.0e10/' ~/.jupyter/jupyter_notebook_config.py
                source /home/ec2-user/anaconda3/bin/deactivate
                source activate JupyterSystemEnv
                jupyter labextension install @jupyter-widgets/jupyterlab-manager@1.1
                source /home/ec2-user/anaconda3/bin/deactivate
                echo "\$(date) Finished!"
                EOF
            
                chown ec2-user:ec2-user /home/ec2-user/setup.sh
                chmod +x /home/ec2-user/setup.sh

                sudo -u ec2-user -i nohup /home/ec2-user/setup.sh >/home/ec2-user/output.log 2>&1 &

  DaskNotebook:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      InstanceType: !Ref SagemakerNotebookInstance
      RoleArn:
        Fn::GetAtt:
          - SagemakerRole
          - Arn
      DefaultCodeRepository: !Ref SagemakerCodeRepo
      DirectInternetAccess: Enabled
      LifecycleConfigName:  !GetAtt NotebookLifecycleConfig.NotebookInstanceLifecycleConfigName
      RootAccess: Enabled
      SecurityGroupIds:
        - Fn::GetAtt:
            - NotebookSecurityGroup
            - GroupId
      SubnetId:
        Ref: PublicSubnet1
      VolumeSizeInGB: 50

Outputs:
  JupyterNotebook:
      Value: !Join
      - ''
      - - https://
        - !Ref 'AWS::Region'
        - .console.aws.amazon.com/sagemaker/home?region=
        - !Ref 'AWS::Region'
        - '#/notebook-instances/openNotebook/'
        - !GetAtt 'DaskNotebook.NotebookInstanceName'
        - '?view=lab'
  DaskECSClusterName:
    Value: !Ref DaskCluster