name: 'Nightly Test (NPU) Multi Node Template'

on:
  workflow_call:
    inputs:
      runner:
        required: false
        type: string
        default: linux-amd64-cpu-8
      test_type:
        required: false
        type: string
        default: perf
        description: perf or accuracy
      test_config_name:
        required: true
        type: string
        description: test config name
      prefill_size:
        required: false
        type: number
        default: 0
        description: number of prefill nodes for pd-separation
      decode_size:
        required: false
        type: number
        default: 0
        description: number of decode nodes for pd-separation
      router_size:
        required: false
        type: number
        default: 0
        description: number of router nodes for pd-separation
      node_size:
        required: false
        type: number
        default: 0
        description: number of nodes for pd-mix
      test_case:
        required: true
        type: string
        description: path of test case file
      image:
        required: false
        type: string
        description: image for pods
        default: "swr.cn-southwest-2.myhuaweicloud.com/base_image/dockerhub/lmsysorg/sglang:main-cann8.5.0-a3"
      install_sglang_from_source:
        required: false
        type: boolean
        default: true
        description: use sglang from source code or from docker image
      prefill_decode_deployment:
        required: true
        type: string
        description: "The deployment of prefill and decode nodes. ['separation', 'mix']"
      transformers_version:
        required: false
        type: string
        default: ""
        description: "The transformers version number for running sglang. Use default version in image if keep empty."

concurrency:
  group: ascend-nightly-multi-node-${{ github.workflow_ref }}-${{ github.ref }}-${{ inputs.test_config_name }}
  cancel-in-progress: true

