name: RL B200 Integration Tests

on:
  push:
    tags:
      - ciflow/b200/*
  schedule:
    - cron: '0 0 * * *'

concurrency:
  group: unit-test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
  cancel-in-progress: true

defaults:
  run:
    shell: bash -l -eo pipefail {0}

permissions:
  id-token: write
  contents: read

jobs:
  rl-kimi-k3:
    if: github.repository_owner == 'pytorch'
    runs-on: linux.dgx.b200.8
    timeout-minutes: 60
    container:
      image: nvidia/cuda:13.0.3-cudnn-devel-ubuntu22.04
      options: --gpus all
    steps:
      - name: Check out repo
        uses: actions/checkout@v7

      - name: Set up uv
        uses: astral-sh/setup-uv@v7
        with:
          python-version: '3.12'
          enable-cache: true

      - name: Run Kimi K3 RL tests
        run: |
          set -eux

          bash .ci/docker/common/install_base.sh

          uv venv --python 3.12
          source .venv/bin/activate

          DRIVER_VERSION=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n 1 || true)
          echo "CUDA driver version: ${DRIVER_VERSION}"

          uv pip install -r requirements.txt
          uv pip install -r torchtitan/rl/requirements.txt pytest
          uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main"
          uv pip install --no-deps "git+https://github.com/thinking-machines-lab/batch_invariant_ops.git@main"
          uv pip install torch torchvision vllm xformers --pre \
            --extra-index-url https://download.pytorch.org/whl/nightly/cu132 \
            --index-strategy unsafe-best-match
          uv pip install "flash-attn-4[cu13]>=4.0.0b31"
          uv pip uninstall tilelang

          export PYTHONPATH="$PWD:${PYTHONPATH:-}"

          python - <<'PY'
          import torch

          capability = torch.cuda.get_device_capability()
          if capability not in ((10, 0), (10, 3)):
              raise RuntimeError(
                  f"B200 RL integration requires SM100 or SM103, got SM{capability[0]}{capability[1]}"
              )
          PY

          OUTPUT_DIR="$RUNNER_TEMP/artifacts-to-be-uploaded/rl"
          mkdir -p "$OUTPUT_DIR"

          RL_TEST_DUMP_FOLDER="$OUTPUT_DIR" torchrun --nproc-per-node=1 -m pytest \
            tests/rl/unit_tests/gpu/test_bitwise_parity.py::TestBitwiseParityKimiK3DebugVarlen -v

          python -m tests.rl.integration_tests.rl "$OUTPUT_DIR" --ngpu 8 \
            --test_suite kda

      - name: Upload artifacts
        if: always()
        uses: actions/upload-artifact@v7
        with:
          name: b200-rl-outputs
          path: ${{ runner.temp }}/artifacts-to-be-uploaded
