AWSTemplateFormatVersion: 2010-09-09
Description: Template to provision OpenSearch cluster and SageMaker Notebook for hybrid search

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: Required Parameters
        Parameters:
          - AppName
    ParameterLabels:
      AppName:
        default: Name of the overall application
      SageMakerNotebookName:
        default: Name of SageMaker Notebook Instance
      
      
Parameters:
  InstanceType:
    Type: String
    Description: Choosing c5.large
    Default: c5.large 
  SageMakerIAMRole:
    Description: Name of IAM role that will be created by this cloud formation template. The role name must _not_ already exist in your AWS account.
    Type: String
    Default: "HybridSearchIAMRole"  
  AppName:
    Default: opensearch-hybrid-search
    Type: String
    AllowedValues: [opensearch-hybrid-search]
    Description: Name of the overall application, this is used while creating the ML model endpoint.
  

Mappings: 
  RegionMap: 
    us-east-1:
      HVM64: ami-041feb57c611358bd
    us-west-2:
      HVM64: ami-07d07d65c47e5aa90
    eu-west-1:
      HVM64: ami-0b9fd8b55a6e3c9d5

Resources:
  DemoInstance:
    Type: 'AWS::EC2::Instance'
    Properties: 
      ImageId: !FindInMap [RegionMap, !Ref "AWS::Region", HVM64]
      InstanceType: !Ref InstanceType
      SecurityGroups: 
        - !Ref DemoSecurityGroup
      IamInstanceProfile: !Ref Ec2InstanceProfile
      UserData:
        Fn::Base64: 
          Fn::Sub:
            - |
              #!/bin/bash
              cd /home/ec2-user
              sudo yum install awscli
              sudo yum install zip
              sudo yum -y install git
              git clone https://github.com/aws-samples/AI-search-with-amazon-opensearch-service.git
              aws s3 cp /home/ec2-user/amazon-sagemaker-bert-finetuning-for-search/opensearch-hybridsearch/hybridsearch.zip s3://${s3Bucket} 
            - s3Bucket: !Ref s3BucketHosting
  
  

  DemoSecurityGroup:
    Type: 'AWS::EC2::SecurityGroup'
    Properties:
      GroupDescription: SG to allow web access 
      
  Ec2InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties: 
      InstanceProfileName: WebServiceApp
      Roles: 
        - !Ref NBRole
  

  CodeRepository:
    Type: AWS::SageMaker::CodeRepository
    Properties:
      GitConfig:
        RepositoryUrl: https://github.com/aws-samples/AI-search-with-amazon-opensearch-service.git

  NotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      NotebookInstanceName: hybrid-search-nb
      InstanceType: ml.m5d.2xlarge
      RoleArn: !GetAtt NBRole.Arn
      DefaultCodeRepository: !GetAtt CodeRepository.CodeRepositoryName

  OpenSearchSagemakerRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: opensearch-sagemaker-role
      Policies:
        - PolicyName: SageMakerAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "sagemaker:InvokeEndpointAsync"
                  - "sagemaker:InvokeEndpoint"
                Resource:
                  - "*"
              - Effect: Allow
                Action:
                  - "bedrock:*"
                Resource:
                  - "*"
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - opensearchservice.amazonaws.com
            Action:
              - 'sts:AssumeRole'

  NBRole:
    Type: AWS::IAM::Role
    Properties:
      Policies:
        - PolicyName: CustomNotebookAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "es:ESHttp*"
                Resource:
                  - !Sub arn:aws:es:${AWS::Region}:${AWS::AccountId}:domain/*
              - Effect: Allow
                Action:
                  - "s3:*"
                  - "dynamodb:*"
                Resource:
                  - "*"
              - Effect: Allow
                Action:
                  - "iam:PassRole"
                Resource:
                  - !Sub arn:aws:iam::${AWS::AccountId}:role/opensearch-sagemaker-role
        - PolicyName: lambda-ssm-full-access 
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "lambda:*"
                Resource:
                  - "*"
              - Effect: Allow
                Action:
                  - "ssm:*"
                Resource:
                  - "*"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AWSCloudFormationReadOnlyAccess
        - arn:aws:iam::aws:policy/AmazonEC2FullAccess
        - arn:aws:iam::aws:policy/AmazonOpenSearchServiceFullAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/IAMFullAccess
        - arn:aws:iam::aws:policy/AmazonS3FullAccess
        - arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
        - arn:aws:iam::aws:policy/AmazonSSMFullAccess
      MaxSessionDuration: 43200
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
          - Effect: Allow
            Principal:
              Service:
              - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
          - Effect: Allow
            Principal:
              Service:
              - ec2.amazonaws.com
            Action:
              - 'sts:AssumeRole'
          

  Role:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Ref SageMakerIAMRole
      Policies:
        - PolicyName: CustomNotebookAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Sid: ReadFromOpenSearch
                Effect: Allow
                Action:
                  - "es:ESHttp*"
                Resource:
                  - !Sub arn:aws:es:${AWS::Region}:${AWS::AccountId}:domain/*
              - Sid: ReadWriteFromECR
                Effect: Allow
                Action:
                  - "ecr:BatchGetImage"
                  - "ecr:BatchCheckLayerAvailability"
                  - "ecr:CompleteLayerUpload"
                  - "ecr:DescribeImages"
                  - "ecr:DescribeRepositories"
                  - "ecr:GetDownloadUrlForLayer"
                  - "ecr:InitiateLayerUpload"
                  - "ecr:ListImages"
                  - "ecr:PutImage"
                  - "ecr:UploadLayerPart"
                  - "ecr:CreateRepository"
                  - "ecr:GetAuthorizationToken"
                  - "ec2:DescribeAvailabilityZones"
                Resource: "*"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/AWSCloudFormationReadOnlyAccess
        - arn:aws:iam::aws:policy/TranslateReadOnly
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'

  s3BucketHosting:
    Type: AWS::S3::Bucket
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: "AES256"
      VersioningConfiguration:
        Status: Enabled
      WebsiteConfiguration:
        IndexDocument: index.html
        ErrorDocument: error.html

  OpenSearchServiceDomain:
    Type: AWS::OpenSearchService::Domain
    Properties:
      AccessPolicies:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              AWS: '*'
            Action: 'es:*'
            Resource: !Sub arn:aws:es:${AWS::Region}:${AWS::AccountId}:domain/*/*
      EngineVersion: 'OpenSearch_2.11'
      ClusterConfig:
        InstanceType: "r6g.2xlarge.search"
      EBSOptions:
        EBSEnabled: True
        VolumeSize: 20
        VolumeType: 'gp3'
      AdvancedSecurityOptions:
        AnonymousAuthEnabled: False
        Enabled: True
        InternalUserDatabaseEnabled: False
        MasterUserOptions:
          MasterUserARN: !GetAtt NBRole.Arn
      NodeToNodeEncryptionOptions:
        Enabled: True
      EncryptionAtRestOptions:
        Enabled: True
        KmsKeyId: alias/aws/es
      DomainEndpointOptions:
        EnforceHTTPS: True
      Tags:
        - {"Key" : "app","Value" : "hybrid-search"}




  EmbeddingEndpoint:
    Type: "AWS::SageMaker::Endpoint"
    Properties:
      EndpointName: !Sub
      - '${AppName}-embedding-gpt-j-6b-${RandomGUID}'
      - { RandomGUID: !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref AWS::StackId ]]]] }
      EndpointConfigName: !GetAtt EmbeddingEndpointConfig.EndpointConfigName

  EmbeddingEndpointConfig:
    Type: "AWS::SageMaker::EndpointConfig"
    Properties:
      EndpointConfigName: !Sub
      - '${AppName}-embedding-gpt-j-6b-endppoint-${RandomGUID}'
      - { RandomGUID: !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref AWS::StackId ]]]] }
      ProductionVariants:
        - InitialInstanceCount: 1
          InitialVariantWeight: 1.0
          InstanceType: "ml.g5.2xlarge"
          ModelName: !GetAtt EmbeddingModel.ModelName
          VariantName: !GetAtt EmbeddingModel.ModelName
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W1200
            reason: Solution does not have KMS encryption enabled by default

  EmbeddingModel:
    Type: "AWS::SageMaker::Model"
    Properties:
      ModelName: !Sub
      - '${AppName}-embedding-gpt-j-6b-model-${RandomGUID}'
      - { RandomGUID: !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref AWS::StackId ]]]] }
      PrimaryContainer:
        ModelDataUrl: !Sub "s3://jumpstart-cache-prod-${AWS::Region}/huggingface-infer/prepack/v1.0.1/infer-prepack-huggingface-textembedding-gpt-j-6b-fp16.tar.gz"
        Image: !Sub "763104351884.dkr.ecr.${AWS::Region}.amazonaws.com/pytorch-inference:1.12.0-gpu-py38"
        Environment: {"TS_DEFAULT_WORKERS_PER_MODEL": "1"}
        Mode: "SingleModel"
      ExecutionRoleArn: !GetAtt Role.Arn

  
