name: PR Test NVIDIA

# Split vendor CI into separate workflows so either side can be rerun without
# waiting for the other side's jobs to finish.
on:
  push:
    branches: [main]
    paths:
      - "python/pyproject.toml"
      - "python/tokenspeed/**"
      - "tokenspeed-kernel/**"
      - "tokenspeed-kernel-amd/**"
      - "tokenspeed-mla/**"
      - "tokenspeed-scheduler/**"
      - "test/**"
      - ".github/workflows/pr-test-amd.yml"
      - ".github/workflows/pr-test-nvidia.yml"
      - ".github/workflows/pr-test-nvidia-arm.yml"
  pull_request:
    types: [opened, synchronize, reopened, ready_for_review, closed]
    branches: [main]
    paths:
      - "python/pyproject.toml"
      - "python/tokenspeed/**"
      - "tokenspeed-kernel/**"
      - "tokenspeed-kernel-amd/**"
      - "tokenspeed-mla/**"
      - "tokenspeed-scheduler/**"
      - "test/**"
      - ".github/workflows/pr-test-amd.yml"
      - ".github/workflows/pr-test-nvidia.yml"
      - ".github/workflows/pr-test-nvidia-arm.yml"
  workflow_dispatch:
    inputs:
      trigger:
        description: "CI task trigger to run"
        required: false
        default: manual
        type: choice
        options:
          - manual
          - debug

env:
  PIP_BREAK_SYSTEM_PACKAGES: 1
  CUDA_VERSION: 13.0.1
  TOKENSPEED_B200_RUNNER_LABEL: ${{ vars.TOKENSPEED_B200_RUNNER_LABEL }}

concurrency:
  group: pr-test-nvidia-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: ${{ github.event_name != 'workflow_dispatch' }}

