whisper-prod.yaml

AWSTemplateFormatVersion: '2010-09-09'
Description: Whisper on EC2 with S3, SQS, Lambda, and Auto Scaling

Resources:
  MyVPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: "10.0.0.0/16"
      Tags:
        - Key: Name
          Value: "My VPC"

  MyInternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: "My IGW"

  MyVPCGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref MyVPC
      InternetGatewayId: !Ref MyInternetGateway

  MyPublicSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      CidrBlock: "10.0.1.0/24"
      VpcId: !Ref MyVPC
      AvailabilityZone: !Select [0, !GetAZs ""]
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: "My Public Subnet"

  MyRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref MyVPC
      Tags:
        - Key: Name
          Value: "My Route Table"

  MyDefaultRoute:
    Type: AWS::EC2::Route
    DependsOn: MyVPCGatewayAttachment
    Properties:
      RouteTableId: !Ref MyRouteTable
      DestinationCidrBlock: "0.0.0.0/0"
      GatewayId: !Ref MyInternetGateway

  MySubnetRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref MyPublicSubnet
      RouteTableId: !Ref MyRouteTable

  MySecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupName: "My Security Group"
      GroupDescription: "My security group"
      VpcId: !Ref MyVPC
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 8501
          ToPort: 8501
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 9000
          ToPort: 9000
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 8000
          ToPort: 8000
          CidrIp: 0.0.0.0/0

  InstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AdministratorAccess
  MyInstanceProfile: 
    Type: "AWS::IAM::InstanceProfile"
    Properties: 
      Path: "/"
      Roles: 
        - !Ref InstanceRole
  MyS3Bucket:
    Type: AWS::S3::Bucket
    DependsOn: LambdaPermission
    Properties:
      BucketName: !Sub 'whisper-bucket-${AWS::AccountId}'

  MySQSQueue:
    Type: AWS::SQS::Queue
    Properties:
      QueueName: whisper-queue

  MySQSQueuePolicy:
    Type: AWS::SQS::QueuePolicy
    Properties:
      PolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              AWS: !GetAtt LambdaExecutionRole.Arn
            Action:
              - sqs:SendMessage
            Resource: !GetAtt MySQSQueue.Arn
      Queues:
        - !Ref MySQSQueue

  LambdaExecutionRole:
    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
        - arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess
        - arn:aws:iam::aws:policy/AmazonSQSFullAccess

  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import json
          import boto3
          import urllib.parse
          import os

          s3 = boto3.client('s3')
          sqs = boto3.client('sqs')

          def lambda_handler(event, context):
              # 获取S3事件记录
              bucket_name = event['Records'][0]['s3']['bucket']['name']
              file_key = urllib.parse.unquote_plus(event['Records'][0]['s3']['object']['key'], encoding='utf-8')

              queue_url = os.environ['QUEUE_URL']

              try:
                  # 判断文件后缀，如果文件不是以 .mp3 .mp4 .m4a 结尾的，则直接return
                  file_extension = file_key.split('.')[-1].lower()
                  allowed_extensions = ['mp3', 'mp4', 'm4a']
                  if file_extension not in allowed_extensions:
                      return {
                      'statusCode': 200,
                      'body': json.dumps('File {file_key} is not an audio or video file. Skipping...')
                  }

                  # 获取文件标签
                  response = s3.get_object_tagging(
                      Bucket=bucket_name,
                      Key=file_key
                  )
                  tags = {tag['Key']: tag['Value'] for tag in response['TagSet']}

                  # 构建SQS消息体
                  message_body = {
                      'bucket': bucket_name,
                      'key': file_key,
                      'tags': tags
                  }

                  # 发送SQS消息
                  sqs.send_message(
                      QueueUrl=queue_url,
                      MessageBody=json.dumps(message_body)
                  )
              except Exception as e:
                  print(e)
                  raise e

      Handler: index.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.10
      Environment:
        Variables:
          QUEUE_URL: !Sub "https://sqs.${AWS::Region}.amazonaws.com/${AWS::AccountId}/whisper-queue"
      FunctionName: !Sub 'whisper-s3-event-${AWS::AccountId}'

  LambdaPermission:
    Type: AWS::Lambda::Permission
    DependsOn: LambdaFunction
    Properties:
      FunctionName: !GetAtt LambdaFunction.Arn
      Action: 'lambda:InvokeFunction'
      Principal: s3.amazonaws.com
      SourceAccount: !Ref 'AWS::AccountId'
      SourceArn: !Sub 'arn:aws:s3:::${BucketName}'

  MyAutoScalingLaunchTemplate:
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateName: MyAutoScalingLaunchTemplate
      LaunchTemplateData:
        ImageId: !FindInMap [RegionMap, !Ref "AWS::Region", AMI]
        InstanceType: !Ref InstanceType
        KeyName: !Ref KeyPairName
        SecurityGroupIds: 
          - !Ref MySecurityGroup
        IamInstanceProfile:
          Arn: !GetAtt MyInstanceProfile.Arn
        BlockDeviceMappings:
          - DeviceName: /dev/sda1
            Ebs:
              VolumeSize: 100
              VolumeType: gp3
        UserData:
          'Fn::Base64': !Sub |
            #!/bin/bash
            sudo apt update
            sudo apt install -y python3-pip nvidia-driver-525 ffmpeg git nvidia-cuda-toolkit awscli
            git clone https://github.com/yanjun-ios/whisper-on-aws-jumpstart /home/ubuntu/whisper
            cd /home/ubuntu/whisper
            pip3 install -r requirements.txt
            # Setup environment  
            echo "HF_TOKEN=${HuggingFaceToken}" >> /home/ubuntu/whisper/.env
            echo "LAUNCH_API=${LaunchAPI}" >> /home/ubuntu/whisper/.env
            echo "LAUNCH_DEMO=${LaunchDemo}" >> /home/ubuntu/whisper/.env
            echo "LAUNCH_EVENT=${LaunchEvent}" >> /home/ubuntu/whisper/.env
            echo "BUCKET_NAME=${BucketName}" >> /home/ubuntu/whisper/.env
            echo "LAMBDA_NAME=whisper-s3-event-${AWS::AccountId}" >> /home/ubuntu/whisper/.env
            echo "SQS_QUEUE_URL=$(aws sqs get-queue-url --queue-name whisper-queue --region ${AWS::Region} --output text --query QueueUrl)" >> /home/ubuntu/whisper/.env
            echo "------Install Finished ------"
            # Add permissions to whisper directory
            sudo chown -R ubuntu:ubuntu /home/ubuntu/whisper/
            sudo chmod -R 755 /home/ubuntu/whisper/
            # Create systemd service
            sudo tee /etc/systemd/system/whisper.service > /dev/null <<EOF
            [Unit]
            Description=Whisper Service
            After=network.target
            [Service]
            User=ubuntu
            ExecStart=/bin/bash -c 'sh /home/ubuntu/whisper/launch.sh'
            WorkingDirectory=/home/ubuntu/whisper
            Restart=always
            [Install]
            WantedBy=multi-user.target
            EOF
            # Reload systemd and enable service
            sudo systemctl daemon-reload
            sudo systemctl enable whisper.service
            # Start the service
            sudo systemctl start whisper.service
            sudo reboot

  MyAutoScalingGroup:
    Type: AWS::AutoScaling::AutoScalingGroup
    Properties:
      VPCZoneIdentifier: 
        - !Ref MyPublicSubnet
      LaunchTemplate:
        LaunchTemplateId: !Ref MyAutoScalingLaunchTemplate
        Version: !GetAtt MyAutoScalingLaunchTemplate.LatestVersionNumber
      MinSize: 1
      MaxSize: 4
      DesiredCapacity: 1
      TargetGroupARNs:
        - !Ref MyTargetGroup
        - !Ref MyTargetGroupApi

  MyTargetGroup:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    Properties:
      VpcId: !Ref MyVPC
      Port: 8501
      Protocol: TCP
      TargetType: instance
  MyTargetGroupApi:
    Type: AWS::ElasticLoadBalancingV2::TargetGroup
    Properties:
      VpcId: !Ref MyVPC
      Port: 9000
      Protocol: TCP
      TargetType: instance

  MyNetworkLoadBalancer:
    Type: AWS::ElasticLoadBalancingV2::LoadBalancer
    Properties:
      Scheme: internet-facing
      Subnets:
        - !Ref MyPublicSubnet
      Type: network
        
  MyListener8501:
    Type: AWS::ElasticLoadBalancingV2::Listener
    Properties:
      DefaultActions:
        - Type: forward
          TargetGroupArn: !Ref MyTargetGroup
      LoadBalancerArn: !Ref MyNetworkLoadBalancer
      Port: 8501
      Protocol: TCP

  MyListener9000:
    Type: AWS::ElasticLoadBalancingV2::Listener
    Properties:
      DefaultActions:
        - Type: forward
          TargetGroupArn: !Ref MyTargetGroupApi
      LoadBalancerArn: !Ref MyNetworkLoadBalancer
      Port: 9000
      Protocol: TCP
  QueueLengthScaleUpPolicy: 
    Type: AWS::AutoScaling::ScalingPolicy
    Properties:
      AdjustmentType: ChangeInCapacity
      AutoScalingGroupName: !Ref MyAutoScalingGroup  
      Cooldown: 60
      ScalingAdjustment: 1 # 每次扩展或缩减的实例数量

  QueueLengthScaleDownPolicy:
    Type: AWS::AutoScaling::ScalingPolicy
    Properties:
      AdjustmentType: ChangeInCapacity
      AutoScalingGroupName: !Ref MyAutoScalingGroup
      Cooldown: 60
      ScalingAdjustment: -1

  ScaleUpAlarmNotification:
    Type: AWS::CloudWatch::Alarm
    Properties:
      AlarmDescription: 'Queue length is greater than the threshold'
      MetricName: ApproximateNumberOfMessagesVisible
      Namespace: AWS/SQS
      Statistic: Average    
      Period: 300
      EvaluationPeriods: 1
      Threshold: 10
      AlarmActions:
        - !Ref QueueLengthScaleUpPolicy
      Dimensions:
        - Name: QueueName
          Value: !GetAtt MySQSQueue.QueueName
      ComparisonOperator: GreaterThanThreshold

  ScaleDownAlarmNotification:
    Type: AWS::CloudWatch::Alarm  
    Properties:
      AlarmDescription: 'Queue length is lower than the threshold'
      MetricName: ApproximateNumberOfMessagesVisible
      Namespace: AWS/SQS
      Statistic: Average
      Period: 300
      EvaluationPeriods: 3
      Threshold: 5
      AlarmActions:
        - !Ref QueueLengthScaleDownPolicy
      Dimensions:
        - Name: QueueName
          Value: !GetAtt MySQSQueue.QueueName  
      ComparisonOperator: LessThanThreshold
  CloudFrontDist:
    Type: AWS::CloudFront::Distribution
    DependsOn: MyNetworkLoadBalancer
    Properties:
      DistributionConfig:
        Origins:
          - DomainName: !GetAtt MyNetworkLoadBalancer.DNSName
            Id: streamlitOrigin
            CustomOriginConfig:
              HTTPPort: 8501
              OriginProtocolPolicy: http-only
        Enabled: true
        DefaultCacheBehavior:
          TargetOriginId: streamlitOrigin
          ViewerProtocolPolicy: allow-all
          AllowedMethods:
            - DELETE
            - GET
            - HEAD
            - OPTIONS
            - PATCH
            - POST
            - PUT
          CachePolicyId: 4135ea2d-6df8-44a3-9df3-4b5a84be39ad
          OriginRequestPolicyId: 216adef6-5c7f-47e4-b989-5492eafa07d3
          ForwardedValues:
            QueryString: false
