cloudformation-template-dev.yaml

AWSTemplateFormatVersion: '2010-09-09'
Description: 'CloudFormation template for EC2 instance for ModelHubStack'

Metadata: 
  AWS::CloudFormation::Interface: 
    StackName: 
      Default: "ModelHubStack"

Parameters:
  InstanceType:
    Type: String
    Default: m5.xlarge
    Description: EC2 instance type

  AMIId:
    Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
    Default: /aws/service/canonical/ubuntu/server/22.04/stable/current/amd64/hvm/ebs-gp2/ami-id
    Description: Ubuntu 22.04 AMI ID

  HuggingFaceHubToken:
    Type: String
    Description: Optional Hugging Face Hub Token
    Default: ""

  WandbApiKey:
    Type: String
    Description: Optional WANDB API Key for view W&B in wandb.ai
    Default: ""

Resources:
  EC2Instance:
    Type: AWS::EC2::Instance
    CreationPolicy:
      ResourceSignal:
        Timeout: PT30M
    Properties:
      Tags:
        - Key: Name
          Value: ModelHubServer
      InstanceType: !Ref InstanceType
      ImageId: !Ref AMIId
      SecurityGroupIds: 
        - !Ref EC2SecurityGroup
      BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
            VolumeSize: 200
            VolumeType: gp3
      IamInstanceProfile: !Ref EC2InstanceProfile
      UserData:
        Fn::Base64:
          !Sub 
            - |
              #!/bin/bash
              # 设置日志文件
              LOG_FILE="/var/log/user-data.log"

              # 函数：记录日志
              log() {
                echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" >> "$LOG_FILE"
              }

              log "Starting UserData script execution"

              # 更新和安装基本软件
              log "Updating apt and installing git"
              sudo apt update && sudo apt install -y git || { log "Failed to update apt or install git"; exit 1; }

              # 检查并安装 AWS CLI
              if ! command -v aws &> /dev/null; then
                log "Installing AWS CLI"
                sudo apt update && sudo apt install -y awscli || { log "Failed to install AWS CLI"; exit 1; }
              else
                log "AWS CLI already installed"
              fi

              # 安装 Node.js
              log "Installing Node.js"
              curl -fsSL https://deb.nodesource.com/setup_18.x | sudo -E bash - && \
              sudo apt install -y nodejs || { log "Failed to install Node.js"; exit 1; }

              log "Installing yarn"
              sudo npm install --global yarn || { log "Failed to install yarn"; exit 1; }

              # 克隆仓库
              log "Cloning repository"
              sudo su ubuntu
              cd /home/ubuntu/
              git clone --recurse-submodule https://github.com/aws-samples/llm_model_hub.git -b dev|| { log "Failed to clone repository"; exit 1; }
              cd /home/ubuntu/llm_model_hub
              yarn install || { log "Failed to run yarn install"; exit 1; }

              # 安装 pm2
              log "Installing pm2"
              sudo yarn global add pm2 || { log "Failed to install pm2"; exit 1; }

              # 等待实例完全启动
              log "Waiting for instance to fully start"
              sleep 30

              # 获取实例元数据
              log "Retrieving instance metadata"
              TOKEN=$(curl -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600")
              EC2_IP=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" -s http://169.254.169.254/latest/meta-data/public-ipv4)
              REGION=$(curl -H "X-aws-ec2-metadata-token: $TOKEN" -s http://169.254.169.254/latest/meta-data/placement/region)

              log "Got IP: $EC2_IP and Region: $REGION"

              # 生成随机密钥
              RANDOM_KEY=$(openssl rand -base64 32 | tr -dc 'a-zA-Z0-9' | fold -w 32 | head -n 1)
              
              # 写入环境变量文件
              log "Writing environment variables"
              echo "REACT_APP_API_ENDPOINT=http://$EC2_IP:8000/v1" > /home/ubuntu/llm_model_hub/.env
              echo "REACT_APP_API_KEY=$RANDOM_KEY" >> /home/ubuntu/llm_model_hub/.env
              echo "REACT_APP_CALCULATOR=https://aws-gpu-memory-caculator.streamlit.app/" >> /home/ubuntu/llm_model_hub/.env

              # 写入后端环境变量
              log "Writing backend environment variables"
              cat << EOF > /home/ubuntu/llm_model_hub/backend/.env
              AK=
              SK=
              role=${SageMakerRoleArn}
              region=$REGION
              db_host=127.0.0.1
              db_name=llm
              db_user=llmdata
              db_password=llmdata
              api_keys=$RANDOM_KEY
              HUGGING_FACE_HUB_TOKEN=${HuggingFaceHubToken}
              WANDB_API_KEY=${WandbApiKey}
              MAX_MODEL_LEN=4096
              EOF

              # 设置权限
              log "Setting proper permissions"
              sudo chown -R ubuntu:ubuntu /home/ubuntu/

              # 生成随机密码并存储
              RANDOM_PASSWORD=$(openssl rand -base64 12 | tr -dc 'a-zA-Z0-9' | fold -w 8 | head -n 1) 
              aws ssm put-parameter --name "/${AWS::StackName}/RandomPassword" --value "$RANDOM_PASSWORD" --type "SecureString" --overwrite --region ${AWS::Region} || { log "Failed to store random password in SSM"; exit 1; }

              # 运行设置脚本
              log "Running setup script"
              cd /home/ubuntu/llm_model_hub/backend
              sudo -u ubuntu bash 01.setup.sh || { log "Failed to run setup script"; exit 1; }
              sleep 30

              # 添加用户到数据库
              log "Adding user to database"
              source ../miniconda3/bin/activate py311
              conda activate py311
              python3 users/add_user.py demo_user $RANDOM_PASSWORD default || { log "Failed to add user to database"; exit 1; }

              # 构建镜像
              log "Building and pushing image"
              cd /home/ubuntu/llm_model_hub/backend/byoc
              sudo -u ubuntu bash build_and_push.sh || { log "Failed to build and push image"; exit 1; }
              sleep 5

              # 上传 dummy tar.gz
              log "Uploading dummy tar.gz"
              sudo -u ubuntu ../../miniconda3/envs/py311/bin/python startup.py || { log "Failed to upload dummy tar.gz"; exit 1; }

              # 启动后端
              log "Starting backend"
              cd /home/ubuntu/llm_model_hub/backend/
              sudo -u ubuntu bash 02.start_backend.sh || { log "Failed to start backend"; exit 1; }
              sleep 15

              # 启动前端
              log "Starting frontend"
              cd /home/ubuntu/llm_model_hub/
              sudo -u ubuntu pm2 start pm2run.config.js || { log "Failed to start frontend"; exit 1; }

              # 安装 cfn-signal
              log "Installing cfn-signal"
              sudo pip install https://s3.amazonaws.com/cloudformation-examples/aws-cfn-bootstrap-py3-latest.tar.gz || { log "Failed to install cfn-signal"; exit 1; }

              # 发送成功信号
              log "Sending success signal to CloudFormation"
              /usr/local/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource EC2Instance --region ${AWS::Region} || { log "Failed to send signal to CloudFormation"; exit 1; }

              echo "RandomPassword=$RANDOM_PASSWORD" >> /etc/environment

              log "UserData script execution completed successfully"
            - SageMakerRoleArn: !GetAtt SageMakerExecutionRole.Arn
              HuggingFaceHubToken: !Ref HuggingFaceHubToken
              WandbApiKey: !Ref WandbApiKey

  EC2SecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allow SSH, port 8000 and 3000
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 8000
          ToPort: 8000
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 3000
          ToPort: 3000
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: 18.237.140.160/29

  EC2Role:
    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/CloudWatchLogsFullAccess
        - arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
      Policies:
        - PolicyName: CloudFormationSignalPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: cloudformation:SignalResource
                Resource: !Sub 'arn:aws:cloudformation:${AWS::Region}:${AWS::AccountId}:stack/${AWS::StackName}/*'
        - PolicyName: SSMParameterAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ssm:PutParameter
                Resource: !Sub 'arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/${AWS::StackName}/*'


  EC2InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
        - !Ref EC2Role

  SageMakerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
      Policies:
        - PolicyName: S3AccessPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - s3:GetObject
                  - s3:PutObject
                  - s3:DeleteObject
                  - s3:ListBucket
                  - s3:CreateBucket
                Resource:
                  - arn:aws:s3:::*
        - PolicyName: SSMSessionManagerPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ssmmessages:CreateControlChannel
                  - ssmmessages:CreateDataChannel
                  - ssmmessages:OpenControlChannel
                  - ssmmessages:OpenDataChannel
                Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:*'            

  SSMConsoleAccessRole:
    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/AmazonSSMFullAccess
      Policies:
        - PolicyName: SSMConsoleAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ssm:StartSession
                  - ssm:DescribeSessions
                  - ssm:TerminateSession
                  - ssm:ResumeSession
                Resource: '*'

                
Outputs:
  InstanceId:
    Description: ID of the EC2 instance
    Value: !Ref EC2Instance
  PublicIP:
    Description: Public IP of the EC2 instance
    Value: !GetAtt EC2Instance.PublicIp
  SageMakerRoleArn:
    Description: ARN of the SageMaker Execution Role
    Value: !GetAtt SageMakerExecutionRole.Arn
  RandomPasswordParameter:
    Description: AWS Systems Manager Parameter name for the random password
    Value: !Sub '/${AWS::StackName}/RandomPassword'