# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

name: Shared Local GPU and CPU Test

on:
  workflow_call:
    inputs:
      test_suite_name:
        description: 'Test suite name (vllm, sglang, trtllm)'
        required: true
        type: string
      test_type:
        description: 'Test type (e.g. single-gpu, multi-gpu)'
        required: true
        type: string
      amd_runner:
        description: 'Runner to execute tests on (amd64 only)'
        required: true
        type: string
      target_tag_plain:
        description: 'Plain runtime image tag prefix from the build workflow'
        required: true
        type: string
      cuda_version:
        description: 'CUDA versions to test as a JSON array'
        required: true
        type: string
      platform:
        description: 'Target platforms to test as a JSON array'
        required: true
        type: string
      run_sanity_check:
        description: 'Whether to run sanity check on the runtime image before executing tests'
        required: false
        type: boolean
        default: true
      run_cpu_only_tests:
        description: 'Whether to run CPU-only tests'
        required: false
        type: boolean
        default: false
      cpu_only_test_markers:
        description: 'CPU-only pytest markers'
        required: false
        type: string
      cpu_only_test_timeout_minutes:
        description: 'Timeout in minutes for CPU tests'
        required: false
        type: number
        default: 10
      run_gpu_tests:
        description: 'Whether to run GPU tests'
        required: false
        type: boolean
        default: true
      gpu_test_markers:
        description: 'GPU pytest markers'
        required: false
        type: string
      gpu_test_timeout_minutes:
        description: 'Timeout in minutes for GPU tests'
        required: false
        type: number
        default: 30
    secrets:
      AWS_DEFAULT_REGION:
        required: true
      AWS_ACCOUNT_ID:
        required: true
      AZURE_ACR_HOSTNAME:
        required: true
      AZURE_ACR_USER:
        required: true
      AZURE_ACR_PASSWORD:
        required: true
      HF_TOKEN:
        required: false

jobs:
  test:
    strategy:
      fail-fast: false
      matrix:
        platform: ${{ fromJson(inputs.platform) }}
        cuda_version: ${{ fromJson(inputs.cuda_version) }}
    name: ${{ inputs.test_type }} ${{ matrix.cuda_version == '' && 'cpu' || format('cuda{0}', matrix.cuda_version) }}, ${{ matrix.platform }}
    runs-on: ${{ matrix.platform == 'amd64' && inputs.amd_runner || 'prod-tester-arm-v1' }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
      - name: Calculate target tag
        id: calculate-target-tag
        shell: bash
        env:
          ECR_REPOSITORY: ${{ secrets.AWS_ACCOUNT_ID }}.dkr.ecr.${{ secrets.AWS_DEFAULT_REGION }}.amazonaws.com/ai-dynamo/dynamo
        run: |
          CUDA_VERSION="${{ matrix.cuda_version }}"
          CUDA_MAJOR=${CUDA_VERSION%%.*}
          if [[ "${{ inputs.target_tag_plain }}" == *"planner"* ]]; then
            IMAGE_TAG=${{ github.sha }}-${{ inputs.target_tag_plain }}
          else
            IMAGE_TAG=${{ github.sha }}-${{ inputs.target_tag_plain }}-cuda${CUDA_MAJOR}
          fi

          RUNTIME_IMAGE=${ECR_REPOSITORY}:${IMAGE_TAG}
          TEST_IMAGE=${ECR_REPOSITORY}:${IMAGE_TAG}-test
          echo "runtime_image=${RUNTIME_IMAGE}" >> $GITHUB_OUTPUT
          echo "test_image=${TEST_IMAGE}" >> $GITHUB_OUTPUT
      - name: Docker Login
        uses: ./.github/actions/docker-login
        with:
          aws_default_region: ${{ secrets.AWS_DEFAULT_REGION }}
          aws_account_id: ${{ secrets.AWS_ACCOUNT_ID }}
          azure_acr_hostname: ${{ secrets.AZURE_ACR_HOSTNAME }}
          azure_acr_user: ${{ secrets.AZURE_ACR_USER }}
          azure_acr_password: ${{ secrets.AZURE_ACR_PASSWORD }}
      - name: Pull relevant images
        shell: bash
        run: |
          source ./.github/scripts/retry_docker.sh
          start_time=$(date +%s)
          retry_pull ${{ steps.calculate-target-tag.outputs.runtime_image }}
          retry_pull ${{ steps.calculate-target-tag.outputs.test_image }}
          retry_pull quay.io/minio/minio
          end_time=$(date +%s)
          duration=$((end_time - start_time))
          echo "⏱️ Image pull duration: ${duration}s"
      - name: Run Sanity Check on Runtime Image
        if: ${{ inputs.run_sanity_check }}
        shell: bash
        run: |
          echo "Running sanity check on image: ${{ steps.calculate-target-tag.outputs.runtime_image }}"

          export WORKSPACE=/workspace

          set +e
          docker run --rm "${{ steps.calculate-target-tag.outputs.runtime_image }}" python ${WORKSPACE}/deploy/sanity_check.py --runtime-check --no-gpu-check
          SANITY_CHECK_EXIT_CODE=$?
          set -e
          if [ ${SANITY_CHECK_EXIT_CODE} -ne 0 ]; then
            echo "ERROR: Sanity check failed - ai-dynamo packages not properly installed"
            exit ${SANITY_CHECK_EXIT_CODE}
          else
            echo "✅ Sanity check passed"
          fi
      - name: Run CPU-only tests (parallelized)
        if: ${{ inputs.run_cpu_only_tests }}
        timeout-minutes: ${{ inputs.cpu_only_test_timeout_minutes }}
        uses: ./.github/actions/pytest
        with:
          image_tag: ${{ steps.calculate-target-tag.outputs.test_image }}
          pytest_marks: ${{ inputs.cpu_only_test_markers }}
          test_suite_name: ${{ inputs.test_suite_name }}
          test_type: "pre_merge_cpu"
          platform_arch: ${{ matrix.platform }}
          hf_token: ${{ secrets.HF_TOKEN }}
          parallel_mode: 'auto'
          dind_as_sidecar: 'true'
      - name: Run GPU tests (sequential)
        timeout-minutes: ${{ inputs.gpu_test_timeout_minutes }}
        if: ${{ matrix.platform == 'amd64' && inputs.run_gpu_tests }}
        uses: ./.github/actions/pytest
        with:
          image_tag: ${{ steps.calculate-target-tag.outputs.test_image }}
          pytest_marks: ${{ inputs.gpu_test_markers }}
          test_suite_name: ${{ inputs.test_suite_name }}
          test_type: "pre_merge_gpu"
          platform_arch: ${{ matrix.platform }}
          hf_token: ${{ secrets.HF_TOKEN }}
          parallel_mode: 'none'
          dind_as_sidecar: 'true'