Outputs:
  LoadBalancerDNS:
    Description: The DNS name of the Network Load Balancer
    Value: !GetAtt MyNetworkLoadBalancer.DNSName
    Export:
      Name: MyLoadBalancerDNS
  CloudFrontDomainName:
    Value: !GetAtt CloudFrontDist.DomainName

Mappings:
  RegionMap:
    us-west-1:
      AMI: 'ami-080e1f13689e07408'
    us-west-2:
      AMI: 'ami-08116b9957a259459'

Parameters:
  KeyPairName:
    Type: AWS::EC2::KeyPair::KeyName
    Description: Name of the SSH key pair to use for the EC2 instance   
  InstanceType:
    Type: String
    Description: EC2 instance type
    Default: g5.xlarge
    AllowedValues:
      - g4dn.xlarge
      - g4dn.2xlarge
      - g4dn.4xlarge
      - g4dn.8xlarge
      - g4dn.12xlarge
      - g4dn.16xlarge
      - g5.xlarge
      - g5.2xlarge
      - g5.4xlarge
      - g5.8xlarge
      - g5.12xlarge
      - g5.24xlarge
  HuggingFaceToken:
    Type: String
    Description: Hugging Face token for authentication
  BucketName:
    Type: String
    Default: 'whisper-bucket-${AWS::AccountId}'
    Description: The Bucket Name for trigger the event,pls replace the ${AWS::AccountId} with your accountid number
  LaunchAPI:
    Type: String
    Default: 'True'
    AllowedValues:
      - 'True'
      - 'False'
    Description: If launch the api for check the output files
  LaunchDemo:
    Type: String
    Default: 'True'
    AllowedValues:
      - 'True'
      - 'False'
    Description: If launch the demo ui
  LaunchEvent:
    Type: String
    Default: 'True'
    AllowedValues:
      - 'True'
      - 'False'
    Description: If launch the event handler work flow