model_training_template.yaml

AWSTemplateFormatVersion: '2010-09-09'
Transform:
  - AWS::LanguageExtensions
  - AWS::Serverless-2016-10-31

Parameters:
  Environment:
    Type: String
    AllowedValues:
      - staging
      - prod
  GitShaHash:
    Type: String
    Default: latest
  ExportTimeout:
    Type: Number
    Default: 900
  ModelDataVersion:
    Type: String
    Default: latest
  ModelTrainingVersion:
    Type: String
    Default: latest

Globals:
  Function:
    Architectures:
      - x86_64
    Timeout: 30
    MemorySize: 128
    AutoPublishAlias: live
    DeploymentPreference:
      Type: AllAtOnce

Conditions:
  IsProd: !Equals [!Ref Environment, "prod"]
  IsStaging: !Equals [!Ref Environment, "staging"]

# TODO: review IAM policies to reduce permissions scope to only what is required

Resources:
  RatingsExportFunction:
    Type: AWS::Serverless::Function
    Properties:
      PackageType: Image
      Role: !GetAtt RatingsExportRole.Arn
      Environment:
        Variables:
          PYTHON_ENV: !Ref Environment
          DESTINATION_BUCKET_NAME: !Ref RatingsExportBucket
          RATING_TABLE_NAME: !Ref RatingTable
          RATING_FILE_NAME: !Sub ratings-${ModelDataVersion}.csv
      Timeout: !Ref ExportTimeout
      LoggingConfig:
        ApplicationLogLevel: INFO
        LogFormat: JSON
    Metadata:
      Dockerfile: ./Dockerfile
      DockerContext: ./ratings_exporter
      DockerTag: !If [ IsProd, !Ref GitShaHash, latest ]

  RatingsExportFunctionFunctionLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${RatingsExportFunction}
      RetentionInDays: !If [ IsProd, 30, 7 ]

  RatingsExportRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - 'lambda.amazonaws.com'
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Policies:
        - PolicyName: !Sub ${AWS::StackName}-RatingsExportBucket
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "s3:GetObject"
                  - "s3:ListBucket"
                  - "s3:PutObject"
                Resource:
                  - !GetAtt RatingsExportBucket.Arn
                  - !Sub
                    - ${BucketArn}/*
                    - BucketArn: !GetAtt RatingsExportBucket.Arn

        - PolicyName: !Sub ${AWS::StackName}-RatingsExportData
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - "dynamodb:BatchGet*"
                  - "dynamodb:DescribeStream"
                  - "dynamodb:DescribeTable"
                  - "dynamodb:Get*"
                  - "dynamodb:Query"
                  - "dynamodb:Scan"
                Resource:
                  - !GetAtt RatingTable.Arn
                  - !Sub
                    - ${TableName}/*
                    - TableName: !GetAtt RatingTable.Arn

  ModelTrainerFunction:
    Type: AWS::Serverless::Function
    Properties:
      PackageType: Image
      Role: !GetAtt ModelTrainerRole.Arn
      Environment:
        Variables:
          PYTHON_ENV: !Ref Environment
          RATING_BUCKET_NAME: !Ref RatingsExportBucket
          RATING_FILE_NAME: !Sub ratings-${ModelDataVersion}.csv
          MODEL_BUCKET_NAME: !Ref ModelBucket
          STACK_NAME: !Sub ${AWS::StackName}
          MODEL_TRAINING_VERSION: !Ref ModelTrainingVersion
          MODEL_DATA_VERSION: !Ref ModelDataVersion
          MODEL_TRAINING_ECR_URL: !Sub
                                    - ${RepositoryUri}:${ModelTrainingVersion}
                                    - RepositoryUri: !GetAtt ModelTrainingRepository.RepositoryUri
                                      ModelTrainingVersion: !Ref ModelTrainingVersion
          SAGEMAKER_ROLE: !GetAtt SageMakerExecutionRole.Arn
      LoggingConfig:
        ApplicationLogLevel: INFO
        LogFormat: JSON
    Metadata:
      Dockerfile: ./Dockerfile
      DockerContext: ./model_trainer
      DockerTag: !If [ IsProd, !Ref GitShaHash, latest ]

  ModelTrainerFunctionLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub /aws/lambda/${ModelTrainerFunction}
      RetentionInDays: !If [ IsProd, 30, 7 ]

  ModelTrainerRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - 'lambda.amazonaws.com'
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Policies:
        - PolicyName: !Sub ${AWS::StackName}-ModelTrainerRatingsExportBucket
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "s3:GetObject"
                  - "s3:ListBucket"
                Resource:
                  - !GetAtt RatingsExportBucket.Arn
                  - !Sub
                    - ${BucketArn}/*
                    - BucketArn: !GetAtt RatingsExportBucket.Arn
        - PolicyName: !Sub ${AWS::StackName}-ModelTrainerModelBucket
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "s3:GetObject"
                  - "s3:ListBucket"
                  - "s3:PutObject"
                Resource:
                  - !GetAtt ModelBucket.Arn
                  - !Sub
                    - ${BucketArn}/*
                    - BucketArn: !GetAtt ModelBucket.Arn
        - PolicyName: !Sub ${AWS::StackName}-ModelTrainerSagemaker
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "sagemaker:*"
                Resource:
                  - '*'
        - PolicyName: !Sub ${AWS::StackName}-ModelTrainerIam
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "iam:PassRole"
                Resource:
                  - !GetAtt SageMakerExecutionRole.Arn
                  - '*'

  SageMakerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Principal:
              Service: "sagemaker.amazonaws.com"
            Action: "sts:AssumeRole"
      Policies:
        - PolicyName: "S3AccessPolicy"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "s3:GetObject"
                  - "s3:ListBucket"
                  - "s3:PutObject"
                Resource:
                  - !GetAtt ModelBucket.Arn
                  - !Sub
                    - ${BucketArn}/*
                    - BucketArn: !GetAtt ModelBucket.Arn
                  - !GetAtt RatingsExportBucket.Arn
                  - !Sub
                    - ${BucketArn}/*
                    - BucketArn: !GetAtt RatingsExportBucket.Arn
        - PolicyName: "ECRAccessPolicy"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "ecr:BatchCheckLayerAvailability"
                  - "ecr:GetDownloadUrlForLayer"
                  - "ecr:BatchGetImage"
                  - "ecr:ListImages"
                Resource:
                  - '*'
        - PolicyName: "CloudwatchAccess"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "cloudwatch:PutMetricData"
                Resource: "*"
                Condition:
                  StringLike:
                    cloudwatch:namespace:
                      - "*SageMaker*"
                      - "*Sagemaker*"
                      - "*sagemaker*"
              - Effect: "Allow"
                Action:
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                  - "logs:CreateLogGroup"
                  - "logs:DescribeLogStreams"
                Resource: "arn:aws:logs:*:*:log-group:/aws/sagemaker/*"
        - PolicyName: "ECRFullAccessPolicy"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "ecr:*"
                Resource: "*"

  ModelTrainingStateMachine:
    Type: AWS::StepFunctions::StateMachine
    Properties:
      StateMachineName: !Sub ${AWS::StackName}-ModelTrainingStateMachine
      RoleArn: !GetAtt ModelTrainingStateMachineRole.Arn
      DefinitionString: !Sub
        - |
            {
              "Comment": "Orchestrates data export and model training kickoff",
              "StartAt": "InvokeRatingsExport",
              "States": {
                "InvokeRatingsExport": {
                  "Type": "Task",
                  "Resource": "${RatingsExportArn}",
                  "Next": "InvokeModelTrainer"
                },
                "InvokeModelTrainer": {
                  "Type": "Task",
                  "Resource": "${ModelTrainerArn}",
                  "End": true
                }
              }
            }
        - RatingsExportArn: !GetAtt RatingsExportFunction.Arn
          ModelTrainerArn: !GetAtt ModelTrainerFunction.Arn

  ModelTrainingStateMachineRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: states.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: MyStateMachinePolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource:
                  - !GetAtt RatingsExportFunction.Arn
                  - !GetAtt ModelTrainerFunction.Arn

  ShowTable:
    Type: AWS::DynamoDB::Table
    DeletionPolicy: Retain
    Properties:
      BillingMode: PAY_PER_REQUEST
      TableClass: STANDARD
      TableName: !Sub ${AWS::StackName}-ShowsV1
      AttributeDefinitions:
        - AttributeName: ShowId
          AttributeType: S
        - AttributeName: Title
          AttributeType: S
      KeySchema:
        - AttributeName: ShowId
          KeyType: HASH
      GlobalSecondaryIndexes:
        - IndexName: TitleGSI
          KeySchema:
            - AttributeName: Title
              KeyType: HASH
          Projection:
            NonKeyAttributes:
              - ShowId
            ProjectionType: "INCLUDE"

  RatingTable:
    Type: AWS::DynamoDB::Table
    DeletionPolicy: Retain
    Properties:
      BillingMode: PAY_PER_REQUEST
      TableClass: STANDARD
      TableName: !Sub ${AWS::StackName}-RatingsV1
      AttributeDefinitions:
        - AttributeName: ShowId
          AttributeType: S
        - AttributeName: UserId
          AttributeType: S
      KeySchema:
        - AttributeName: ShowId
          KeyType: HASH
        - AttributeName: UserId
          KeyType: RANGE

  PredictionTable:
    Type: AWS::DynamoDB::Table
    DeletionPolicy: Retain
    Properties:
      BillingMode: PAY_PER_REQUEST
      TableClass: STANDARD
      TableName: !Sub ${AWS::StackName}-PredictionsV1
      AttributeDefinitions:
        - AttributeName: PredictionId
          AttributeType: S
      KeySchema:
        - AttributeName: PredictionId
          KeyType: HASH

  ModelTrainingRepository:
    Type: AWS::ECR::Repository
    Properties:
      RepositoryName: !Sub ${AWS::StackName}/modeltrainingrepo
      RepositoryPolicyText:
        Version: "2012-10-17"
        Statement:
          - Sid: AllowSageMakerAccess
            Effect: Allow
            Principal:
              Service: "sagemaker.amazonaws.com"
            Action:
              - "ecr:*"
              - "sagemaker:*"

  ModelBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub ${AWS::StackName}-modelbucket

  RatingsExportBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub ${AWS::StackName}-ratingsexport

Outputs:
  ShowTableArn:
    Value: !GetAtt ShowTable.Arn
    Export:
      Name: !Sub ${AWS::StackName}-ShowTableArn
  ShowTableName:
    Value: !Ref ShowTable
    Export:
      Name: !Sub ${AWS::StackName}-ShowTableName
  PredictionTableArn:
    Value: !GetAtt PredictionTable.Arn
    Export:
      Name: !Sub ${AWS::StackName}-PredictionTableArn
  PredictionTableName:
    Value: !Ref PredictionTable
    Export:
      Name: !Sub ${AWS::StackName}-PredictionTableName
  RatingTableArn:
    Value: !GetAtt RatingTable.Arn
    Export:
      Name: !Sub ${AWS::StackName}-RatingTableArn
  RatingTableName:
    Value: !Ref RatingTable
    Export:
      Name: !Sub ${AWS::StackName}-RatingTableName
  ModelBucket:
    Value: !GetAtt ModelBucket.Arn
    Export:
      Name: !Sub ${AWS::StackName}-ModelBucketArn
  ModelBucketName:
    Value: !Ref ModelBucket
    Export:
      Name: !Sub ${AWS::StackName}-ModelBucketName
  RatingsExportBucket:
    Value: !GetAtt RatingsExportBucket.Arn
    Export:
      Name: !Sub ${AWS::StackName}-RatingsExportBucketArn
  SageMakerExecutionRole:
    Value: !GetAtt SageMakerExecutionRole.Arn
    Export:
      Name: !Sub ${AWS::StackName}-SageMakerExecutionRoleArn