# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

Description: This template will deploy a SageMaker Endpoint

Parameters:

  SageMakerProjectName:
    Type: String
    Description: Name of the project
    MinLength: 1
    MaxLength: 32
    AllowedPattern: ^[a-zA-Z](-*[a-zA-Z0-9])*

  SageMakerProjectId:
    Type: String
    Description: Service generated Id of the project.

  ModelPackageName:
    Type: String
    Description: The trained Model Package Name

  EndpointInstanceCount:
    Type: Number
    Description: Number of instances to launch for the endpoint.
    MinValue: 1

  EndpointInstanceType:
    Type: String
    Description: The ML compute instance type for the endpoint.

  ExecutionRoleName:
    Type: String
    Description: Name of the execution role to run the model

  EnvName:
    Type: String
    AllowedPattern: '[a-z0-9\-]*'
    Description: Please specify your data science environment name.  Used as a suffix for environment resource names.

  EnvType:
    Description: System Environment (e.g. dev, staging, prod). Used as a suffix for environment resource names.
    Type: String
    Default: dev
  
  # SSM parameter key is "${EnvName}-${EnvType}-private-subnet-ids"
  SageMakerSubnetIds:
    Type: AWS::SSM::Parameter::Value<String>
    Description: VPC subnets

  # SSM parameter key is "${EnvName}-${EnvType}-sagemaker-sg-ids"
  SageMakerSecurityGroupIds:
    Type: AWS::SSM::Parameter::Value<String>
    Description: Security groups for SageMaker Studio and SageMaker workloads
  
  VolumeKmsKeyArn:
    Type: String
    Description: The ARN of a KMS key that Amazon SageMaker uses to encrypt data on the storage volume attached endpoint

  Accounts:
    Description: Comma-delimited list of target accounts where the model will be deployed
    Type: String
    Default: ''

Resources:
  SageMakerModel:
    Type: AWS::SageMaker::Model
    Properties:
      PrimaryContainer:
        ModelPackageName: !Ref ModelPackageName
      ExecutionRoleArn: !Sub 'arn:aws:iam::${AWS::AccountId}:role/${ExecutionRoleName}'
      VpcConfig:
        SecurityGroupIds: !Split [',', !Ref SageMakerSecurityGroupIds]
        Subnets: !Split [',', !Ref SageMakerSubnetIds]

      Tags:
        - Key: sagemaker:deployment-stage
          Value: !Ref EnvType
        - Key: sagemaker:project-id
          Value: !Ref SageMakerProjectId
        - Key: sagemaker:project-name
          Value: !Ref SageMakerProjectName
        - Key: EnvironmentName
          Value: !Ref EnvName
        - Key: EnvironmentType
          Value: !Ref EnvType

  SageMakerEndpointConfig:
    Type: AWS::SageMaker::EndpointConfig
    Properties:
      ProductionVariants:
        - InitialInstanceCount: !Ref EndpointInstanceCount
          InitialVariantWeight: 1.0
          InstanceType: !Ref EndpointInstanceType
          ModelName: !GetAtt SageMakerModel.ModelName
          VariantName: AllTraffic
      # KmsKeyId: !Ref VolumeKmsKeyArn
      Tags:
        - Key: sagemaker:deployment-stage
          Value: !Ref EnvType
        - Key: sagemaker:project-id
          Value: !Ref SageMakerProjectId
        - Key: sagemaker:project-name
          Value: !Ref SageMakerProjectName
        - Key: EnvironmentName
          Value: !Ref EnvName
        - Key: EnvironmentType
          Value: !Ref EnvType

  SageMakerEndpoint:
    Type: AWS::SageMaker::Endpoint
    Properties:
      EndpointName: !Sub '${SageMakerProjectName}-${SageMakerProjectId}-${EnvType}'
      EndpointConfigName: !GetAtt SageMakerEndpointConfig.EndpointConfigName
      Tags:
        - Key: sagemaker:deployment-stage
          Value: !Ref EnvType
        - Key: sagemaker:project-id
          Value: !Ref SageMakerProjectId
        - Key: sagemaker:project-name
          Value: !Ref SageMakerProjectName
        - Key: EnvironmentName
          Value: !Ref EnvName
        - Key: EnvironmentType
          Value: !Ref EnvType