name: Test Parallel (v2)

on:
  workflow_call:
    inputs:
      resolved_short:
        required: true
        type: string
        description: Short SHA of upstream PyTorch (from _build-new.yml outputs)

defaults:
  run:
    shell: bash

jobs:
  # ── Config: read whitelist YAML → build matrix JSON ──────────────────────
  # Replaces the former _test-collect-new.yml (which downloaded PyTorch source
  # and ran shard_test_files.py).  Now we only read shard counts — the actual
  # file scanning/classification happens inside test-npu.sh at runtime.
  config:
    runs-on: ubuntu-latest
    outputs:
      matrix_json: ${{ steps.build.outputs.matrix_json }}
    steps:
      - uses: actions/checkout@v6
        with:
          ref: ci-test
      - name: Build test matrix from classification config
        id: build
        run: |
          python3 -c "
          import json

          # Load whitelist YAML (minimal parser: no PyYAML needed)
          import sys; sys.path.insert(0, '.github/scripts/v2')
          from shard_test_files import load_categories_config
          config = load_categories_config('.github/config/nightly_v2_test_whitelist.yml')

          include = []
          for cat_name, cat_cfg in config.get('categories', {}).items():
              n = cat_cfg.get('shards', 1)
              runner = cat_cfg.get('runner', 'linux-aarch64-a3-8')
              # Derive NPU count from runner label convention: ...-{n}
              # e.g. linux-aarch64-a3-8 → 8, linux-aarch64-a3-16 → 16
              npu_count = int(runner.rsplit('-', 1)[-1]) if runner else 8
              devices_per_proc = cat_cfg.get('devices_per_proc', 1)
              for s in range(1, n + 1):
                  include.append({
                      'category': cat_name,
                      'shard': s,
                      'num_shards': n,
                      'runner': runner,
                      'npu_count': npu_count,
                      'devices_per_proc': devices_per_proc,
                  })

          matrix = {'include': include}
          print(json.dumps(matrix))
          " > /tmp/matrix.json

          MATRIX=$(cat /tmp/matrix.json)
          echo "matrix_json=${MATRIX}" >> "$GITHUB_OUTPUT"
          echo "Test matrix: ${MATRIX}"

  # ── Test: matrix-driven fan-out ──────────────────────────────────────────
  # Each entry in the matrix becomes a parallel _test-exec-new.yml instance.
  # Shard count and shard index come directly from the whitelist YAML config.
  test:
    needs: config
    strategy:
      matrix: ${{ fromJson(needs.config.outputs.matrix_json) }}
      fail-fast: false
    uses: ./.github/workflows/_test-exec-new.yml
    secrets: inherit
    with:
      resolved_short: ${{ inputs.resolved_short }}
      category: ${{ matrix.category }}
      shard: ${{ matrix.shard }}
      num_shards: ${{ matrix.num_shards }}
      runner: ${{ matrix.runner }}
      npu_count: ${{ matrix.npu_count }}
      devices_per_proc: ${{ matrix.devices_per_proc }}

  # ── Report: aggregate all shard artifacts ─────────────────────────────────
  report:
    needs: [config, test]
    if: always() && needs.config.result == 'success'
    uses: ./.github/workflows/_test-report-new.yml
    secrets: inherit