<% set megatron_tests = [
    {'test_file': 'test_full_disk_weight_update.py', 'num_gpus': 4},
    {'test_file': 'test_quick_start_glm4_9B.py', 'num_gpus': 8, 'enable_eval': '0'},
    {'test_file': 'test_glm4.7_30B_A3B_pd_mooncake.py', 'num_gpus': 8},
    {'test_file': 'test_qwen3_30B_A3B.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1', 'enable_eval': '0'},
    {'test_file': 'test_qwen3.6_35B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'},
    {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1', 'enable_eval': '0'},
    {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8, 'enable_eval': '0'},
    {'test_file': 'test_qwen3_4B_ppo_disaggregate.py', 'num_gpus': 8, 'enable_eval': '0'},
    {'test_file': 'test_qwen3_4B_ppo_train_critic_only.py', 'num_gpus': 8, 'enable_eval': '0'},
    {'test_file': 'test_ppo_logprob_entropy_gpu.py', 'num_gpus': 2},
    {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8, 'enable_eval': '0'},
    {'test_file': 'test_moonlight_16B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'},
    {'test_file': 'test_mimo_7B_mtp_only_grad.py', 'num_gpus': 8},
    {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8},
    {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8},
    {'test_file': 'test_qwen2.5_0.5B_opd_sglang.py', 'num_gpus': 8},
    {'test_file': 'test_qwen3_4B_external_pd.py', 'num_gpus': 6},
    {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4},
    {'test_file': 'test_qwen2.5_0.5B_fanout_short.py', 'num_gpus': 4},
    {'test_file': 'test_qwen3_4B_streaming_partial_rollout.py', 'num_gpus': 8},
    {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4},
    {'test_file': 'test_qwen3.5_0.8B_gsm8k_async_short.py', 'num_gpus': 4},
    {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer gpu --load-optimizer gpu', 'num_gpus': 8},
    {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer gpu --load-optimizer cpu', 'num_gpus': 8},
    {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer cpu --load-optimizer cpu', 'num_gpus': 8},
    {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer cpu --load-optimizer gpu', 'num_gpus': 8},
    {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8},
] %>
<% set jobs = {
    'e2e-test-sglang-config': {
      'label': 'run-ci-sglang-config',
      'tests': [
        {'test_file': 'test_qwen2.5_0.5B_sglang_config.py', 'num_gpus': 8},
        {'test_file': 'test_qwen2.5_0.5B_sglang_config_distributed.py', 'num_gpus': 8},
        {'test_file': 'test_sglang_config_mixed_offload.py', 'num_gpus': 8},
        {'test_file': 'test_sglang_config_mixed_offload_ft.py', 'num_gpus': 8},
      ],
    },
    'e2e-test-megatron': {
      'label': 'run-ci-megatron',
      'tests': megatron_tests,
    },
    'e2e-test-precision': {
      'label': 'run-ci-precision',
      'tests': [
        {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8},
      ],
    },
    'e2e-test-ckpt': {
      'label': 'run-ci-ckpt',
      'tests': [
        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer gpu --load-optimizer gpu', 'num_gpus': 8},
        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer gpu --load-optimizer cpu', 'num_gpus': 8},
        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer cpu --load-optimizer cpu', 'num_gpus': 8},
        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--save-optimizer cpu --load-optimizer gpu', 'num_gpus': 8},
        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8},
      ],
    },

    'cpu-unittest': {
      'label': 'run-ci-cpu-unittest',
      'always': True,
      'cpu': True,
      'tests': [
        {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0},
        {'test_file': 'utils/test_megatron_server_arguments.py', 'num_gpus': 0},
        {'test_file': 'test_dp_schedule.py', 'num_gpus': 0},
        {'test_file': 'test_cp_utils.py', 'num_gpus': 0},
        {'test_file': 'test_metric_report.py', 'num_gpus': 0},
        {'test_file': 'test_metric_report_dist.py', 'num_gpus': 0},
        {'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0},
        {'test_file': 'test_logprob_response_spans.py', 'num_gpus': 0},
        {'test_file': 'test_value_temperature.py', 'num_gpus': 0},
        {'test_file': 'test_cispo_loss.py', 'num_gpus': 0},
        {'test_file': 'test_ppo_logprob_entropy.py', 'num_gpus': 0},
        {'test_file': 'test_rm_f1.py', 'num_gpus': 0},
        {'test_file': 'test_rm_gpqa.py', 'num_gpus': 0},
        {'test_file': 'test_rm_math.py', 'num_gpus': 0},
        {'test_file': 'test_rm_math_dapo.py', 'num_gpus': 0},
        {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0},
        {'test_file': 'test_sample.py', 'num_gpus': 0},
        {'test_file': 'test_rollout_validation.py', 'num_gpus': 0},
        {'test_file': 'test_placement_group.py', 'num_gpus': 0},
        {'test_file': 'test_external_sglang_engines.py', 'num_gpus': 0},
        {'test_file': 'test_empty_colocated_weight_bucket.py', 'num_gpus': 0},
        {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0},
        {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0},
        {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0},
        {'test_file': 'plugin_contracts/test_plugin_path_loading_contracts.py', 'num_gpus': 0},
        {'test_file': 'plugin_contracts/test_plugin_generate_contracts.py', 'num_gpus': 0},
      ],
    },

    'agent-test': {
      'label': 'run-ci-agent',
      'always': True,
      'cpu': True,
      'tests': [
        {'test_file': 'test_agent/test_trajectory_manager_branching.py', 'num_gpus': 0},
        {'test_file': 'test_agent/test_adapters.py', 'num_gpus': 0},
        {'test_file': 'test_agent/test_harness.py', 'num_gpus': 0},
        {'test_file': 'test_agent/test_agent_rollout_cpu.py', 'num_gpus': 0},
      ],
    },

    'e2e-test-image': {
      'label': 'run-ci-image',
      'image': 'slimerl/slime-test:latest',
      'tests': megatron_tests,
    },
} %>
name: PR Test

on:
  # Push to main triggers the default GitHub-hosted jobs (cheap CPU/unit and
  # agent checks), catching PR-pair regressions where two PRs pass
  # individually but main is broken after both land. GPU jobs stay label-gated
  # — see each job's `if:` below — so push events never burn the self-hosted
  # fleet.
  push:
    branches: [main]
  pull_request:
    branches: [main]
    types: [opened, reopened, synchronize, labeled]
  workflow_dispatch:
    inputs:
      infinite_run:
        description: 'Run training infinitely'
        required: false
        type: boolean
        default: false

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

jobs:
<% for job_name, config in jobs.items() %>
  << job_name >>:
<% if config.get('always') %>
    if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' || github.event_name == 'push'
<% else %>
    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, '<< config.label >>'))
<% endif %>
<% if config.get('cpu') %>
    runs-on: ubuntu-latest
<% else %>
    runs-on: self-hosted
<% endif %>
    strategy:
      fail-fast: false
      matrix:
        info:
<< config.tests | tojson(indent=2) | indent(10, true) >>
    defaults:
      run:
        working-directory: ${{ github.workspace }}
    env:
      GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
      SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }}
      SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }}
      SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }}
      SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }}

    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
