hail-jupyter.yml

AWSTemplateFormatVersion: "2010-09-09"
Description: "SageMaker Jupyter Instance for use with Hail on EMR"


Parameters:

  pAllowSsmShell:
    AllowedValues:
      - "false"
      - "true"
    Default: "false"
    Description: "Required - Allow Systems Manager (SSM) shell access to EMR nodes.  Must also be set to true on EMR CloudFormation stacks."
    Type: "String"

  pIngressCidr:
    AllowedPattern: "(\\d{1,3})\\.(\\d{1,3})\\.(\\d{1,3})\\.(\\d{1,3})/(\\d{1,2})"
    Description: "Required - CIDR is allowed full security group ingress.  Examples - VPC CIDR, VPN CIDR, etc."
    Type: "String"

  pInstanceType:
    AllowedValues:
      - "ml.t2.medium"
      - "ml.t2.large"
      - "ml.t2.xlarge"
      - "ml.t2.2xlarge"
      - "ml.t3.medium"
      - "ml.t3.large"
      - "ml.t3.xlarge"
      - "ml.t3.2large"
    Default: "ml.t3.medium"
    Description: "Required"
    Type: "String"

  pName:
    Description: "Required - Used as the name of the notebook instance and S3 backup location.  Recommend username.  Example: aperry"
    Type: "String"

  pS3bucketJupyter:
    Description: "Required - This parameter is set by the S3 stack.  No changes should be necessary.  S3 bucket with Jupyter scripts, common notebooks, and home directory backups."
    Type: "AWS::SSM::Parameter::Value<String>"
    Default: "/hail/s3/sagemaker"

  pS3bucketHail:
    Description: "Required - This parameter is set by the S3 stack.  No changes should be necessary.  Notebook instance will have full read."
    Type: "AWS::SSM::Parameter::Value<String>"
    Default: "/hail/s3/hail"

  pSubnetId:
    Description: "Required - Target subnet for Notebook server.  Private subnet is strongly suggested as browser access will still be available via the AWS console."
    Type: "AWS::EC2::Subnet::Id"

  pTagEnvironment:
    AllowedValues:
      - "production"
      - "staging"
      - "sandbox"
      - "test"
      - "development"
      - "qa"
      - "dr"
    Default: "development"
    Description: "Optional - Environment type for default resource tagging."
    Type: "String"

  pEfsFsapId:
    Default: ""
    Description: "Optional - EFS file system acccess point.  If provided, IAM access will be granted to mount the FSAP.  E.g. fsap-0b3ec479999999999"
    Type: "String"

  pTagOwner:
    Default: ""
    Description: "Optional - Owner of the resources.  Person/Department, etc."
    Type: "String"

  pVolumeSize:
    Default: "20"
    Description: "Required - Size (GB)."
    MaxValue: "500"
    MinValue: "5"
    Type: "Number"

  pVpcId:
    Description: "Required - Notebook is deployed in this VPC."
    Type: "AWS::EC2::VPC::Id"


Metadata:

  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: "Network Details"
        Parameters:
          - "pVpcId"
          - "pSubnetId"
          - "pIngressCidr"
      - Label:
          default: "Instance Details"
        Parameters:
          - "pName"
          - "pInstanceType"
          - "pVolumeSize"
          - "pAllowSsmShell"
      - Label:
          default: "S3 Resources"
        Parameters:
          - "pS3bucketHail"
          - "pS3bucketJupyter"
      - Label:
          default: "EFS Settings"
        Parameters:
          - "pEfsFsapId"
      - Label:
          default: "Tagging"
        Parameters:
          - "pTagEnvironment"
          - "pTagOwner"
    ParameterLabels:
      pAllowSsmShell:
        default: "Allow SSM Shell Access to EMR Nodes"
      pVpcId:
        default: "VPC ID"
      pSubnetId:
        default: "Subnet ID"
      pIngressCidr:
        default: "Ingress CIDR"
      pEfsFsapId:
        default: "File System Access Point ID"
      pName:
        default: "Instance Name"
      pInstanceType:
        default: "Instance Type"
      pVolumeSize:
        default: "Volume Size"
      pS3bucketHail:
        default: "SSM Parameter for Hail Bucket"
      pS3bucketJupyter:
        default: "SSM Parameter for SageMaker Bucket"
      pTagOwner:
        default: "Owner"
      pTagEnvironment:
        default: "Environment"


