AWSTemplateFormatVersion: '2010-09-09'
Description: Deploy a publicly accessible EC2 Inf2 instance for LLM

Parameters:
  EnvironmentName:
    Description: An environment name that will be prefixed to resource names
    Type: String
    Default: 'llm-inf2'

#https://stackoverflow.com/questions/76853086/is-there-anyway-that-i-can-check-if-the-listawsec2subnetid-is-empty-in-c
  VpcId:
    Description: ID of the VPC to deploy resources into (keep default value "New" to create a new VPC)
    Type: String
    Default: 'New'

  PublicSubnetId:
    Description: ID of the public subnet to deploy resources into (keep default value "New" to create a new one)
    Type: String
    Default: 'New'

  InstanceType:
    Description: EC2 instance type
    Type: String
    Default: inf2.xlarge
    AllowedValues:
      - inf2.xlarge
      - inf2.8xlarge

  VolumeSize:
    Description: EBS volume size (GB) for the instance
    Type: Number
    MinValue: 80
    Default: 80

  KeyPairName:
    Description: Name of an existing EC2 KeyPair to enable SSH access. Leave blank for no SSH access
    Type: String

Mappings:
  AmiRegionMap:
    us-east-1:
      AMI: 'ami-0e001c9271cf7f3b9' # Ubuntu 22.04 LTS
    us-west-2:
      AMI: 'ami-03c983f9003cb9cd1' # Ubuntu 22.04 LTS

