e2e-preset-test.yml

name: E2E Preset Test

on:
    workflow_run:
        workflows: ["Build and Push Preset Models"]
        types: 
            - completed
    workflow_dispatch:
        inputs:
            force-run-all:
                type: boolean
                default: false
                description: "Test all models for E2E"

env:
    GO_VERSION: "1.20"
    BRANCH_NAME: ${{ github.head_ref || github.ref_name}} 

permissions:
    id-token: write
    contents: read

jobs:
  determine-models:
    if: github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success'
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.affected_models.outputs.matrix }}
      is_matrix_empty: ${{ steps.check_matrix_empty.outputs.is_empty }}
      full_matrix: ${{ steps.images.outputs.full_matrix }}
    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
            submodules: true
            fetch-depth: 0

      - name: Set FORCE_RUN_ALL Flag
        run: echo "FORCE_RUN_ALL=${{ github.event_name == 'workflow_dispatch' && github.event.inputs.force-run-all == 'true' }}" >> $GITHUB_ENV

      # This script should output a JSON array of model names
      - name: Determine Affected Models
        id: affected_models
        run: |
            PR_BRANCH=${{ env.BRANCH_NAME }} \
            FORCE_RUN_ALL=${{ env.FORCE_RUN_ALL }} \
            python3 .github/workflows/kind-cluster/determine_models.py

      - name: Print Determined Models
        run: |
            echo "Output from determine_models: ${{ steps.affected_models.outputs.matrix }}"
    
      - name: Check if Matrix is Empty
        id: check_matrix_empty
        run: |
            if [ "${{ steps.affected_models.outputs.matrix }}" == "[]" ] || [ -z "${{ steps.affected_models.outputs.matrix }}" ]; then
                echo "is_empty=true" >> $GITHUB_OUTPUT
            else
                echo "is_empty=false" >> $GITHUB_OUTPUT
            fi
        
      - name: Add Config info for Testing
        if: steps.check_matrix_empty.outputs.is_empty == 'false'
        id: images
        run: |         
            # Read the additional configurations from e2e-preset-configs.json
            CONFIGS=$(cat .github/e2e-preset-configs.json | jq -c '.matrix.image')
            
            # Pseudocode for combining matrices
            # COMBINED_MATRIX = []
            # for model in MATRIX:
            #     for config in CONFIGS:
            #         if config['name'] == model['name']:
            #             combined = {**model, **config}
            #             COMBINED_MATRIX.append(combined)
            #             break

            COMBINED_MATRIX=$(echo '${{ steps.affected_models.outputs.matrix }}' | jq --argjson configs "$CONFIGS" -c '
                map(. as $model | $configs[] | select(.name == $model.name) | $model + .)
            ')

            echo "full_matrix=$COMBINED_MATRIX" >> $GITHUB_OUTPUT
      
      - name: Print Combined Matrix
        if: steps.check_matrix_empty.outputs.is_empty == 'false'
        run: |
            echo "Combined Matrix:"
            echo '${{ steps.images.outputs.full_matrix }}'

  e2e-preset-tests:
    needs: determine-models
    if: needs.determine-models.outputs.is_matrix_empty == 'false' && (github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success')
    runs-on: ubuntu-latest
    environment: e2e-test
    strategy:
      fail-fast: false
      matrix:
        # Ex matrix element:
        # {"name":"falcon-40b","type":"text-generation","version":"#",
        # "runtime":"tfs","tag":"0.0.1","node-count":1,
        # "node-vm-size":"Standard_NC96ads_A100_v4", "node-osdisk-size":400}
        model: ${{fromJson(needs.determine-models.outputs.full_matrix)}}
    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
            submodules: true
            fetch-depth: 0

      - name: Set OSS Flag
        run: echo "MODEL_IS_OSS=${{ matrix.model.OSS }}" >> $GITHUB_ENV

      - name: 'Az CLI login'
        uses: azure/login@v1.6.1
        with:
            client-id: ${{ secrets.AZURE_CLIENT_ID }}
            tenant-id: ${{ secrets.AZURE_TENANT_ID }}
            allow-no-subscriptions: true
      
      - name: 'Set subscription'
        run: az account set --subscription ${{secrets.AZURE_SUBSCRIPTION_ID}}

      - name: 'Check if Image exists in Test ACR'
        id: check_test_image
        run: |
            ACR_NAME=${{ secrets.ACR_AMRT_USERNAME }}
            IMAGE_NAME=${{ matrix.model.name }}
            TAG=${{ matrix.model.tag }}

            # Use '|| true' to prevent script from exiting with an error if the repository is not found
            TAGS=$(az acr repository show-tags -n $ACR_NAME --repository $IMAGE_NAME --output tsv || true)

            if [[ -z "$TAGS" ]]; then
                echo "Image $IMAGE_NAME:$TAG or repository not found in $ACR_NAME."
                echo "IMAGE_EXISTS=false" >> $GITHUB_OUTPUT
            else
                if echo "$TAGS" | grep -q "^$TAG$"; then
                    echo "IMAGE_EXISTS=true" >> $GITHUB_OUTPUT
                else
                    echo "IMAGE_EXISTS=false" >> $GITHUB_OUTPUT
                    echo "Image $IMAGE_NAME:$TAG not found in $ACR_NAME."
                fi
            fi
    
      - name: 'Set subscription'
        run: az account set --subscription ${{secrets.PROD_ACR_SUB_ID}}
    
      - name: 'Check if Image exists in Prod ACR'
        id: check_prod_image
        run: |
            ACR_NAME=${{ secrets.PROD_ACR_USERNAME }}
            az acr login --name $ACR_NAME --expose-token

            IMAGE_NAME=unlisted/aks/kaito/kaito-${{ matrix.model.name }}
            TAG=${{ matrix.model.tag }}
        
            # Use '|| true' to prevent script from exiting with an error if the repository is not found
            TAGS=$(az acr repository show-tags -n $ACR_NAME --repository $IMAGE_NAME --output tsv || true)

            if [[ -z "$TAGS" ]]; then
                echo "Image $IMAGE_NAME:$TAG or repository not found in $ACR_NAME."
                echo "IMAGE_EXISTS=false" >> $GITHUB_OUTPUT
            else
                if echo "$TAGS" | grep -q "^$TAG$"; then
                    echo "IMAGE_EXISTS=true" >> $GITHUB_OUTPUT
                else
                    echo "IMAGE_EXISTS=false" >> $GITHUB_OUTPUT
                    echo "Image $IMAGE_NAME:$TAG not found in $ACR_NAME."
                fi
            fi
    
      - name: Check if Image is Test and Prod ACRs
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'true'
        run: |
            echo "Skipping: Image already exists in both Test and Prod ACRs, remember to bump tag"

      - name: Set up kubectl context
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
          az aks get-credentials --resource-group llm-test --name GitRunner
    
      - name: Get Nodepool Name
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        id: get_nodepool_name
        run: |
            NAME_SUFFIX=${{ matrix.model.name }}
            NAME_SUFFIX_WITHOUT_DASHES=${NAME_SUFFIX//-/}  # Removing all '-' symbols
            
            if [ ${#NAME_SUFFIX_WITHOUT_DASHES} -gt 12 ]; then
                TRUNCATED_NAME_SUFFIX=${NAME_SUFFIX_WITHOUT_DASHES: -12}
            else
                TRUNCATED_NAME_SUFFIX=$NAME_SUFFIX_WITHOUT_DASHES
            fi
            echo "Nodepool Name: $TRUNCATED_NAME_SUFFIX"
            echo "NODEPOOL_NAME=$TRUNCATED_NAME_SUFFIX" >> $GITHUB_OUTPUT

      - name: Create Nodepool
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
            NODEPOOL_EXIST=$(az aks nodepool show \
                            --name ${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }} \
                            --cluster-name GitRunner \
                            --resource-group llm-test \
                            --query 'name' -o tsv || echo "")
            echo "NODEPOOL_EXIST: $NODEPOOL_EXIST"
            if [ -z "$NODEPOOL_EXIST" ]; then
                az aks nodepool add \
                    --name ${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }} \
                    --cluster-name GitRunner \
                    --resource-group llm-test \
                    --node-count ${{ matrix.model.node-count }} \
                    --node-vm-size ${{ matrix.model.node-vm-size }} \
                    --node-osdisk-size ${{ matrix.model.node-osdisk-size }} \
                    --labels pool=${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }} \
                    --node-taints sku=gpu:NoSchedule \
                    --aks-custom-headers UseGPUDedicatedVHD=true
            else
                NODEPOOL_STATE=$(az aks nodepool show \
                                --name ${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }} \
                                --cluster-name GitRunner \
                                --resource-group llm-test \
                                --query 'provisioningState' -o tsv)
                echo "NODEPOOL_STATE: $NODEPOOL_STATE"          
                if [ "$NODEPOOL_STATE" != "Succeeded" ]; then
                    echo "Nodepool exists but is not in a Succeeded state. Please check manually."
                    exit 1
                else
                    echo "Nodepool already exists and is in a running state."
                fi
            fi

      - name: Create Service
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: kubectl apply -f presets/test/manifests/${{ matrix.model.name }}/${{ matrix.model.name }}-service.yaml
      
      - name: Retrieve External Service IP
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        id: get_ip
        run: |
            while [[ -z $SERVICE_IP ]]; do 
                SERVICE_IP=$(kubectl get svc ${{ matrix.model.name }} -o=jsonpath='{.status.loadBalancer.ingress[0].ip}')
                sleep 5
            done 
            echo "Service IP is $SERVICE_IP"
            echo "SERVICE_IP=$SERVICE_IP" >> $GITHUB_OUTPUT

      - name: Get Resource Type
        id: resource
        run: |
            RESOURCE_TYPE=$(echo "${{ matrix.model.name }}" | grep -q "llama" && echo "statefulset" || echo "deployment")
            echo "RESOURCE_TYPE=$RESOURCE_TYPE" >> $GITHUB_OUTPUT
      
      - name: Replace IP and Deploy Resource to K8s
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
            sed -i "s/MASTER_ADDR_HERE/${{ steps.get_ip.outputs.SERVICE_IP }}/g" presets/test/manifests/${{ matrix.model.name }}/${{ matrix.model.name }}.yaml
            sed -i "s/TAG_HERE/${{ matrix.model.tag }}/g" presets/test/manifests/${{ matrix.model.name }}/${{ matrix.model.name }}.yaml
            sed -i "s/REPO_HERE/${{ secrets.ACR_AMRT_USERNAME }}/g" presets/test/manifests/${{ matrix.model.name }}/${{ matrix.model.name }}.yaml
            kubectl apply -f presets/test/manifests/${{ matrix.model.name }}/${{ matrix.model.name }}.yaml

      - name: Wait for Resource to be ready
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
            kubectl rollout status ${{steps.resource.outputs.RESOURCE_TYPE}}/${{ matrix.model.name }} --timeout=1800s
        
      - name: Test home endpoint
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
            curl http://${{ steps.get_ip.outputs.SERVICE_IP }}:80/

      - name: Test healthz endpoint
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
            curl http://${{ steps.get_ip.outputs.SERVICE_IP }}:80/healthz
    
      - name: Test inference endpoint
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false'
        run: |
            if [[ "${{ matrix.model.name }}" == *"llama"* && "${{ matrix.model.name }}" == *"-chat"* ]]; then
                echo "Testing inference for ${{ matrix.model.name }}"
                curl -X POST \
                -H "Content-Type: application/json" \
                -d '{
                    "input_data": {
                        "input_string": [
                            [
                                {
                                    "role": "system",
                                    "content": "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe."
                                },
                                {
                                    "role": "user",
                                    "content": "Write a brief birthday message to John"
                                }
                            ]
                        ]
                    }
                }' \
                http://${{ steps.get_ip.outputs.SERVICE_IP }}:80/chat
            elif [[ "${{ matrix.model.name }}" == *"llama"* ]]; then
                echo "Testing inference for ${{ matrix.model.name }}"
                curl -X POST \
                -H "Content-Type: application/json" \
                -d '{
                    "prompts": [
                        "I believe the meaning of life is",
                        "Simply put, the theory of relativity states that ",
                        "A brief message congratulating the team on the launch: Hi everyone, I just ",
                        "Translate English to French: sea otter => loutre de mer, peppermint => menthe poivrée, plush girafe => girafe peluche, cheese =>"
                    ],
                    "parameters": {
                        "max_gen_len": 128
                    }
                }' \
                http://${{ steps.get_ip.outputs.SERVICE_IP }}:80/generate
            else
                echo "Testing inference for ${{ matrix.model.name }}"
                curl -X POST \
                -H "accept: application/json" \
                -H "Content-Type: application/json" \
                -d '{
                    "prompt":"Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
                    "return_full_text": false,
                    "clean_up_tokenization_spaces": false, 
                    "prefix": null,
                    "handle_long_generation": null,
                    "generate_kwargs": {
                            "max_length":200,
                            "min_length":0,
                            "do_sample":true,
                            "early_stopping":false,
                            "num_beams":1,
                            "num_beam_groups":1,
                            "diversity_penalty":0.0,
                            "temperature":1.0,
                            "top_k":10,
                            "top_p":1,
                            "typical_p":1,
                            "repetition_penalty":1,
                            "length_penalty":1,
                            "no_repeat_ngram_size":0,
                            "encoder_no_repeat_ngram_size":0,
                            "bad_words_ids":null,
                            "num_return_sequences":1,
                            "output_scores":false,
                            "return_dict_in_generate":false,
                            "forced_bos_token_id":null,
                            "forced_eos_token_id":null,
                            "remove_invalid_values":null
                        }
                    }' \
                http://${{ steps.get_ip.outputs.SERVICE_IP }}:80/chat                
            fi
      
      - name: Move from Test to Prod ACR
        if: steps.check_test_image.outputs.IMAGE_EXISTS == 'true' && steps.check_prod_image.outputs.IMAGE_EXISTS == 'false' && github.event_name == 'workflow_dispatch' && env.MODEL_IS_OSS == 'true'
        run: |
            # This should only run if:
            # 1. All prior steps have succeeed (Given)
            # 2. Image exists in test ACR repo but not Prod
            # 3. Workflow was triggered manually (workflow_dispatch)
            # 4. Image is OSS (MIT/Apache2.0)

            az account set --subscription ${{secrets.PROD_ACR_SUB_ID}}

            TEST_ACR_NAME=${{ secrets.ACR_AMRT_USERNAME }}
            PROD_ACR_NAME=${{ secrets.PROD_ACR_USERNAME }}
            IMAGE_NAME=${{ matrix.model.name }}
            TAG=${{ matrix.model.tag }}

            # Formulate the source image reference
            SOURCE_IMAGE=$IMAGE_NAME:$TAG
            DEST_IMAGE=unlisted/aks/kaito/kaito-$IMAGE_NAME:$TAG

            # Import the image from Test ACR to Prod ACR
            az acr import \
                --name $PROD_ACR_NAME \
                --source $SOURCE_IMAGE \
                --image $DEST_IMAGE \
                --registry /subscriptions/${{secrets.AZURE_SUBSCRIPTION_ID}}/resourceGroups/${{secrets.TEST_ACR_RG}}/providers/Microsoft.ContainerRegistry/registries/$TEST_ACR_NAME
            
      - name: Cleanup
        if: always()
        run: |
            # Only proceed if RESOURCE_TYPE is set (else resource wasn't created)
            if [ -n "${{ steps.resource.outputs.RESOURCE_TYPE }}" ]; then
                # Use RESOURCE_TYPE from the previous step
                RESOURCE_TYPE=${{ steps.resource.outputs.RESOURCE_TYPE }}
                
                # Check and Delete K8s Resource (Deployment or StatefulSet)
                if kubectl get $RESOURCE_TYPE ${{ matrix.model.name }} > /dev/null 2>&1; then
                    kubectl delete $RESOURCE_TYPE ${{ matrix.model.name }}
                fi
            fi

            # Check and Delete K8s Service if it exists
            if kubectl get svc ${{ matrix.model.name }} > /dev/null 2>&1; then
                kubectl delete svc ${{ matrix.model.name }}
            fi
        
            # Check and Delete AKS Nodepool if it exists            
            if [ -n "${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }}" ]; then
                NODEPOOL_EXIST=$(az aks nodepool show \
                                --name ${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }} \
                                --cluster-name GitRunner \
                                --resource-group llm-test \
                                --query 'name' -o tsv || echo "")

                if [ -n "$NODEPOOL_EXIST" ]; then
                    az aks nodepool delete \
                    --name ${{ steps.get_nodepool_name.outputs.NODEPOOL_NAME }} \
                    --cluster-name GitRunner \
                    --resource-group llm-test
                fi
            fi