functional_tests_hetero_train.yml

name: Common Functional Tests - Heterogeneous Training

on:
  workflow_call:
    inputs:
      platform:
        required: true
        type: string
        description: Platform name (e.g., cuda, default)
      test_matrix:
        required: true
        type: string
        description: JSON array of test configurations
      image:
        required: true
        type: string
      runs_on:
        required: true
        type: string
      container_volumes:
        required: true
        type: string
      container_options:
        required: true
        type: string
      pkg_mgr:
        required: false
        type: string
        description: Package manager (pip, uv, conda). Default uv.
        default: "uv"
      env_name:
        required: false
        type: string
        description: Conda environment name (for conda only)
        default: ""
      env_path:
        required: false
        type: string
        description: Environment path (venv path for uv, conda installation path for conda)
        default: "/opt/venv"

jobs:
  functional_test_hetero_train:
    defaults:
      run:
        shell: bash
    env:
      PROJECT_ROOT: ${{ github.workspace }}/FlagScale
    runs-on: ${{ fromJson(inputs.runs_on) }}
    strategy:
      fail-fast: false
      matrix:
        test_config: ${{ fromJson(inputs.test_matrix) }}
    container:
      image: ${{ inputs.image }}
      ports:
        - 80
      volumes: ${{ fromJson(inputs.container_volumes) }}
      options: ${{ inputs.container_options }}

    steps:
      - name: Checkout code (attempt 1)
        id: checkout_attempt_1
        uses: actions/checkout@v4
        continue-on-error: true
        with:
          repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
          ref: ${{ github.event.pull_request.head.ref || github.ref }}
          path: FlagScale
          ssh-strict: true
          ssh-user: git
          persist-credentials: false
          clean: true
          sparse-checkout-cone-mode: true
          fetch-tags: false
          show-progress: true
          lfs: false
          submodules: false
          set-safe-directory: true

      - name: Checkout code (attempt 2)
        if: steps.checkout_attempt_1.outcome == 'failure'
        id: checkout_attempt_2
        uses: actions/checkout@v4
        continue-on-error: true
        with:
          repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
          ref: ${{ github.event.pull_request.head.ref || github.ref }}
          path: FlagScale
          ssh-strict: true
          ssh-user: git
          persist-credentials: false
          clean: true
          sparse-checkout-cone-mode: true
          fetch-tags: false
          show-progress: true
          lfs: false
          submodules: false
          set-safe-directory: true

      - name: Checkout code (attempt 3)
        if: steps.checkout_attempt_2.outcome == 'failure'
        id: checkout_attempt_3
        uses: actions/checkout@v4
        with:
          repository: ${{ github.event.pull_request.head.repo.full_name || github.repository }}
          ref: ${{ github.event.pull_request.head.ref || github.ref }}
          path: FlagScale
          ssh-strict: true
          ssh-user: git
          persist-credentials: false
          clean: true
          sparse-checkout-cone-mode: true
          fetch-tags: false
          show-progress: true
          lfs: false
          submodules: false
          set-safe-directory: true

      - name: Set safe directory
        run: |
          git config --global --add safe.directory $PROJECT_ROOT

      - name: Check environment info
        run: cd $PROJECT_ROOT && bash ./tests/test_utils/runners/check_env.sh

      - name: Install dependencies for heterogeneous training
        run: |
          set -euo pipefail
          cd $PROJECT_ROOT

          bash ./tests/test_utils/runners/setup_training_test_env.sh \
            --platform '${{ inputs.platform }}' \
            --pkg-mgr '${{ inputs.pkg_mgr }}' \
            --env-name '${{ inputs.env_name }}' \
            --env-path '${{ inputs.env_path }}' \
            --install-cli
        timeout-minutes: 30

      - name: Run functional tests
        id: functional_test
        run: |
          set -euo pipefail
          cd $PROJECT_ROOT

          PLATFORM='${{ inputs.platform }}'
          DEVICE='${{ matrix.test_config.device }}'
          TASK='${{ matrix.test_config.task }}'
          MODEL='${{ matrix.test_config.model }}'
          CASE='${{ matrix.test_config.case }}'
          PKG_MGR='${{ inputs.pkg_mgr }}'
          ENV_NAME='${{ inputs.env_name }}'
          ENV_PATH='${{ inputs.env_path }}'

          echo "Running functional tests for heterogeneous training"
          echo "Platform: $PLATFORM"
          echo "Device: $DEVICE"
          echo "Task: $TASK"
          echo "Model: $MODEL"
          echo "Case: ${CASE:-all}"
          echo "Package Manager: $PKG_MGR"
          echo "Environment Name: $ENV_NAME"
          echo "Environment Path: $ENV_PATH"
          echo "Project root: $PROJECT_ROOT"

          # Source environment utilities
          source ./tools/install/utils/pyenv_utils.sh

          # Activate environment based on package manager
          case "$PKG_MGR" in
            conda)
              if [ -n "$ENV_NAME" ]; then
                activate_conda "$ENV_NAME" "$ENV_PATH" || echo "⚠️  Conda activation failed"
              fi
              ;;
            uv)
              if [ -n "$ENV_PATH" ] && [ -d "$ENV_PATH" ]; then
                activate_uv_env "$ENV_PATH" || echo "⚠️  UV activation failed"
              fi
              ;;
            pip)
              echo "ℹ️  Running tests with pip/system Python"
              ;;
          esac

          # Display Python environment info
          echo "Python location: $(which python)"
          echo "Python version: $(python --version)"

          # Run functional tests using run_tests.sh
          bash "$PROJECT_ROOT/tests/test_utils/runners/run_tests.sh" \
            --platform "$PLATFORM" \
            --device "$DEVICE" \
            --type functional \
            --task "$TASK" \
            --model "$MODEL" \
            --list "$CASE"
          exit_code=$?

          if [ $exit_code -eq 0 ]; then
            echo "✅ Functional tests passed for $PLATFORM/$DEVICE/$TASK/$MODEL/$CASE"
          else
            echo "❌ Functional tests failed for $PLATFORM/$DEVICE/$TASK/$MODEL/$CASE (exit code: $exit_code)"
          fi

          echo "exit_code=$exit_code" >> $GITHUB_OUTPUT
          exit $exit_code
        timeout-minutes: 30

      - name: Upload Functional Test Logs
        if: always() && steps.functional_test.outcome == 'failure'
        uses: actions/upload-artifact@v4
        with:
          name: functional_tests-logs-${{ inputs.platform }}-${{ github.run_id }}-${{ matrix.test_config.task }}-${{ matrix.test_config.model }}
          path: ${{ env.PROJECT_ROOT }}/tests/functional_tests/${{ matrix.test_config.task }}/${{ matrix.test_config.model }}/test_results
          retention-days: 7
          if-no-files-found: warn