Conditions:

  CondAllowSsmShell: !Equals [!Ref pAllowSsmShell, "true"]
  CondEfsFsapId: !Not [ !Equals [ !Ref pEfsFsapId, "" ] ]

Resources:

  instance:
    Type: "AWS::SageMaker::NotebookInstance"
    Properties:
      DirectInternetAccess: "Disabled"
      InstanceType: !Ref pInstanceType
      RoleArn: !GetAtt role.Arn
      LifecycleConfigName: !GetAtt lifecycleConfiguration.NotebookInstanceLifecycleConfigName
      NotebookInstanceName: !Ref pName
      RootAccess: Enabled
      SubnetId: !Ref pSubnetId
      VolumeSizeInGB: !Ref pVolumeSize
      SecurityGroupIds:
        - !Ref sg
      Tags:
        - Key: "Name"
          Value: !Ref pName
        - Key: "environment"
          Value: !Ref pTagEnvironment
        - Key: "owner"
          Value: !Ref pTagOwner

  lifecycleConfiguration:
    Type: "AWS::SageMaker::NotebookInstanceLifecycleConfig"
    Properties:
      NotebookInstanceLifecycleConfigName: !Sub "${pName}"
      OnCreate:
        - Content:
            Fn::Base64: !Sub |
              #!/bin/bash
              set -ex

              BUCKET="${pS3bucketJupyter}"
              EC2_USER_BIN="/home/ec2-user/SageMaker/bin/"
              sudo -u ec2-user -i <<EOF
              set -ex
              /usr/bin/aws s3 cp s3://"$BUCKET"/scripts/ "$EC2_USER_BIN" --recursive
              find "$EC2_USER_BIN" -type f -exec chmod -v 700 {} \;
              EOF
      OnStart:
        - Content:
            Fn::Base64: !Sub |
              #!/bin/bash
              set -ex

              BUCKET="${pS3bucketJupyter}"
              DIRECTORY="/home/ec2-user/SageMaker/"

              # Copy previously backed up user notebooks
              /usr/bin/aws s3 cp s3://$BUCKET/${pName}/ $DIRECTORY --recursive
              # Reset bin permissions
              find "$DIRECTORY/bin/" -type f -exec chmod -v 700 {} \;
              # Copy common organization notebooks.  This will overwrite common notebooks previously saved to the home directory
              /usr/bin/aws s3 cp s3://$BUCKET/common-notebooks/ $DIRECTORY/common-notebooks/ --recursive

              # Use the system AWS, since ec2-user may change conda environments
              CMD="/usr/bin/aws s3 sync --exclude 'bin/*' --exclude '.sparkmagic/*' --exclude 'efs/*' --exclude '*.ipynb_checkpoints/*' $DIRECTORY s3://$BUCKET/${pName}/ > /home/ec2-user/s3-backup.log 2>&1"
              chown "ec2-user":"ec2-user" $DIRECTORY --recursive
              # Backup the notebook details to S3 on 5 minute intervals
              cronjob="*/5 * * * * $CMD"
              echo "$cronjob" | crontab -u "ec2-user" -

              # Conda/Pip additions
              sudo -u ec2-user -i <<EOF
              source /home/ec2-user/anaconda3/bin/activate JupyterSystemEnv
              # Required for "%manage_spark"
              jupyter labextension install @jupyter-widgets/jupyterlab-manager
              source /home/ec2-user/anaconda3/bin/deactivate

              # Prevent iopub throttling
              sed -i.bak 's/^#c.NotebookApp.iopub_data_rate_limit.*$/c.NotebookApp.iopub_data_rate_limit = 1000000/' ~/.jupyter/jupyter_notebook_config.py
              # Session Manager for SSM connections to EMR nodes, if desired.
              sudo yum -y install https://s3.amazonaws.com/session-manager-downloads/plugin/latest/linux_64bit/session-manager-plugin.rpm
              EOF

              # Execute user's specific startup operations, if present
              CUSTOM_USER_STARTUP_SCRIPT="$DIRECTORY/custom-user-startup.sh"
              if [ -f "$CUSTOM_USER_STARTUP_SCRIPT" ]; then source "$CUSTOM_USER_STARTUP_SCRIPT"; fi



  role:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Principal:
              Service:
                - "sagemaker.amazonaws.com"
            Action:
              - "sts:AssumeRole"
      Path: "/"
      Policies:
        - PolicyName: "s3FullUserPathControl"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "s3:*"
                Resource:
                  - !Sub "arn:aws:s3:::${pS3bucketJupyter}/${pName}/*"
              - Effect: "Allow"
                Action:
                  - "s3:ListObject"
                  - "s3:ListBucket"
                Resource: !Sub "arn:aws:s3:::${pS3bucketJupyter}"
        - PolicyName: "s3ListPaths"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "s3:ListBucket"
                Resource:
                  - !Sub "arn:aws:s3:::${pS3bucketJupyter}"
                Condition:
                  StringLike:
                    s3:prefix:
                      - !Sub "${pName}/"
                      - "scripts/"
                      - "common-notebooks/"
              - Effect: "Allow"
                Action:
                  - "s3:GetObject"
                Resource:
                  - !Sub "arn:aws:s3:::${pS3bucketJupyter}/scripts/*"
                  - !Sub "arn:aws:s3:::${pS3bucketJupyter}/common-notebooks/*"
        - PolicyName: "s3hailFullRead"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "s3:Get*"
                  - "s3:List*"
                Resource:
                  - !Sub "arn:aws:s3:::${pS3bucketHail}/*"
                  - !Sub "arn:aws:s3:::${pS3bucketHail}"
        - PolicyName: "cloudwatchWrite"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogStream"
                  - "logs:CreateLogGroup"
                  - "logs:PutLogEvents"
                Resource: "*"
        - PolicyName: "emrListClusters"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "elasticmapreduce:ListInstances"
                  - "elasticmapreduce:ListBootstrapActions"
                  - "elasticmapreduce:ListClusters"
                  - "elasticmapreduce:DescribeCluster"
                  - "elasticmapreduce:ListInstanceGroups"
                Resource: "*"
      RoleName: !Sub "sagemaker-jupyter-${pName}"

  policyAllowSsmShell:
    Type: "AWS::IAM::Policy"
    Condition: "CondAllowSsmShell"
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Action:
              - "ssm:StartSession"
            Resource: "arn:aws:ec2:*:*:instance/*"
            Condition:
              StringLike:
                "ssm:resourceTag/allow-emr-ssm":
                  - "true"
      PolicyName: "ssm-session-by-tag"
      Roles:
        - !Ref role

  policyEfsFsapFullAccess:
    Condition: "CondEfsFsapId"
    Type: "AWS::IAM::Policy"
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Sid: "efsFsapFull"
            Action:
              - "elasticfilesystem:Client*"
            Effect: "Allow"
            Resource: "*"
            Condition:
              StringEquals:
                "elasticfilesystem:AccessPointArn": !Sub "arn:aws:elasticfilesystem:${AWS::Region}:${AWS::AccountId}:access-point/${pEfsFsapId}"
      PolicyName: "efs-fsap-full-access"
      Roles:
        - !Ref role


  sg:
    Type: "AWS::EC2::SecurityGroup"
    Properties:
      VpcId: !Ref pVpcId
      GroupDescription: "Security group for Sagemaker Notebook to connect to Hail cluster."
      SecurityGroupIngress:
        - IpProtocol: "-1"
          FromPort: -1
          ToPort: -1
          CidrIp: !Ref pIngressCidr
          Description: "Full Access"
      SecurityGroupEgress:
        - IpProtocol: "-1"
          FromPort: -1
          ToPort: -1
          CidrIp: "0.0.0.0/0"
          Description: "all"
      Tags:
        - Key: "Name"
          Value: !Sub "sagemaker-jupyter-${pName}"
        - Key: "environment"
          Value: !Ref pTagEnvironment
        - Key: "owner"
          Value: !Ref pTagOwner