Skip to content

Update pyrefly pre-commit hook. #6245

Update pyrefly pre-commit hook.

Update pyrefly pre-commit hook. #6245

Workflow file for this run

name: Tests
on:
pull_request:
push:
branches:
- main
- release/*
concurrency:
group: ${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
cancel-in-progress: true
jobs:
load-matrix:
runs-on: ubuntu-latest
outputs:
matrix: ${{ steps.set-matrix.outputs.matrix }}
steps:
- name: Checkout repository
uses: actions/checkout@v6
- name: Load matrix from file
id: set-matrix
run: |
matrix=$(cat .github/matrix.json | jq -c .)
echo "matrix=$matrix" >> $GITHUB_OUTPUT
test:
needs: load-matrix
strategy:
fail-fast: false
matrix: ${{ fromJSON(needs.load-matrix.outputs.matrix) }}
name: test-${{ matrix.runtime-version }}-py${{ matrix.python-version }}-${{ matrix.pytorch-version }}-${{ matrix.alias }}
container:
image: ${{ matrix.image }}
options: ${{ matrix.container-options }}
runs-on: ${{ matrix.runner }}
defaults:
run:
shell: bash -l {0}
steps:
- name: Run NVIDIA command
if: startsWith(matrix.image, 'nvidia')
run: |
echo "Detected NVIDIA image"
nvidia-smi || echo "nvidia-smi not found"
- name: Run ROCm command
if: startsWith(matrix.image, 'rocm')
run: |
echo "Detected ROCm image"
rocminfo || echo "rocminfo not found"
- name: Check out code
uses: actions/checkout@v6
- name: Install system dependencies
run: |
set -eux
apt-get update
apt-get install -y libdw1 curl wget git pkg-config zlib1g-dev build-essential
- name: Install NVSHMEM
if: contains(matrix.alias, 'distributed')
run: |
set -euxo pipefail
GPU_COUNT=$(nvidia-smi -L | wc -l)
if [ "$GPU_COUNT" -ne 4 ]; then
echo "Error: Expected 4 GPUs but found $GPU_COUNT"
exit 1
fi
curl -L https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/common/install_cuda.sh -o install_cuda.sh
chmod +x install_cuda.sh
source install_cuda.sh
install_nvshmem 13 3.4.5
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
python-version: ${{ matrix.python-version }}
enable-cache: true
- name: Create virtual environment
run: |
uv venv --python ${{ matrix.python-version }}
- name: Get current month
id: date
run: echo "month=$(date +'%Y-%m')" >> $GITHUB_OUTPUT
- name: Cache dependencies
id: cache
uses: actions/cache@v4
with:
path: |
~/.cache/uv
~/.venv
key: ${{ matrix.python-version }}-${{ matrix.runtime-version }}-${{ matrix.pytorch-version }}-${{ hashFiles('.github/workflows/test.yml') }}-${{ steps.date.outputs.month }}
- name: Install PyTorch
run: |
source .venv/bin/activate
if [[ "${{ matrix.pytorch-version }}" == "pytorch-2.9" ]]; then
# Install stable 2.9 from test channel
uv pip install -U "torch==2.9.*" --index-url https://download.pytorch.org/whl/${{ matrix.runtime-version }}
else
# Default to nightly
if [[ "${{ matrix.runtime-version }}" == "cu128" ]]; then
# Install nvidia-nvshmem-cu12 from cu129 index (missing on cu128)
uv pip install -U --pre nvidia-nvshmem-cu12 --index-url https://download.pytorch.org/whl/nightly/cu129
fi
uv pip install -U --pre torch --index-url https://download.pytorch.org/whl/nightly/${{ matrix.runtime-version }}
fi
- name: Install Triton
if: contains(matrix.alias, 'cpu') || (steps.cache.outputs.cache-hit != 'true' && matrix.pytorch-version != 'pytorch-2.9')
run: |
set -x
source .venv/bin/activate
apt-get update
apt-get install -y git
apt-get install -y clang-20 clang++-20 zlib1g-dev
export CC=clang-20
export CXX=clang++-20
mkdir -p /tmp/$USER
cd /tmp/$USER
uv pip uninstall triton pytorch-triton || true
rm -rf triton/ || true
if [[ "${{ matrix.alias }}" == *cpu* ]]; then
git clone --recursive -b main-merged https://github.com/triton-lang/triton-cpu.git triton
else
git clone https://github.com/triton-lang/triton.git triton
fi
cd triton/
uv pip install -r python/requirements.txt
MAX_JOBS=$(nproc) TRITON_PARALLEL_LINK_JOBS=2 uv pip install .
cd /tmp/$USER
rm -rf triton/
python -c "import triton; print(f'Triton version: {triton.__version__}')"
- name: Install Helion
run: |
source .venv/bin/activate
uv pip install setuptools ninja
SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,surrogate]'
python -c "import helion; print(helion.__name__)"
- name: Run Tests
run: |
set -o pipefail
source .venv/bin/activate
# Conditionally enable ref-eager and golden-accept/dtype-assert test modes
if [[ "${{ matrix.dtype-asserts }}" == "true" ]]; then export HELION_DEBUG_DTYPE_ASSERTS=1; fi
if [[ "${{ matrix.expecttest-accept }}" == "true" ]]; then export EXPECTTEST_ACCEPT=1; fi
if [[ "${{ matrix.ref-eager }}" == "true" ]]; then export HELION_INTERPRET=1; fi
if [[ "${{ contains(matrix.alias, 'cpu') }}" == "true" ]]; then export TRITON_CPU_BACKEND=1; fi
# -rf: print failed tests
# --timeout: max allowed time for each test
TEST_PATH=$([[ "${{ contains(matrix.alias, 'distributed') }}" == "true" ]] && echo "test/test_examples_dist.py" || echo ".")
EXTRA_FLAGS=$([[ "${{ contains(matrix.alias, 'distributed') }}" == "true" ]] && echo "-rs" || echo "--ignore=test/test_examples_dist.py")
# For distributed tests, fail if any test is skipped, failed, or has an error
SKIP_CHECK=$([[ "${{ contains(matrix.alias, 'distributed') }}" == "true" ]] && echo "! grep -qE '(SKIPPED|FAILED|ERROR)'" || echo "cat")
pytest -rf --timeout=60 $EXTRA_FLAGS $TEST_PATH | tee >(eval $SKIP_CHECK)
test-notebooks:
name: test-notebooks-cu128-py3.12-pytorch-2.9-a10g
container:
image: nvidia/cuda:12.8.1-devel-ubuntu24.04
options: --gpus all
runs-on: linux.g5.4xlarge.nvidia.gpu
defaults:
run:
shell: bash -l {0}
steps:
- name: Run NVIDIA command
run: |
echo "Detected NVIDIA image"
nvidia-smi || echo "nvidia-smi not found"
- name: Check out code
uses: actions/checkout@v6
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
python-version: "3.12"
enable-cache: true
- name: Create virtual environment
run: |
uv venv --python 3.12
- name: Install pip in venv
run: |
source .venv/bin/activate
uv pip install pip
- name: Get current month
id: date
run: echo "month=$(date +'%Y-%m')" >> $GITHUB_OUTPUT
- name: Cache dependencies
id: cache
uses: actions/cache@v4
with:
path: |
~/.cache/uv
~/.venv
key: notebooks-3.12-cu128-${{ hashFiles('.github/workflows/test.yml') }}-${{ steps.date.outputs.month }}
- name: Install notebook execution tools
run: |
source .venv/bin/activate
# Install jupyter for executing notebooks
uv pip install jupyter nbconvert pytest numpy
- name: Run Notebook Tests
run: |
source .venv/bin/activate
# Execute notebook using jupyter nbconvert
# The notebook's subprocess pip install will install torch and helion
jupyter nbconvert --to notebook --execute --inplace \
--ExecutePreprocessor.timeout=600 \
notebooks/softmax.ipynb