unit_tests_common.yml

name: Common Unit Tests

on:
  workflow_call:
    inputs:
      platform:
        required: true
        type: string
      device:
        required: true
        type: string
      image:
        required: true
        type: string
      runs_on:
        required: true
        type: string
      container_volumes:
        required: true
        type: string
      container_options:
        required: true
        type: string
      # Platform-specific environment setup script path (from platform config)
      setup_script:
        required: false
        type: string
        default: ''
      # Platform-specific build environment variables (JSON object from config)
      build_env:
        required: false
        type: string
        default: '{}'

jobs:
  unit_test:
    defaults:
      run:
        shell: bash
    runs-on: ${{ fromJson(inputs.runs_on) }}
    strategy:
      fail-fast: false
      matrix:
        test_group:
          - name: pytorch_debug
            path: "qa/L0_pytorch_debug_unittest/test.sh"
            test_type: "debug"
          - name: pytorch_unittest
            path: "qa/L0_pytorch_unittest/test.sh"
            test_type: "unittest"
          - name: pytorch_distributed_unittest
            path: "qa/L1_pytorch_distributed_unittest/test.sh"
            test_type: "unittest"
          - name: pytorch_onnx_unittest
            path: "qa/L1_pytorch_onnx_unittest/test.sh"
            test_type: "unittest"
    name: unit-${{ inputs.device }}-${{ matrix.test_group.name }}
    container:
      image: ${{ inputs.image }}
      volumes: ${{ fromJson(inputs.container_volumes) }}
      options: --pull never ${{ inputs.container_options }}

    steps:
      # Cuda requires git safe.directory configuration and 3 checkout attempts to handle submodule-heavy repos
      - name: Configure Git Safe Directory on Cuda
        if: inputs.platform == 'cuda'
        run: /usr/bin/git config --global safe.directory '*'

      - name: Checkout Source Code on Cuda (attempt 1)
        id: checkout1
        if: inputs.platform == 'cuda'
        uses: actions/checkout@v4
        continue-on-error: true
        with:
          fetch-depth: 0
          submodules: recursive
          set-safe-directory: true

      - name: Checkout Source Code on Cuda (attempt 2)
        id: checkout2
        if: inputs.platform == 'cuda' && steps.checkout1.outcome == 'failure'
        uses: actions/checkout@v4
        continue-on-error: true
        with:
          fetch-depth: 0
          submodules: recursive
          set-safe-directory: true

      - name: Checkout Source Code on Cuda (attempt 3)
        id: checkout3
        if: inputs.platform == 'cuda' && steps.checkout2.outcome == 'failure'
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
          submodules: recursive
          set-safe-directory: true

      # Metax requires to clean vscode-remote-container
      - name: Configure Clean Git Env on Metax
        if: inputs.platform == 'metax'
        run: |
          git config --global --unset-all credential.helper 2>/dev/null || true
          git config --system --unset-all credential.helper 2>/dev/null || true

      # Metax no need submodules
      - name: Checkout Source Code on Metax
        if: inputs.platform == 'metax'
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
      
      - name: Environment Setup
        if: inputs.setup_script != ''
        run: |
          bash $GITHUB_WORKSPACE/${{ inputs.setup_script }}
      
      - name: Execute Tests
        working-directory: ${{ github.workspace }}
        run: |
          set -euo pipefail

          # Load platform-specific environment variables
          while IFS='=' read -r key value; do
            [ -n "$key" ] && export "$key=$value"
          done < <(echo '${{ inputs.build_env }}' | python3 -c "
          import json, sys
          env = json.load(sys.stdin)
          for k, v in env.items():
              print(f'{k}={v}')
          ")
          
          # Activate conda environment
          if ${{inputs.platform == 'metax'}}; then
            source /opt/conda/etc/profile.d/conda.sh
            conda activate base
          else
            source /opt/miniconda3/etc/profile.d/conda.sh
            conda activate flagscale-train
          fi
          echo "PATH=$PATH" >> $GITHUB_ENV

          export TE_PATH=$GITHUB_WORKSPACE
          export TE_LIB_PATH=$(python3 -c "import site; print(site.getsitepackages()[0])")
          export PYTHONPATH=$GITHUB_WORKSPACE:${PYTHONPATH:-}
          export PATH=${CUDA_HOME:-/usr/local/cuda}/bin:$PATH
          export LD_LIBRARY_PATH=${CUDA_HOME:-/usr/local/cuda}/lib:${LD_LIBRARY_PATH:-}
          
          # check envs before running tests
          echo "TE_PATH=$TE_PATH"
          echo "TE_LIB_PATH=$TE_LIB_PATH"
          echo "PYTHONPATH=$PYTHONPATH"
          echo "PATH=$PATH"
          echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"
          
          # Ensure log directory exists regardless of volume mount state
          mkdir -p /logs

          # Coverage setup: install once + configure collection via PYTEST_ADDOPTS
          COVERAGE_ENABLED=false
          if pip3 install coverage pytest-cov --quiet 2>/dev/null; then
            export PYTEST_ADDOPTS="--cov=transformer_engine/pytorch --cov=transformer_engine/debug --cov=transformer_engine/plugin --cov-append --cov-report="
            COVERAGE_ENABLED=true
          else
            echo "WARNING: Failed to install coverage/pytest-cov, coverage collection disabled"
          fi

          if [[ "${{ matrix.test_group.name }}" != *"debug"* ]]; then
            # Fail fast on backend/API mismatch before running the full test group.
            # Skip for debug group (does not use FP8/optimizer symbols).
            python3 -c "import sys, importlib; import transformer_engine.common as _te_common; tex = importlib.import_module('transformer_engine_torch'); required=['multi_tensor_scale','multi_tensor_compute_scale_and_scale_inv']; missing=[n for n in required if not hasattr(tex, n)]; print('[TE check] module:', tex); print('[TE check] file:', getattr(tex, '__file__', 'N/A')); print('[TE check] missing:', ', '.join(missing) if missing else 'none'); sys.exit(1 if missing else 0)"
          fi

          bash ${{ matrix.test_group.path }}
          exit_code=$?

          # Combine coverage fragments and generate JSON report
          if [ "$COVERAGE_ENABLED" = "true" ]; then
            python3 -m coverage combine --keep 2>/dev/null || true
            python3 -m coverage json \
              -o "coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.json" \
              --include="transformer_engine/pytorch/*,transformer_engine/debug/*,transformer_engine/plugin/*" \
              --omit="*/setup.py,*/transformer_engine/plugin/core/_build_config.py" \
              -i 2>/dev/null \
              || echo "WARNING: No coverage data found"

            if [ -f .coverage ]; then
              cp .coverage "coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.coverage"
            fi
          fi
          exit $exit_code
        timeout-minutes: 60

      - name: Upload Coverage Artifact
        uses: actions/upload-artifact@v4
        continue-on-error: true
        with:
          name: coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}
          path: |
            coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.json
            coverage-${{ inputs.platform }}-${{ inputs.device }}-${{ matrix.test_group.name }}.coverage

  aggregate_coverage:
    name: unit-${{ inputs.device }}-pytorch-unittest-coverage
    needs: unit_test
    if: always()
    defaults:
      run:
        shell: bash
    runs-on: ${{ fromJson(inputs.runs_on) }}
    container:
      image: ${{ inputs.image }}
      volumes: ${{ fromJson(inputs.container_volumes) }}
      options: --pull never ${{ inputs.container_options }}

    steps:
      - name: Checkout Source Code
        uses: actions/checkout@v4
        with:
          fetch-depth: 0
          set-safe-directory: true

      - name: Download Coverage Artifacts
        uses: actions/download-artifact@v4
        continue-on-error: true
        with:
          pattern: coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch_*
          path: coverage-artifacts
          merge-multiple: true

      - name: Aggregate PyTorch Coverage
        id: aggregate
        run: |
          set -euo pipefail

          if ${{inputs.platform == 'metax'}}; then
            source /opt/conda/etc/profile.d/conda.sh
            conda activate base
          else
            source /opt/miniconda3/etc/profile.d/conda.sh
            conda activate flagscale-train
          fi

          python3 -m pip install coverage --quiet

          mkdir -p coverage-raw
          coverage_count=0
          for coverage_file in coverage-artifacts/*.coverage; do
            [ -f "$coverage_file" ] || continue
            coverage_count=$((coverage_count + 1))
            cp "$coverage_file" "coverage-raw/.coverage.$coverage_count"
          done

          if [ "$coverage_count" -eq 0 ]; then
            echo "WARNING: No raw coverage files found for aggregation"
            echo "coverage_aggregated=false" >> "$GITHUB_OUTPUT"
            exit 0
          fi

          python3 -m coverage combine coverage-raw
          python3 -m coverage json \
            -o "coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest.json" \
            --include="transformer_engine/pytorch/*,transformer_engine/debug/*,transformer_engine/plugin/*" \
            --omit="*/setup.py,*/transformer_engine/plugin/core/_build_config.py" \
            -i

          echo "coverage_aggregated=true" >> "$GITHUB_OUTPUT"

      - name: Upload Aggregated Coverage Artifact
        if: steps.aggregate.outputs.coverage_aggregated == 'true'
        uses: actions/upload-artifact@v4
        continue-on-error: true
        with:
          name: coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest
          path: coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest.json

      - name: Upload Coverage Report to FlagCICD
        if: steps.aggregate.outputs.coverage_aggregated == 'true'
        uses: flagos-ai/FlagOps/actions/post-pytest-report@v2
        continue-on-error: true
        env:
          NO_PROXY: "flagcicd-inner.flagos.net"
        with:
          backend_url: 'http://flagcicd-inner.flagos.net:8000/metrics/'
          user_id: '000000000000000000'
          report_path: 'coverage-${{ inputs.platform }}-${{ inputs.device }}-pytorch-unittest.json'
          fail_on_error: 'false'