<% if config.get('cpu') %>

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.10'
          cache: 'pip'

      - name: Install dependencies
        shell: bash
        run: |
          pip install torch --index-url https://download.pytorch.org/whl/cpu
          pip install pytest numpy packaging pyyaml omegaconf tqdm httpx requests ray pybase64 pylatexenc sympy aiohttp pillow safetensors
<% if config.get('extra_pip_deps') %>
          pip install << config.extra_pip_deps >>
<% endif %>

      - name: Install
        shell: bash
        run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps
<% else %>
<% endif %>

      - name: Execute
        shell: bash
        run: |
<% if config.get('cpu') %>
          TEST_PATH="${{ matrix.info.test_file }}"
          if [[ "$TEST_PATH" != tests/* ]]; then
            TEST_PATH="tests/$TEST_PATH"
          fi
          TEST_ARGS="${{ matrix.info.test_args || '' }}"
          if [[ -n "$TEST_ARGS" ]]; then
            read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS")
          else
            TEST_ARGS_ARRAY=()
          fi
          if [ "${{ matrix.info.num_gpus }}" = "0" ]; then
            python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
          else
            python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
          fi
<% else %>
          docker run --pull=always --rm \
            --privileged \
            --cap-add SYS_NICE \
            --security-opt seccomp=unconfined \
            --network host \
            --gpus all \
            --ipc=host \
            --shm-size=16g \
            --ulimit memlock=-1 \
            --ulimit stack=67108864 \
            --memory=0 \
            --memory-swap=0 \
            -e http_proxy \
            -e https_proxy \
            -e HTTP_PROXY \
            -e HTTPS_PROXY \
            -e GITHUB_COMMIT_NAME \
            -e WANDB_API_KEY \
            -e SLIME_TEST_ENABLE_INFINITE_RUN \
            -e SLIME_TEST_USE_DEEPEP \
            -e SLIME_TEST_USE_FP8_ROLLOUT \
            -e SLIME_TEST_ENABLE_EVAL \
            -e TEST_FILE="${{ matrix.info.test_file }}" \
            -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \
            -e NUM_GPUS="${{ matrix.info.num_gpus }}" \
            -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \
            -v /mnt/nvme0n1/slime_ci:/data/slime_ci \
            -v /mnt/nvme0n1/slime_ci/models:/root/models \
            -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \
            -w "$GITHUB_WORKSPACE" \
            << config.image if config.image else 'slimerl/slime:latest' >> \
            bash -lc '
              set -euo pipefail
              pip install -e . --no-deps --break-system-packages
              TEST_PATH="$TEST_FILE"
              if [[ "$TEST_PATH" != tests/* ]]; then
                TEST_PATH="tests/$TEST_PATH"
              fi
              if [[ -n "$TEST_ARGS" ]]; then
                read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS")
              else
                TEST_ARGS_ARRAY=()
              fi
              if [ "$NUM_GPUS" = "0" ]; then
                python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
              else
                python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
              fi
            '
<% endif %>
<% endfor %>

  e2e-test-changed-detect:
    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-changed'))
    runs-on: self-hosted
    outputs:
      matrix: ${{ steps.detect.outputs.matrix }}
      has_tests: ${{ steps.detect.outputs.has_tests }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          fetch-depth: 0

      - name: Detect changed tests
        id: detect
        shell: bash
        run: |
          CHANGED=$(git diff --name-only --diff-filter=AM origin/main...HEAD -- 'tests/test_*.py' 'tests/plugin_contracts/test_*.py' || true)
          if [ -z "$CHANGED" ]; then
            echo "No new or modified test files found."
            echo "has_tests=false" >> $GITHUB_OUTPUT
            echo 'matrix={"info":[]}' >> $GITHUB_OUTPUT
          else
            echo "Changed test files:"
            echo "$CHANGED"
            MATRIX="["
            FIRST=true
            for filepath in $CHANGED; do
              # Extract NUM_GPUS from the test file, default to 8
              NGPU=$(grep -oP '^NUM_GPUS\s*=\s*\K\d+' "$filepath" | head -1)
              NGPU=${NGPU:-8}
              if [ "$FIRST" = true ]; then FIRST=false; else MATRIX+=","; fi
              MATRIX+="{\"test_file\":\"$filepath\",\"num_gpus\":$NGPU}"
            done
            MATRIX+="]"
            echo "has_tests=true" >> $GITHUB_OUTPUT
            echo "matrix={\"info\":$MATRIX}" >> $GITHUB_OUTPUT
            echo "Generated matrix: $MATRIX"
          fi

  e2e-test-changed:
    needs: e2e-test-changed-detect
    if: needs.e2e-test-changed-detect.outputs.has_tests == 'true'
    runs-on: self-hosted
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.e2e-test-changed-detect.outputs.matrix) }}
    defaults:
      run:
        working-directory: ${{ github.workspace }}
    env:
      GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
      SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }}
      SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }}
      SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }}
      SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }}

    steps:
      - name: Checkout repository
        uses: actions/checkout@v4

      - name: Execute
        shell: bash
        run: |
          docker run --pull=always --rm \
            --privileged \
            --cap-add SYS_NICE \
            --security-opt seccomp=unconfined \
            --network host \
            --gpus all \
            --ipc=host \
            --shm-size=16g \
            --ulimit memlock=-1 \
            --ulimit stack=67108864 \
            --memory=0 \
            --memory-swap=0 \
            -e http_proxy \
            -e https_proxy \
            -e HTTP_PROXY \
            -e HTTPS_PROXY \
            -e GITHUB_COMMIT_NAME \
            -e WANDB_API_KEY \
            -e SLIME_TEST_ENABLE_INFINITE_RUN \
            -e SLIME_TEST_USE_DEEPEP \
            -e SLIME_TEST_USE_FP8_ROLLOUT \
            -e SLIME_TEST_ENABLE_EVAL \
            -e TEST_FILE="${{ matrix.info.test_file }}" \
            -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \
            -e NUM_GPUS="${{ matrix.info.num_gpus }}" \
            -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \
            -v /mnt/nvme0n1/slime_ci:/data/slime_ci \
            -v /mnt/nvme0n1/slime_ci/models:/root/models \
            -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \
            -w "$GITHUB_WORKSPACE" \
            slimerl/slime:latest \
            bash -lc '
              set -euo pipefail
              pip install -e . --no-deps --break-system-packages
              TEST_PATH="$TEST_FILE"
              if [[ "$TEST_PATH" != tests/* ]]; then
                TEST_PATH="tests/$TEST_PATH"
              fi
              if [[ -n "$TEST_ARGS" ]]; then
                read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS")
              else
                TEST_ARGS_ARRAY=()
              fi
              if [ "$NUM_GPUS" = "0" ]; then
                python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
              else
                python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
              fi
            '