Description: SageMaker notebook instance, S3 Bucket and IAM roles to be used for Lookout for Equipment Getting Started tutorial

Resources:
  S3Bucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName: !Join
        - "-"
        - - "lookoutequiment-getting-started"
          - !Select
            - 0
            - !Split
              - "-"
              - !Select
                - 2
                - !Split
                  - "/"
                  - !Ref "AWS::StackId"

  SageMakerIamRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
                - lookoutequipment.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"
        - "arn:aws:iam::aws:policy/AmazonS3FullAccess"
        - "arn:aws:iam::aws:policy/IAMFullAccess"
        - "arn:aws:iam::aws:policy/AmazonLookoutEquipmentFullAccess"
      Path: /
      Policies:
        - PolicyName: LookoutForEquipment
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: 'lookoutequipment:*'
                Resource: '*'
                Sid: "LookoutEquipmentAPIAccess"
              - Effect: Allow
                Action: 'iam:PassRole'
                Resource: '*'
              - Effect: Allow
                Action:
                  - 's3:Get*'
                  - 's3:List*'
                  - 's3:PutObject'
                Resource: !Join
                  - ""
                  - - "arn:aws:s3:::"
                    - !Ref S3Bucket
                    - "*"
                Sid: TrainingAndLabelDataAccess

  NotebookInstance:
    Type: 'AWS::SageMaker::NotebookInstance'
    Properties:
      InstanceType: ml.t3.medium
      RoleArn: !GetAtt SageMakerIamRole.Arn
      PlatformIdentifier: notebook-al2-v1
      NotebookInstanceName: !Join
        - "-"
        - - "lookout-equipment-getting-started"
          - !Select
            - 0
            - !Split
              - "-"
              - !Select
                - 2
                - !Split
                  - "/"
                  - !Ref "AWS::StackId"

      VolumeSizeInGB: 20
      DefaultCodeRepository: 'https://github.com/aws-samples/amazon-lookout-for-equipment.git'
      LifecycleConfigName: !GetAtt NotebookLifecycleConfig.NotebookInstanceLifecycleConfigName

  NotebookLifecycleConfig:
    Type: 'AWS::SageMaker::NotebookInstanceLifecycleConfig'
    Properties:
      NotebookInstanceLifecycleConfigName: 'install-label-studio'
      OnStart:
        - Content:
            Fn::Base64: !Sub |
                  #!/bin/bash
                  set -e
                  #################################################
                  ## INSTALL DEPENDENCIES AND PULL CONTAINER IMAGE
                  #################################################
                  
                  echo ==INSTALLING DEPENDENCIES==
                  /home/ec2-user/anaconda3/envs/JupyterSystemEnv/bin/pip install -U jupyter-server-proxy
                  
                  /usr/bin/docker pull heartexlabs/label-studio:latest
                  #################################################
                  ### INTEGRATE LABEL-STUDIO WITH JUPYTER
                  #################################################
                  echo ==UPDATING JUPYTER SERVER CONFIG==
                  # Strip out the name of the notebook from the instance bootstrap and the region details from aws/config
                  INSTANCE_NAME=$(grep  notebook_uri   /etc/opt/ml/sagemaker-notebook-instance-config.json| sed 's/.*(\(.*\))/\1/' | awk -F'|' '{print $1}')
                  AWS_REGION=$(cat /home/ec2-user/.aws/config | grep region | awk '{print $NF}')
                  
                  
                  cat >>/home/ec2-user/.jupyter/jupyter_notebook_config.py <<EOC
                  c.ServerProxy.servers = {
                    'label-studio': {
                        'launcher_entry': {
                              'enabled': True,
                              'title': 'Label Studio',
                        },
                        'command': ['/usr/bin/docker', 'run', '-t', '-p', '{port}:8080', '-v', '/home/ec2-user/SageMaker/mydata:/label-studio/data', 'heartexlabs/label-studio:latest', 'label-studio', '--host', 'https://$INSTANCE_NAME.notebook.$AWS_REGION.sagemaker.aws/proxy/{port}'],
                        'absolute_url': False,
                        'timeout': 60
                    }
                  }
                  EOC
                  echo ==INSTALL SUCCESSFUL. RESTARTING JUPYTER==
                  # RESTART THE JUPYTER SERVER
                  systemctl restart jupyter-server

Outputs:
  NotebookInstanceId:
    Value: !Ref NotebookInstance
  S3Bucket:
    Value: !Ref S3Bucket
    Description: S3 Bucket for object storage