AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  sagemaker-auto-shutdown

  Sample SAM Template for sagemaker-auto-shutdown

Parameters:
  EndpointExcludeTagKey:
    Type: String
    Default: env
    AllowedPattern: '^([\p{L}\p{Z}\p{N}_.:/=+\-@]*)$'
    Description: 'Endpoint with this tag will not be deleted.'
  EndpointExcludeTagValue:
    Type: String
    Default: prod
    AllowedPattern: '^([\p{L}\p{Z}\p{N}_.:/=+\-@]*)$'
    Description: 'Endpoint with this tag will not be deleted.'
  NotebookExcludeTagKey:
    Type: String
    Default: env
    AllowedPattern: '^([\p{L}\p{Z}\p{N}_.:/=+\-@]*)$'
    Description: 'Notebook with this tag will not be stopped.'
  NotebookExcludeTagValue:
    Type: String
    Default: prod
    AllowedPattern: '^([\p{L}\p{Z}\p{N}_.:/=+\-@]*)$'
    Description: 'Notebook with this tag will not be stopped.'

# More info about Globals: https://github.com/awslabs/serverless-application-model/blob/master/docs/globals.rst
Globals:
  Function:
    Timeout: 300 
    MemorySize: 128
    Environment:
      Variables:
        ENDPOINT_EXCLUDE_TAG: !Sub '{ "Key": "${EndpointExcludeTagKey}", "Value": "${EndpointExcludeTagValue}" }'
        NOTEBOOK_EXCLUDE_TAG: !Sub '{ "Key": "${NotebookExcludeTagKey}", "Value": "${NotebookExcludeTagValue}" }'
        ENDPOINT_SPEC: '[{ "EndpointConfigName": "Example-Endpoint-Config", "EndpointName": "Example-Endpoint"}]'
        MAX_COUNT: 100 # Max number of endpoint or notebook to be cleaned in one run. This is to reduce the blast radius in case of misconfiguration.
        LOG_LEVEL: INFO

Resources:
  CleanSageMakerResourcesFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction
    Properties:
      CodeUri: cleaner/
      Handler: app.lambda_handler
      Runtime: python3.9
      Architectures:
        - x86_64
      Policies:
        - Statement:
            - Effect: Allow
              Action:
                - sagemaker:ListEndpoints
                - sagemaker:DescribeEndpoint
                - sagemaker:DescribeEndpointConfig
                - sagemaker:ListNotebookInstances
                - sagemaker:ListTags
              Resource: '*'
         - Statement:
             - Effect: Allow
               Action:
                 - sagemaker:DeleteEndpoint
                 - sagemaker:StopNotebookInstance
               Resource: '*'
                  
      Events:
        CloudWatchEvent:
          Type: Schedule
          Properties:
            Schedule: cron(0 20 * * ? *) # Every day at 20:00 UTC
            
  CreateSageMakerEndpointFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction
    Properties:
      CodeUri: creater/
      Handler: app.lambda_handler
      Runtime: python3.9
      Architectures:
        - x86_64
      Policies:
        - Statement:
            - Effect: Allow
              Action:
                - sagemaker:CreateEndpoint
              Resource: '*'
      Events:
        CloudWatchEvent:
          Type: Schedule
          Properties:
            Schedule: cron(0 7 * * ? *) # Every day at 7:00 UTC

Outputs:
  # ServerlessRestApi is an implicit API created out of Events key under Serverless::Function
  # Find out more about other implicit resources you can reference within SAM
  # https://github.com/awslabs/serverless-application-model/blob/master/docs/internals/generated_resources.rst#api
  CleanSageMakerResourcesFunction:
    Description: "Clean SageMaker Resources Lambda Function ARN"
    Value: !GetAtt CleanSageMakerResourcesFunction.Arn
  CleanSageMakerResourcesFunctionIamRole:
    Description: "Implicit IAM Role created for lambda function"
    Value: !GetAtt CleanSageMakerResourcesFunctionRole.Arn