Outputs:

  OpenSearchDomainEndpoint:
    Description: OpenSearch domain endpoint
    Value:
      'Fn::GetAtt':
        - OpenSearchServiceDomain
        - DomainEndpoint

  OpenSearchDomainName:
    Description: OpenSearch domain name
    Value: !Ref OpenSearchServiceDomain

  Region:
    Description: Deployed Region
    Value: !Ref AWS::Region


  s3BucketHostingBucketName:
    Description: S3 bucket name for frontend hosting
    Value: !Ref s3BucketHosting

  S3BucketSecureURL:
    Value: !Join
      - ''
      - - 'https://'
        - !GetAtt
          - s3BucketHosting
          - DomainName
    Description: Name of S3 bucket to hold website content

  SageMakerNotebookURL:
    Description: SageMaker Notebook Instance
    Value: !Join
      - ''
      - - 'https://console.aws.amazon.com/sagemaker/home?region='
        - !Ref AWS::Region
        - '#/notebook-instances/openNotebook/'
        - !GetAtt NotebookInstance.NotebookInstanceName
        - '?view=classic'

      
  EmbeddingEndpointName:
    Description: Name of the LLM endpoint
    Value: !GetAtt EmbeddingEndpoint.EndpointName

  SecurityGroupID:
    Description: ID of the Security group of the ec2 instance
    Value: !GetAtt DemoSecurityGroup.GroupId

  

   