Resources:

  VPC:
    Type: AWS::EC2::VPC
    Condition: CreateNewVPC
    Properties:
      CidrBlock: 10.0.0.0/16
      EnableDnsHostnames: true
      EnableDnsSupport: true
      InstanceTenancy: default
      Tags:
        - Key: Name 
          Value: !Ref EnvironmentName

  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Condition: CreateNewVPC

  InternetGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    Condition: CreateNewVPC 
    Properties:
      InternetGatewayId: !Ref InternetGateway
      VpcId: !Ref VPC

  PublicSubnet:
    Type: AWS::EC2::Subnet
    Condition: CreateNewSubnetId
    Properties:
      VpcId: !If [CreateNewVPC, !Ref VPC, !Ref VpcId]
      AvailabilityZone: !Select [0, !GetAZs '']
      CidrBlock: 10.0.0.0/24
      MapPublicIpOnLaunch: true
      Tags: 
        - Key: Name
          Value: !Sub ${EnvironmentName}-Public-Subnet

  
  ## PUBLIC ROUTING
  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
  
  PublicRoute:
    Type: AWS::EC2::Route
    DependsOn: InternetGatewayAttachment
    Properties:
      RouteTableId: !Ref PublicRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway
 
  PublicSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnet
      RouteTableId: !Ref PublicRouteTable

  SecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allow SSH and custom ports
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp  
          FromPort: 3000
          ToPort: 3000
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp  
          FromPort: 8080
          ToPort: 8080
          CidrIp: 0.0.0.0/0
      VpcId: !If [CreateNewVPC, !Ref VPC, !Ref VpcId]


  EC2Instance:
    Type: AWS::EC2::Instance
    Properties:
      ImageId: !FindInMap [AmiRegionMap, !Ref 'AWS::Region', AMI]
      InstanceType: !Ref 'InstanceType'
      KeyName: 
        Fn::If:
          - HasKeyPairName
          - !Ref 'KeyPairName'
          - !Ref 'AWS::NoValue'
      BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
            VolumeSize: !Ref 'VolumeSize'
            VolumeType: gp3
      NetworkInterfaces:
        - AssociatePublicIpAddress: "true"
          DeviceIndex: "0"
          GroupSet: 
            - !Ref SecurityGroup
          SubnetId: !If [CreateNewSubnetId, !Ref PublicSubnet, !Ref PublicSubnetId]
      UserData:
        Fn::Base64: !Sub |
          #!/bin/bash
          set -x

          export HOME=/home/ubuntu

          # Install Docker

          apt-get update -yqqq
          apt-get install docker -y
          apt-get install -y docker.io
          apt-get install docker-compose -y
          usermod -aG docker ubuntu
          groupadd docker
          apt-get install git -y


          sudo su -c '

          # Configure Linux for Neuron repository updates
          . /etc/os-release
          sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
          deb https://apt.repos.neuron.amazonaws.com jammy main
          EOF
          wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -

          # Update OS packages 
          sudo apt-get update -yqqq


          # Install OS headers 
          # sudo apt-get install linux-headers-$(uname -r) -y

          # install Neuron Driver
          sudo apt-get install aws-neuronx-dkms=2.* -y

          # Install Neuron Runtime 
          sudo apt-get install aws-neuronx-collectives=2.* -y
          sudo apt-get install aws-neuronx-runtime-lib=2.* -y

          # Install Neuron Tools 
          sudo apt-get install aws-neuronx-tools=2.* -y

          # Add PATH
          echo "export PATH=/opt/aws/neuron/bin:$PATH" >> ~/.bashrc


          # install 
          #echo "deb http://security.ubuntu.com/ubuntu jammy-security main" | sudo tee \
          #    /etc/apt/sources.list

          # Update OS packages 
          #sudo apt-get update -y

          #sudo apt-get install libc6 -y
          #sudo DEBIAN_FRONTEND=noninteractive apt-get install libc6 -yqqq


          # Download data
          cd /home/ubuntu/
          wget https://zz-common.s3.amazonaws.com/tmp/tgiui/20240501/work.tar 
          tar -xvf work.tar
          rm work.tar

          cd /home/ubuntu
          # download s5cmd 
          wget https://zz-common.s3.amazonaws.com/tmp/tgiui/s5cmd
          chmod +x ./s5cmd

          # download hfdownloader
          wget https://zz-common.s3.amazonaws.com/tmp/tgiui/hfdownloader
          chmod +x ./hfdownloader


          # download models
          #cd ~/work/
          #mkdir -p ~/work/models/NousResearch/Llama-2-7b-chat-hf
          #./hfdownloader -m cszhzleo/Llama-2-7b-chat-hf-nc2-bs1-token1024-neuron-218 -c 16 --#silentMode > /dev/null
          #cd downloads/cszhzleo_Llama-2-7b-chat-hf-nc2-bs1-token1024-neuron-218
          #mv * ~/work/models/NousResearch/Llama-2-7b-chat-hf/


          cd ~/work/
          mkdir -p ~/work/models/mistralai/Mistral-7B-Instruct-v0.2
          ~/work/hfdownloader -m cszhzleo/Mistral-7B-Instruct-v0.2-nc2-bs1-token1024-neuron-218 -c 16 --silentMode > /dev/null
          cd downloads/cszhzleo_Mistral-7B-Instruct-v0.2-nc2-bs1-token1024-neuron-218
          mv * ~/work/models/mistralai/Mistral-7B-Instruct-v0.2/

          # codellama/CodeLlama-7b-Instruct-hf
          cd ~/work/
          mkdir -p ~/work/models/codellama/CodeLlama-7b-Instruct-hf
          ~/work/hfdownloader -m cszhzleo/CodeLlama-7b-Instruct-hf-nc2-bs1-token1024-neuron-218 -c 16 --silentMode > /dev/null
          cd downloads/cszhzleo_CodeLlama-7b-Instruct-hf-nc2-bs1-token1024-neuron-218
          mv * ~/work/models/codellama/CodeLlama-7b-Instruct-hf/

          cd ~/work/
          mkdir -p ~/work/models/NousResearch/Meta-Llama-3-8B-Instruct
          ~/work/hfdownloader -m cszhzleo/Meta-Llama-3-8B-Instruct-nc2-bs1-token1024-neuron-218 -c 16 --silentMode > /dev/null
          cd downloads/cszhzleo_Meta-Llama-3-8B-Instruct-nc2-bs1-token1024-neuron-218
          mv * ~/work/models/NousResearch/Meta-Llama-3-8B-Instruct/

          cd ~/work/
          rm -rf downloads


          # Download docker image - mongo
          docker pull mongo:latest

          # Download docker image - tgi
          docker pull cszhzleo/tgi-2.18.2:latest
          docker tag cszhzleo/tgi-2.18.2 tgi-2.18.2
          docker rmi cszhzleo/tgi-2.18.2:latest

          # Install Node
          curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.7/install.sh | bash

          # update env

          echo "export NVM_DIR=$HOME/.nvm" >> ~/.bashrc
          [ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"  # This loads nvm
          [ -s "$NVM_DIR/bash_completion" ] && \. "$NVM_DIR/bash_completion" 

          sudo chmod +x $HOME/.nvm/nvm.sh
          . $HOME/.nvm/nvm.sh

          . $HOME/.bashrc
          cd /home/ubuntu
          nvm install 20.12.2


          cd work
          chmod 777 start.sh
          chmod 777 startui.sh
          chmod 664 tgiui.service
          sudo cp tgiui.service /etc/systemd/system/
          sudo systemctl daemon-reload
          sudo systemctl enable tgiui.service
          cd web
          npm install
          cd ..
          ./startui.sh
          ' ubuntu


Conditions:
  CreateNewVPC: !Equals [!Ref VpcId, 'New']
  CreateNewSubnetId: !Equals [!Ref PublicSubnetId, 'New']
  HasKeyPairName: !Not [!Equals [!Ref KeyPairName, '']]

Outputs:

  InstancePublicIp:
    Description: Public IP of the EC2 instance
    Value: !GetAtt EC2Instance.PublicIp
    Export:
      Name: !Sub ${EnvironmentName}-InstancePublicIp

  WebServerPublicEndpoint:
    Description: Public endpoint for the web server
    Value: !Sub "http://${EC2Instance.PublicIp}:3000/"
    Export: 
      Name: !Sub "${EnvironmentName}-WebServerPublicEndpoint"
