name: RL H100 Integration Tests

on:
  push:
    branches: [ main ]
    tags:
      - ciflow/rl/*
  schedule:
    # Runs every 12 hours
    - cron: '0 */12 * * *'

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:
  set-matrix:
    if: github.repository_owner == 'pytorch' || github.event_name != 'schedule'
    uses: ./.github/workflows/set-matrix.yaml
    with:
      runner-cuda: mt-l-bx86iamx-176-1800-h100-8
      gpu-arch: cuda
      is-experimental: true

  rl-kimi-k3:
    needs: set-matrix
    if: ${{ needs.set-matrix.outputs.matrix != '' && fromJSON(needs.set-matrix.outputs.matrix).include[0] != null }}
    uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main
    strategy:
      fail-fast: false
      matrix: ${{ fromJSON(needs.set-matrix.outputs.matrix) }}
    with:
      runner: ${{ matrix.runner }}
      gpu-arch-type: ${{ matrix.gpu-arch-type }}
      gpu-arch-version: ${{ matrix.gpu-arch-version }}
      docker-image: 308535385114.dkr.ecr.us-east-1.amazonaws.com/torchtitan/torchtitan-ubuntu-22.04-clang12:rl-${{ needs.set-matrix.outputs.docker-hash }}
      repository: pytorch/torchtitan
      upload-artifact: h100-rl-outputs
      timeout: 60
      script: |
        set -eux

        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
        conda activate "${CONDA_ENV}"

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

        pip config --user set global.progress_bar off

        python -m pip install uv

        uv pip install --no-deps "git+https://github.com/thinking-machines-lab/batch_invariant_ops.git@main"

        uv pip install --upgrade torch torchvision vllm xformers --pre \
          --extra-index-url https://download.pytorch.org/whl/nightly/cu132 \
          --index-strategy unsafe-best-match \
          --constraint .ci/docker/requirements.txt
        python -c '
        from importlib.metadata import version
        wheels = {name: version(name) for name in ("torch", "torchvision", "vllm")}
        print(wheels)
        assert all(".dev" in wheel and "+cu132" in wheel for wheel in wheels.values()), wheels
        '

        uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main"
        uv pip install "flash-attn-4[cu13]>=4.0.0b31"
        uv pip install flash-attn-3 \
          --extra-index-url https://download.pytorch.org/whl/test/cu130 \
          --index-strategy unsafe-best-match
        uv pip uninstall tilelang
        uv pip install nvidia-cuda-nvcc

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

        python - <<'PY'
        import torch

        capability = torch.cuda.get_device_capability()
        if capability != (9, 0):
            raise RuntimeError(
                f"H100 RL integration requires SM90, got SM{capability[0]}{capability[1]}"
            )
        PY

        sudo mkdir -p "$RUNNER_ARTIFACT_DIR"
        sudo chown -R "$(id -u):$(id -g)" "$RUNNER_ARTIFACT_DIR"
        OUTPUT_DIR="$RUNNER_ARTIFACT_DIR/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