jobs:
  cancel-on-close:
    if: github.event_name == 'pull_request' && github.event.action == 'closed'
    runs-on: ubuntu-latest
    steps:
      - name: Cancel superseded PR Test runs
        run: |
          echo "Cancelling older PR Test NVIDIA runs for PR #${{ github.event.pull_request.number }}"

  scan:
    # Closed PRs only need cancel-on-close. For all other events, run jobs for
    # pull_request events or for push/manual events in lightseekorg/tokenspeed.
    if: |
      (github.event_name != 'pull_request' || github.event.action != 'closed') &&
      (
        github.event_name == 'pull_request' ||
        github.repository == 'lightseekorg/tokenspeed'
      )
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.scan.outputs.matrix }}
      install_tokenspeed_mla_from_source: ${{ steps.detect_mla.outputs.changed }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Install scan dependency
        # pypi.org occasionally returns "Content-Type: Unknown" for the simple
        # index, which strict pip>=22 refuses to parse. Retry a few times.
        run: |
          for i in 1 2 3 4 5; do
            python3 -m pip install pyyaml && exit 0
            echo "pyyaml install attempt $i failed, retrying in 10s..."
            sleep 10
          done
          exit 1

      - name: Build task matrix
        id: scan
        run: |
          MATRIX=$(python3 test/ci_system/pipeline.py scan \
            --root test/ci \
            --trigger "${{ github.event_name == 'workflow_dispatch' && github.event.inputs.trigger || 'per-commit' }}" \
            --runner-group "nvidia-x86")
          echo "matrix=${MATRIX}" >> "$GITHUB_OUTPUT"

      # Compute this once on the hosted runner (where `gh` is pre-installed)
      # and propagate the result to every matrix job via job output. The
      # downstream unit-test job reads it as an env var and `install_deps.sh`
      # gates its in-tree `tokenspeed-mla` reinstall on it. Without this
      # detection, CI would always exercise the pinned PyPI wheel and edits
      # under `tokenspeed-mla/` would be silently uncovered.
      - name: Detect tokenspeed-mla source changes
        id: detect_mla
        env:
          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        run: |
          set -euo pipefail
          case "${{ github.event_name }}" in
            pull_request)
              files=$(gh api \
                "/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }}/files" \
                --paginate -q '.[].filename')
              ;;
            push)
              files=$(gh api \
                "/repos/${{ github.repository }}/compare/${{ github.event.before }}...${{ github.event.after }}" \
                -q '.files[].filename')
              ;;
            *)
              files=""
              ;;
          esac
          if printf '%s\n' "$files" | grep -q '^tokenspeed-mla/'; then
            echo "tokenspeed-mla touched by this run; in-tree source will be installed."
            echo "changed=1" >> "$GITHUB_OUTPUT"
          else
            echo "tokenspeed-mla untouched; keeping pinned PyPI wheel."
            echo "changed=0" >> "$GITHUB_OUTPUT"
          fi

  unit-test:
    needs: scan
    if: |
      (github.event_name != 'pull_request' || github.event.action != 'closed') &&
      (
        github.event_name == 'pull_request' ||
        github.repository == 'lightseekorg/tokenspeed'
      ) &&
      (
        github.event_name != 'pull_request' ||
        github.event.pull_request.draft == false
      )
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.scan.outputs.matrix) }}
    name: ${{ matrix.name }} (${{ matrix.runner }})
    runs-on: ${{ matrix.runner }}
    timeout-minutes: 60
    continue-on-error: ${{ matrix.optional }}
    env:
      # When the scan job determined that this run touched `tokenspeed-mla/`,
      # `install_deps.sh` will reinstall the in-tree source on top of the
      # pinned PyPI wheel. See the matching detect step in the `scan` job.
      INSTALL_TOKENSPEED_MLA_FROM_SOURCE: ${{ needs.scan.outputs.install_tokenspeed_mla_from_source }}
      # Authenticate Hugging Face Hub downloads so checkpoint pulls during
      # `ts serve` startup are not throttled by the anonymous rate limit. The
      # 600s per-file `runtime-1gpu` timeout is otherwise exceeded on cold
      # runners (e.g. `amd-mi35x-1gpu-test` loading `openai/gpt-oss-120b` +
      # EAGLE3 draft).
      HF_TOKEN: ${{ secrets.HF_TOKEN }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        with:
          path: run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }}

      - name: Set work directory
        run: |
          echo "WORK_DIR=${{ github.workspace }}/run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }}" >> "$GITHUB_ENV"

      - name: Install pipeline dependency
        # See note in the scan job — pypi.org "Content-Type: Unknown" flake.
        run: |
          for i in 1 2 3 4 5; do
            python3 -m pip install pyyaml && exit 0
            echo "pyyaml install attempt $i failed, retrying in 10s..."
            sleep 10
          done
          exit 1

      # When the diff touches `tokenspeed-mla/`, ask install_deps.sh to
      # rebuild + install that package from the in-tree source after the
      # pinned PyPI wheel has been installed (transitively via
      # tokenspeed-kernel). Without this, CI keeps testing whichever
      # `tokenspeed-mla` version is pinned in
      # `tokenspeed-kernel/python/requirements/cuda-thirdparty.txt`, not
      # the code on the PR/branch. When the diff doesn't touch
      # `tokenspeed-mla/` we keep using the pinned PyPI wheel.
      - name: Install CI task
        run: |
          cd "${{ env.WORK_DIR }}"
          mkdir -p .ci-artifacts
          python3 test/ci_system/pipeline.py execute \
            --config "${{ matrix.config }}" \
            --runner "${{ matrix.runner }}" \
            --work-dir "${{ env.WORK_DIR }}" \
            --print-plan \
            --only-stage install \
            --keep-runner-state \
            --result-json .ci-artifacts/result.json

      - name: Execute CI task
        run: |
          cd "${{ env.WORK_DIR }}"
          mkdir -p .ci-artifacts
          python3 test/ci_system/pipeline.py execute \
            --config "${{ matrix.config }}" \
            --runner "${{ matrix.runner }}" \
            --work-dir "${{ env.WORK_DIR }}" \
            --print-plan \
            --skip-stage install \
            --reuse-runner-state \
            --result-json .ci-artifacts/result.json

      - name: Upload task result
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: pr-test-${{ matrix.name }}-${{ matrix.runner }}-${{ github.run_id }}-${{ github.run_attempt }}
          path: ${{ env.WORK_DIR }}/.ci-artifacts/result.json

      - name: Cleanup work directory
        if: always()
        run: rm -rf "${{ env.WORK_DIR }}"

  finish:
    if: |
      always() &&
      (github.event_name != 'pull_request' || github.event.action != 'closed') &&
      (
        github.event_name == 'pull_request' ||
        github.repository == 'lightseekorg/tokenspeed'
      )
    needs: [unit-test]
    runs-on: ubuntu-latest
    steps:
      - name: Finish
        run: echo "PR Test finished."