AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: >
  SAM App to demo Step functions and Sagemaker CICD pipe line
Parameters:
  StateMachineName:
    Description: Name of State Machine
    Type: String
    Default: MLOpsStepFunction
  S3ModelBucket:
    Description: Bucket to store model and data
    Type: String
    Default: mlops-cicd
Resources:
  SNSApprovalEmail:
    Type: AWS::SNS::Topic
    Properties:       
      TopicName: SNSApprovalEmail     
  HttpAPI:
    Type: AWS::Serverless::HttpApi
    Properties:      
      StageName: v1      
  MLOpsStateMachine:
    Type: AWS::Serverless::StateMachine # More info about State Machine Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-statemachine.html
    Properties:
      Name: !Sub ${StateMachineName}-${AWS::AccountId}
      DefinitionUri: ../statemachine/mlops.asl.json
      DefinitionSubstitutions:
        CreateAndEmailLinkFnName: !Ref CreateAndEmailLinks
        UpdateSagemakerEndpointAPI: !GetAtt UpdateSagemakerEndpointAPI.Arn
        SagemakerRoleArn: !GetAtt SageMakerRole.Arn
        S3ModelBucket: !Ref S3ModelBucket
      Role: !GetAtt StepFunctionsRole.Arn
  CreateAndEmailLinks:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: ../functions/create_and_email_accept_reject_links/
      Handler: app.lambda_handler
      Runtime: python3.10
      Policies:
      - SNSPublishMessagePolicy:
          TopicName: SNSApprovalEmail      
      Environment:
        Variables:
          HttpApiID: !Ref HttpAPI
          SNSArn: !Ref SNSApprovalEmail
  UpdateSagemakerEndpointAPI:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      CodeUri: ../functions/update_sagemakerEndpoint_API/
      Handler: app.lambda_handler
      Runtime: python3.10
      Timeout: 600
      Policies:
      - SNSPublishMessagePolicy:
          TopicName: SNSApprovalEmail
      - Version: '2012-10-17' # Policy Document
        Statement:
          - Effect: Allow
            Action:
              - lambda:UpdateFunctionConfiguration
            Resource: !Sub arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${TestSagemakerEndpoint}
      - Version: '2012-10-17' # Policy Document
        Statement:
          - Effect: Allow
            Action:
              - sagemaker:DescribeEndpoint
            Resource: '*'  
      Environment:
        Variables:
          Endpoint_FunctionName: !Ref TestSagemakerEndpoint
          HttpApiID: !Ref HttpAPI
          SNSArn: !Ref SNSApprovalEmail          
  RespondEmailLinks:    
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      CodeUri: ../functions/respond_to_links/
      Handler: app.lambda_handler
      Runtime: python3.10
      Policies:
      - Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Action:
            - states:DescribeActivity
            - states:DeleteActivity
            - states:GetActivityTask
            - states:SendTaskSuccess
            - states:SendTaskFailure
            - states:SendTaskHeartbeat          
          Resource: !Sub "arn:${AWS::Partition}:states:${AWS::Region}:${AWS::AccountId}:stateMachine:${StateMachineName}-${AWS::AccountId}"         
      Events:
        ExplicitApi: # warning: creates a public endpoint
          Type: HttpApi
          Properties:
            ApiId: !Ref HttpAPI
            Method: GET
            Path: /respond
  TestSagemakerEndpoint:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      CodeUri: ../functions/api_sagemaker_endpoint/
      Handler: app.lambda_handler
      Runtime: python3.10
      Policies:
        - Version: '2012-10-17' # Policy Document
          Statement:
            - Effect: Allow
              Action:
                - sagemaker:InvokeEndpoint
              Resource: '*'         
      Events:
        ExplicitApi:
          Type: HttpApi
          Properties:
            ApiId: !Ref HttpAPI
            Method: POST
            Path: /invokeSagemakerAPI
  SageMakerRole:
      Type: AWS::IAM::Role
      Properties:       
        AssumeRolePolicyDocument: 
          Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Principal:
                Service:
                  - sagemaker.amazonaws.com
              Action:
                - 'sts:AssumeRole'
        Description: String
        ManagedPolicyArns:
          - !Ref SagemakerPolicy    
        Path: /      
  SagemakerPolicy:
      Type: 'AWS::IAM::ManagedPolicy'
      Properties:
        Description: Step function use policy
        PolicyDocument: !Sub
            - |-
              {
                      "Version": "2012-10-17",
                      "Statement": [
                        {
                              "Effect": "Allow",
                              "Action": [
                                  "cloudwatch:PutMetricData",                                  
                                  "logs:CreateLogGroup",
                                  "logs:CreateLogStream",
                                  "logs:DescribeLogStreams",
                                  "logs:PutLogEvents",
                                  "logs:GetLogEvents",                                                                                             
                                  "sagemaker:*"                                                             
                              ],
                              "Resource": [
                                  "*"
                              ]
                        },
                        {
                            "Effect": "Allow",
                            "Action": [
                                "iam:PassRole"
                            ],
                            "Resource": "*",
                            "Condition": {
                                "StringEquals": {
                                    "iam:PassedToService": "sagemaker.amazonaws.com"
                                }
                            }
                        },
                        {
                            "Effect": "Allow",
                            "Action": [
                                "ecr:GetAuthorizationToken",
                                "ecr:BatchCheckLayerAvailability",
                                "ecr:GetDownloadUrlForLayer",
                                "ecr:BatchGetImage"
                            ],
                            "Resource": ["*"]
                        },                        
                        {
                            "Effect": "Allow",
                            "Action": [
                                "s3:ListBucket",
                                "s3:GetObject",
                                "s3:PutObject",
                                "s3:DeleteObject"
                            ],
                            "Resource": [
                                "arn:aws:s3:::${S3ModelBucket}",
                                "arn:aws:s3:::${S3ModelBucket}/*"
                            ]
                        }
                      ]
              }
            - {
                AccountID: !Ref AWS::AccountId, 
                Region: !Ref AWS::Region,
                S3ModelBucket: !Ref S3ModelBucket  
              } 
  StepFunctionsRole:
      Type: AWS::IAM::Role
      Properties:       
        AssumeRolePolicyDocument: 
          Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Principal:
                Service:
                  - states.amazonaws.com
              Action:
                - 'sts:AssumeRole'
        Description: String
        ManagedPolicyArns:
          - !Ref StepFunctionsPolicy    
        Path: /      
  StepFunctionsPolicy:
      Type: 'AWS::IAM::ManagedPolicy'
      DependsOn: SageMakerRole
      Properties:
        Description: Step function use policy
        PolicyDocument: !Sub
            - |-
              {
                      "Version": "2012-10-17",
                      "Statement": [
                        {
                              "Effect": "Allow",
                              "Action": [
                                  "lambda:InvokeFunction",                                  
                                  "states:StartExecution"                                                                 
                              ],
                              "Resource": [
                                  "*"
                              ]
                          },                          
                          {
                            "Effect": "Allow",
                            "Action": [
                              "sagemaker:CreateTrainingJob",
                              "sagemaker:DescribeTrainingJob",
                              "sagemaker:StopTrainingJob",
                              "sagemaker:CreateModel",
                              "sagemaker:CreateTransformJob",
                              "sagemaker:CreateEndpointConfig",
                              "sagemaker:CreateEndpoint",
                              "sagemaker:AddTags"
                            ],
                            "Resource": [
                              "arn:aws:sagemaker:${Region}:${AccountID}:training-job/*",
                              "arn:aws:sagemaker:${Region}:${AccountID}:model/*",
                              "arn:aws:sagemaker:${Region}:${AccountID}:transform-job/*",
                              "arn:aws:sagemaker:${Region}:${AccountID}:endpoint-config/*",
                              "arn:aws:sagemaker:${Region}:${AccountID}:endpoint/*"
                            ]
                          },
                          {
                            "Effect": "Allow",
                            "Action": [
                              "sagemaker:ListTags"
                            ],
                            "Resource": [
                              "*"
                            ]
                          },
                          {
                            "Effect": "Allow",
                            "Action": [
                              "iam:PassRole"
                            ],
                            "Resource": [
                              "${SageMakerRole}"
                            ],
                            "Condition": {
                              "StringEquals": {
                                "iam:PassedToService": "sagemaker.amazonaws.com"
                              }
                            }
                          },
                          {
                            "Effect": "Allow",
                            "Action": [
                              "events:PutTargets",
                              "events:PutRule",
                              "events:DescribeRule"
                            ],
                            "Resource": [
                              "arn:aws:events:${Region}:${AccountID}:rule/StepFunctionsGetEventsForSageMakerTrainingJobsRule",
                              "arn:aws:events:${Region}:${AccountID}:rule/StepFunctionsGetEventsForSageMakerTransformJobsRule"
                            ]
                          }  
                      ]
                }
            - {
                AccountID: !Ref AWS::AccountId, 
                Region: !Ref AWS::Region,
                SageMakerRole: !GetAtt SageMakerRole.Arn
                
              }
Outputs:
  MLOpsStateMachineArn:
    Description: "State machine ARN"
    Value: !Ref MLOpsStateMachine
  SNSApprovalEmail:
    Description: "SNS ARN"
    Value: !Ref SNSApprovalEmail