build_tests.yaml

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License

name: Build Tests

on:
  workflow_dispatch:
    inputs:
      tpuType:
        description: 'TPU Type'
        required: true
        default: 'v4-8'
        type: choice
        options:
          - v4-8
          - v5litepod-8
  push:
    branches: ["main","develop"]
  pull_request: # By default this runs for types assigned, opened and synchronize.

env:
  # Names must be unique in parallel running tests.
  TPU_TYPE: ${{ inputs.tpuType || 'v4-8' }}
  TPU_CLUSTER_NAME: build-xpk-nodepools
  TPU_FILESTORE_CLUSTER_NAME: build-xpk-filestore
  WORKLOAD_NAME: xpktest-build-${{ github.run_attempt }}
  PATHWAYS_WORKLOAD_NAME: xpkpw-build-${{ github.run_attempt }}
  RUN_ID: "pr-${{ github.event.number }}"
  PROJECT_ID: ${{secrets.PROJECT_NAME}}
  A3_MEGA_TEST_CLUSTER_NAME: "xpk-mega-ctk-int"
  A3_ULTRA_TEST_CLUSTER_NAME: "xpk-ultra-ctk-int"
  GKE_ML_TEST_CLUSTER_NAME: "xpk-gke-ml"
  ZONE: us-central2-a
  REGION: us-central2
  STORAGE_NAME: test-storage
  CLUSTER_ARGUMENTS: "--network=${{secrets.NETWORK_NAME}} --subnetwork=${{secrets.SUBNETWORK_NAME}} --maintenance-window=23:50"
  FS_STORAGE_NAME: ${{secrets.INSTANCE_NAME}}-test-storage
  FS_DELETE_WORKLOAD: "fs-delete-workload"
  FS_READ_WORKLOAD: "fs-read-workload"
  FS_WRITE_WORKLOAD: "fs-write-workload"
  PW_CLUSTER_ARGUMENTS: "--network=${{secrets.NETWORK_NAME}} --subnetwork=${{secrets.SUBNETWORK_NAME}} --maintenance-window=23:50"

