AWSTemplateFormatVersion: '2010-09-09'
Description: 'CloudFormation template for AutoGluon-Cloud to run on SageMaker'

Resources:
  AGCloudSageMakerBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName: !Sub
        - '${AWS::StackName}-bucket-${suffix}'
        - suffix: !Select [0, !Split ['-', !Select [2, !Split ['/', !Ref 'AWS::StackId']]]]
      VersioningConfiguration:
        Status: Enabled
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true

  AGCloudSageMakerExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Sub '${AWS::StackName}-execution-role'
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: ''
            Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
              AWS: !Sub 'arn:aws:iam::${AWS::AccountId}:root'
            Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/AmazonSageMakerFullAccess'
      Policies:
        - PolicyName: !Sub '${AWS::StackName}-custom-policy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: SageMaker
                Effect: Allow
                Action:
                  - sagemaker:DescribeEndpoint
                  - sagemaker:DescribeEndpointConfig
                  - sagemaker:DescribeModel
                  - sagemaker:DescribeTrainingJob
                  - sagemaker:DescribeTransformJob
                  - sagemaker:CreateArtifact
                  - sagemaker:CreateEndpoint
                  - sagemaker:CreateEndpointConfig
                  - sagemaker:CreateModel
                  - sagemaker:CreateTrainingJob
                  - sagemaker:CreateTransformJob
                  - sagemaker:DeleteEndpoint
                  - sagemaker:DeleteEndpointConfig
                  - sagemaker:DeleteModel
                  - sagemaker:UpdateArtifact
                  - sagemaker:UpdateEndpoint
                  - sagemaker:InvokeEndpoint
                  - sagemaker:ListTags
                Resource:
                  - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:endpoint/ag-cloudpredictor*'
                  - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:endpoint-config/ag-cloudpredictor*'
                  - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:model/autogluon-inference*'
                  - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:training-job/ag-cloudpredictor*'
                  - !Sub 'arn:aws:sagemaker:*:${AWS::AccountId}:transform-job/ag-cloudpredictor*'

              - Sid: IAM
                Effect: Allow
                Action:
                  - iam:PassRole
                Resource:
                  - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*'

              - Sid: CloudWatchDescribe
                Effect: Allow
                Action:
                  - logs:DescribeLogStreams
                Resource:
                  - !Sub 'arn:aws:logs:*:${AWS::AccountId}:log-group:*'

              - Sid: CloudWatchGet
                Effect: Allow
                Action:
                  - logs:GetLogEvents
                Resource:
                  - !Sub 'arn:aws:logs:*:${AWS::AccountId}:log-group:*:log-stream:*'

              - Sid: S3Object
                Effect: Allow
                Action:
                  - s3:PutObject
                  - s3:PutObjectAcl
                  - s3:GetObject
                  - s3:GetObjectAcl
                  - s3:AbortMultipartUpload
                Resource:
                  - !Sub 'arn:aws:s3:::${AGCloudSageMakerBucket}/*'
                  - !Sub 'arn:aws:s3:::${AGCloudSageMakerBucket}'
                  - 'arn:aws:s3:::*SageMaker*'
                  - 'arn:aws:s3:::*Sagemaker*'
                  - 'arn:aws:s3:::*sagemaker*'

              - Effect: Allow
                Action:
                  - s3:CreateBucket
                  - s3:GetBucketLocation
                  - s3:ListBucket
                  - s3:GetBucketCors
                  - s3:PutBucketCors
                  - s3:GetBucketAcl
                  - s3:PutObjectAcl
                Resource:
                  - !Sub 'arn:aws:s3:::${AGCloudSageMakerBucket}/*'
                  - !Sub 'arn:aws:s3:::${AGCloudSageMakerBucket}'
                  - 'arn:aws:s3:::*SageMaker*'
                  - 'arn:aws:s3:::*Sagemaker*'
                  - 'arn:aws:s3:::*sagemaker*'

              - Sid: ListEvents
                Effect: Allow
                Action:
                  - s3:ListAllMyBuckets
                  - sagemaker:ListEndpointConfigs
                  - sagemaker:ListEndpoints
                  - sagemaker:ListTransformJobs
                  - sagemaker:ListTrainingJobs
                  - sagemaker:ListModels
                  - sagemaker:ListDomains
                Resource:
                  - '*'

              - Effect: Allow
                Action:
                  - sagemaker:CreateFlowDefinition
                  - sagemaker:DescribeFlowDefinition
                  - sagemaker:DeleteFlowDefinition
                  - sagemaker:ListFlowDefinitions
                Resource:
                  - 'arn:aws:sagemaker:*:*:flow-definition/*'
                Condition:
                  StringEqualsIfExists:
                    'sagemaker:WorkteamType':
                      - private-crowd
                      - vendor-crowd

              - Sid: Others
                Effect: Allow
                Action:
                  - ecr:BatchGetImage
                  - ecr:Describe*
                  - ecr:GetAuthorizationToken
                  - ecr:GetDownloadUrlForLayer
                  - logs:CreateLogDelivery
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:DeleteLogDelivery
                  - logs:Describe*
                  - logs:GetLogDelivery
                  - logs:GetLogEvents
                  - logs:ListLogDeliveries
                  - logs:PutLogEvents
                  - logs:PutResourcePolicy
                  - logs:UpdateLogDelivery
                Resource:
                  - '*'

Outputs:
  BucketName:
    Description: S3 bucket where AutoGluon-Cloud will save trained predictors
    Value: !Ref AGCloudSageMakerBucket

  RoleARN:
    Description: ARN of the created IAM role for AutoGluon-Cloud to run on SageMaker
    Value: !GetAtt AGCloudSageMakerExecutionRole.Arn