jobs:
  e2e:
    name: ${{ inputs.test_config_name }}
    runs-on: ${{ inputs.runner }}
    container:
      image: swr.ap-southeast-1.myhuaweicloud.com/base_image/ascend-ci/sglang:main-x86
      env:
        KUBECONFIG: /root/.cache/.cache/kb.yaml
        KUBECTL: /root/.cache/.cache/kubectl
        NAMESPACE: sgl-project
        KUBE_JOB_NAME: ascend-sglang-${{ inputs.test_type }}-test
        SGLANG_IS_IN_CI: true
        ASCEND_E2E_TEST_CONFIG_PATH: python/sglang/test/ascend/e2e
        SGLANG_USE_MODELSCOPE: true
        HF_ENDPOINT: https://hf-mirror.com
        TRANSFORMERS_VERBOSITY: "error"
        GDN_ATTN_BACKEND_TRITON: 1
    steps:
      - name: Install dependencies
        run: |
          pip3 install -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple jinja2
          pip3 install -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple psutil
          # Install kubernetes
          pip3 install -i https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple kubernetes
          cp $KUBECTL /usr/local/sbin/

      - name: Checkout code
        uses: actions/checkout@v4

      - name: Prepare code for testing
        run: |
          # copy source code to shared-disk
          current_path=$(pwd)
          target_path=/root/.cache/tests/sglang
          rm -rf ${target_path}
          mkdir -p ${target_path}
          cp -r ${current_path}/* ${target_path}/

      - name: Clear resources
        run: |
          cd ${ASCEND_E2E_TEST_CONFIG_PATH}
          kubectl delete -f ./k8s_multi_pd_*.yaml --ignore-not-found=true || true

          pod_name_prefix="${KUBE_JOB_NAME}-"
          echo "kube name space: $NAMESPACE, pod name prefix: ${pod_name_prefix}"
          while true; do
            if kubectl get po -A -n $NAMESPACE | grep -q "${pod_name_prefix}"; then
              echo "Found exist sglang job, sleeping for 30 seconds..."
              sleep 30
              kubectl get pods | grep "${pod_name_prefix}" | awk '{print $1}' | xargs kubectl delete pod -n $NAMESPACE || true
            else
              echo "No sglang job exist, start test case..."
              break
            fi
          done

      - name: Run test
        timeout-minutes: 300
        run: |
          # sglang_source_relative_path is shared-disk path
          sglang_source_relative_path=tests/sglang
          sglang_source_path=/root/.cache/${sglang_source_relative_path}
          echo "Source code path: ${sglang_source_path}"

          test_case="${{ inputs.test_case }}"
          if [ ! -f "${sglang_source_path}/${test_case}" ]; then
            echo "The testcase does not exit: ${sglang_source_path}/${test_case}"
            exit 1
          fi

          # prepare for output data
          current_date=$(date +%Y%m%d)
          test_data_output_path=/root/.cache/tests/output/${{ inputs.test_type }}/${current_date}
          mkdir -p ${test_data_output_path}
          tc_name=${test_case##*/}
          tc_name=${tc_name%.*}
          metrics_data_path=${test_data_output_path}/${tc_name}
          mkdir -p ${metrics_data_path}
          echo "Metrics file path: ${metrics_data_path}"

          prefill_decode_deployment="${{ inputs.prefill_decode_deployment }}"
          kube_job_type=""
          if [ "$prefill_decode_deployment" = "separation" ];then
            kube_job_type=multi-pd-separation
          elif [ "$prefill_decode_deployment" = "mix" ];then
            kube_job_type=multi-pd-mix
          else
            echo "Unsupported deployment type: ${prefill_decode_deployment}"
            exit 1
          fi

          transformers_version="${{ inputs.transformers_version }}"
          cd ${ASCEND_E2E_TEST_CONFIG_PATH}
          CMD="python3 -u run_npu_e2e_test.py \
            --env ci \
            --image ${{ inputs.image }} \
            --sglang-source-relative-path ${sglang_source_relative_path} \
            --metrics-data-file ${metrics_data_path} \
            --test-case ${test_case} \
            --kube-name-space ${NAMESPACE} \
            --kube-job-type ${kube_job_type} \
            --kube-job-name-prefix ${KUBE_JOB_NAME}"

          if [ "$prefill_decode_deployment" = "separation" ];then
            CMD="${CMD} \
              --prefill-size ${{ inputs.prefill_size }} \
              --decode-size ${{ inputs.decode_size }} \
              --router-size ${{ inputs.router_size }}"
          elif [ "$prefill_decode_deployment" = "mix" ];then
            CMD="${CMD} --node-size ${{ inputs.node_size }}"
          else
            echo "Unsupported deployment type: ${prefill_decode_deployment}"
            exit 1
          fi

          if [ "$SGLANG_IS_IN_CI" = "true" ] || [ "$SGLANG_IS_IN_CI" = "True" ];then
            echo "Run test in ci."
            CMD="${CMD} --sglang-is-in-ci"
          fi

          install_sglang_from_source="${{ inputs.install_sglang_from_source }}"
          if [ "$install_sglang_from_source" = "true" ] || [ "$install_sglang_from_source" = "True" ];then
            echo "Install sglang from source."
            CMD="${CMD} --install-sglang-from-source"
            commit_id=${{ github.sha }}
            echo "commit id: ${commit_id}" > ${test_data_output_path}/commit_id
          else
            echo "Use sglang from image: ${{ inputs.image }}"
          fi

          if [ "$transformers_version" != "" ];then
            CMD="${CMD} --transformers-version ${transformers_version}"
          else
            echo "Use default transformers version in image."
          fi

          echo "Run command: ${CMD}"
          eval "${CMD}"

      - name: Upload metrics
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: metrics-${{ inputs.test_config_name }}
          path: /tmp/metrics.json
          retention-days: 7

      - name: Post process
        if: always()
        run: |
          cd ${ASCEND_E2E_TEST_CONFIG_PATH}
          kubectl get pods -n $NAMESPACE | grep $KUBE_JOB_NAME
          kubectl delete -f ./k8s_multi_pd_*.yaml --ignore-not-found=true || true

          pod_name_prefix="${KUBE_JOB_NAME}-"
          echo "kube name space: $NAMESPACE, pod name prefix: ${pod_name_prefix}"
          while true; do
            if kubectl get po -A -n $NAMESPACE | grep -q "${pod_name_prefix}"; then
              echo "Found exist sglang job, sleeping for 30 seconds..."
              sleep 30
              kubectl get pods | grep "${pod_name_prefix}" | awk '{print $1}' | xargs kubectl delete pod -n $NAMESPACE || true
            else
              echo "No sglang job exist, start test case..."
              break
            fi
          done