jobs:
  run-unit-tests:
    runs-on: [ubuntu-22.04]
    concurrency: # We support one build or nightly test to run at a time currently.
      group: build-test-cluster-group
      cancel-in-progress: false
    steps:
    - uses: actions/checkout@v4
    - uses: actions/setup-python@v5
      with:
        python-version: '3.10'
    - uses: google-github-actions/setup-gcloud@v2
      with:
        version: '>= 363.0.0'
        install_components: 'beta, gke-gcloud-auth-plugin'
    - name: Install dependencies
      run : make install-dev
    - name: Run unit tests
      run: make run-unittests

  run-integration-tests:
    runs-on: [ubuntu-22.04]
    needs: [run-unit-tests]
    concurrency: # We support one build or nightly test to run at a time currently.
      group: build-test-cluster-group
      cancel-in-progress: false
    steps:
    - uses: actions/checkout@v4
    - uses: actions/setup-python@v5
      with:
        python-version: '3.10'
    - uses: 'google-github-actions/auth@v2'
      with:
        credentials_json: '${{ secrets.GCP_SA_KEY }}'
    - uses: google-github-actions/setup-gcloud@v2
      with:
        version: '>= 363.0.0'
        install_components: 'beta,gke-gcloud-auth-plugin, gke-gcloud-auth-plugin'
    - name: Verify gcp setup
      run: gcloud info
    - name: Install dependencies
      run : make install-dev
    - name: "Set auth cidr"
      run: echo "AUTH_CIDR=$(curl api.ipify.org)/32" >> $GITHUB_ENV
    - name: "Set GCLOUD_CFG_PATH"
      run: echo "GCLOUD_CFG_PATH=/home/runner/work/xpk/xpk/" >> $GITHUB_ENV
    - name: "Copy credentials"
      run: cp $GOOGLE_APPLICATION_CREDENTIALS $GCLOUD_CFG_PATH/application_default_credentials.json
    - name: "Set DEPLOYMENT_DIR"
      run: echo "DEPLOYMENT_DIR=$HOME/deployment" >> $GITHUB_ENV
    - name: Create deployment dir
      run: mkdir -p $DEPLOYMENT_DIR
    - name: Run integration tests
      run: make run-integrationtests

  cluster-create-and-delete:
    runs-on: [ubuntu-22.04]
    needs: [run-integration-tests]
    concurrency: # We support one nightly test and one build test for each branch to run at a time currently.
      group: build-test-cluster-group-${{ github.ref }}
      cancel-in-progress: false
    steps:
    - name: Initialize RUN_ID env var
      run: |
        if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then
          RUN_ID="dispatch"
        elif [ "${{ github.ref }}" == "refs/heads/main" ]; then
          RUN_ID="main"
        elif [ "${{ github.ref }}" == "refs/heads/develop" ]; then
          RUN_ID="develop"
        else
          RUN_ID="pr-${{ github.event.number }}"
        fi
        echo "RUN_ID=$RUN_ID" >> $GITHUB_ENV
    - name: Update cluster name with TPU_TYPE and RUN_ID
      run: echo "TPU_CLUSTER_NAME=$TPU_CLUSTER_NAME-$RUN_ID" >> $GITHUB_ENV
    - uses: actions/checkout@v4
    - uses: actions/setup-python@v5
      with:
        python-version: '3.10'
    - name: Initialize RUN_ID env var
      run: |
          if [ "${{ github.event_name }}" == "workflow_dispatch" ]; then
            RUN_ID="dispatch"
          elif [ "${{ github.ref }}" == "refs/heads/main" ]; then
            RUN_ID="main"
          else
            RUN_ID="pr-${{ github.event.number }}"
          fi
          echo "RUN_ID=$RUN_ID" >> $GITHUB_ENV
    - name: Update cluster name with TPU_TYPE and RUN_ID
      run: echo "TPU_CLUSTER_NAME=$TPU_CLUSTER_NAME-$RUN_ID" >> $GITHUB_ENV
    - name: Install expect package
      run: sudo apt-get install expect
    - uses: 'google-github-actions/auth@v2'
      with:
        credentials_json: '${{ secrets.GCP_SA_KEY }}'
    - uses: google-github-actions/setup-gcloud@v2
      with:
        version: '>= 363.0.0'
        install_components: 'beta,gke-gcloud-auth-plugin'
    - name: Generate random seed
      run: |
        RANDOM_SEED=$((RANDOM % 10000))  # Generate a random number between 0 and 9999
        echo "RANDOM_SEED=$RANDOM_SEED" >> $GITHUB_ENV
    - name: Verify gcp setup
      run: gcloud info
    - name: Set Google Cloud CLI properties to a unused zone to verify --zone arg is passed properly in commands.
      run: |
        gcloud config set compute/zone us-east4-a
        gcloud config get compute/zone
    - name: Install xpk dependencies
      run: |
        make install 
        echo $PWD/bin >> "$GITHUB_PATH"
    - name: install kubectl and auth plugin
      run: gcloud components install kubectl && gcloud components install gke-gcloud-auth-plugin
    - name: Check xpk installation
      run: xpk --help
    - name: Create a private Pathways-enabled XPK Cluster with 2x v4-8 nodepools. Larger num-nodes to avoid master resizing.
      run: |
        python3 xpk.py cluster create-pathways --cluster $TPU_CLUSTER_NAME --private --tpu-type=v4-8  --num-slices=2 \
        --zone=us-central2-b --default-pool-cpu-machine-type=n1-standard-16 --default-pool-cpu-num-nodes=16 \
        --reservation='${{ secrets.GCP_TPU_V4_RESERVATION }}' --enable-workload-identity --enable-gcsfuse-csi-driver --enable-gcpfilestore-csi-driver --custom-cluster-arguments="${PW_CLUSTER_ARGUMENTS}"
    - name: List out the nodepools on the cluster
      run: python xpk.py cluster describe --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep -P '^(?=.*NODEPOOL_NAME)(?=.*SLICE)(?=.*TYPE)(?=.*EXPECTED_HEALTHY_NODES)(?=.*ACTUAL_HEALTHY_NODES)(?=.*TOTAL_NODES)'
    - name: Authenticate Docker
      run: gcloud auth configure-docker --quiet
    - name: Create auto-mount Storage instance
      run: |
        python3 xpk.py storage attach $STORAGE_NAME --cluster=$TPU_CLUSTER_NAME --zone=us-central2-b --type=gcsfuse \
        --auto-mount=true \
        --mount-point='/test-mount-point' --readonly=false --manifest='./tests/data/pv-pvc-templates.yaml'
    - name: List and verify existing Storages
      run: python3 xpk.py storage list --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | tee output.txt | grep 'test-storage' || (echo 'No storage found' && cat output.txt && exit 1)
    - name: Create test script to execute in workloads
      run: |
        echo -e \
        '#!/bin/bash \n
        echo "Hello world from a test script!"
        cd ~/../test-mount-point && echo "Hello world from a Github Action CI/CD test script!" > '$RANDOM_SEED'.txt' \
        > test.sh
    - name: Create a test script to execute in jobs.
      run: echo -e '#!/bin/bash \nsleep 3 \necho "Hello world from a test script!"' > job_test.sh
    - name: Run a job
      run: |
        python xpk.py run --cluster $TPU_CLUSTER_NAME --zone=us-central2-b workload.sh | awk '/Starting log streaming for pod xpk-def-app-profile-slurm-[a-zA-Z0-9]+-[0-9]+-[a-zA-Z0-9]+.../,/Job logs streaming finished./ { print }'
    # TODO: fix Delete the job failing step
    # - name: Delete the job
    #   run: |
    #     JOB_NAME=$(python3 xpk.py job ls --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep 'xpk-def-app-profile-slurm-' | head -1 | awk '{print $1}')
    #     kubectl delete job ${JOB_NAME}
    - name: Run a base-docker-image workload
      run: |
        python3 xpk.py workload create --cluster $TPU_CLUSTER_NAME --workload $WORKLOAD_NAME  --command "bash test.sh"  \
        --tpu-type=v4-8 --num-slices=2 --zone=us-central2-b
    - name: Run xpk inspector with the workload created above
      run: python3 xpk.py inspector --cluster $TPU_CLUSTER_NAME --zone=us-central2-b  --workload $WORKLOAD_NAME
    - name: Wait for workload completion and confirm it succeeded
      run: python3 xpk.py workload list --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --wait-for-job-completion $WORKLOAD_NAME --timeout 300
    - name: Run a Pathways workload on Ubuntu base image
      run: |
        python3 xpk.py workload create-pathways --cluster $TPU_CLUSTER_NAME --workload $PATHWAYS_WORKLOAD_NAME \
        --docker-image='marketplace.gcr.io/google/ubuntu2004' --tpu-type=v4-8 --num-slices=2 --zone=us-central2-b \
        --command "echo \"Hello world from a test script! \""
    - name: Verify if the file was created in the GCS bucket
      run: gsutil cp gs://xpk-ci-cd-tests/$RANDOM_SEED.txt .
    - name: Check if the file contains desired content
      run: grep 'Hello world from a Github Action CI/CD test script!' $RANDOM_SEED.txt
    - name: Wait for Pathways workload completion and confirm it succeeded
      run: python3 xpk.py workload list --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --wait-for-job-completion $PATHWAYS_WORKLOAD_NAME --timeout 300
    - name: List out the workloads on the cluster
      run: python3 xpk.py workload list --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
    - name: Run xpk info
      run: python3 xpk.py info --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | tee output.txt | grep -P "^(?=.*QUEUE)(?=.*PENDING_WORKLOADS)(?=.*ADMITTED_WORKLOADS)(?=.*2x$TPU_TYPE:google.com/tpu)(?=.*cpu-rm:cpu)(?=.*cpu-rm:memory)(?=.*cpu-proxy:cpu)(?=.*cpu-proxy:memory)(?=.*cpu-user:cpu)(?=.*cpu-user:memory)" || (echo 'Invalid command output' && cat output.txt && exit 1)
    - name: Delete the workload on the cluster
      run: python3 xpk.py workload delete --workload $WORKLOAD_NAME --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
    - name: Delete the Pathways workload on the cluster
      run: python3 xpk.py workload delete --workload $PATHWAYS_WORKLOAD_NAME --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
    - name: Delete created GCS file
      run: gsutil rm gs://xpk-ci-cd-tests/$RANDOM_SEED.txt
    - name: Delete existing Storage
      run: python3 xpk.py storage delete $STORAGE_NAME  --cluster $TPU_CLUSTER_NAME --zone=us-central2-b
    - name: Create test script to execute in batch
      run: echo -e '#!/bin/bash \n#SBATCH --unknown-flag=value\n echo "Hello world from a test script!"' > batch.sh
    - name: Run a batch job on the cluster
      run: python3 xpk.py batch --cluster $TPU_CLUSTER_NAME --zone=us-central2-b batch.sh --ignore-unknown-flags --array 1-5 --nodes 2 --ntasks 3 --time 60
    - name: List out the jobs on the cluster
      run: python3 xpk.py job ls --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep 'xpk-def-app-profile-slurm-'
    - name: Get created job name
      run: |
        JOB_NAME=$(python3 xpk.py job ls --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep 'xpk-def-app-profile-slurm-' | head -1 | awk '{print $1}')
        echo "JOB_NAME=${JOB_NAME}" >> $GITHUB_ENV
    - name: Check created job
      run: |
        kubectl get job ${JOB_NAME} -o jsonpath='{.metadata.labels}' | grep '"kueue.x-k8s.io/max-exec-time-seconds":"3600"'
        job_spec=$(kubectl get job ${JOB_NAME} -o jsonpath='{.spec}')
        echo "$job_spec" | grep '"completions":2'
        echo "$job_spec" | grep '"parallelism":2'
        echo "$job_spec" | jq '.template.spec.containers | length' | grep 3
    - name: Get job info for the last job created on the cluster
      run: python3 xpk.py job info ${JOB_NAME} | grep -e "Entrypoint environment variables template:" -e "Job name:" -e "Labels:" -e "Mounts:" -e "Pods:" -e "Profile:" -e "Script name:" | wc -l | grep "7"
    - name: Cancel the batch job on the cluster
      run: |
        python3 xpk.py job cancel ${JOB_NAME} --cluster $TPU_CLUSTER_NAME --zone=us-central2-b | grep "job.batch/${JOB_NAME} deleted"
    - name: Create shell and exit it immediately
      run: |
        cat <<'EOF' >> create-shell.exp
        ##!/usr/bin/expect
        spawn python3 ./xpk.py shell
        expect "/ # "
        send "exit\n"
        EOF
        chmod +x ./create-shell.exp
        expect ./create-shell.exp
    # TODO - deflake tests
    # - name: Check if shell exists and is running
    #   run: kubectl get pods | grep xpk-def-app-profile-interactive- | grep Running
    - name: Stop the shell
      run: python3 xpk.py shell stop
    - name: Delete the cluster created
      if: always()
      run: python xpk.py cluster delete --cluster $TPU_CLUSTER_NAME --zone=us-central2-b --force