Skip to content

CI - with Numpy/Scipy nightly wheels (nightly) #67

CI - with Numpy/Scipy nightly wheels (nightly)

CI - with Numpy/Scipy nightly wheels (nightly) #67

name: CI - with Numpy/Scipy nightly wheels (nightly)
# This configures a github action that runs the JAX test suite against nightly development builds
# of numpy and scipy, in order to catch issues with new package versions prior to their release.
# Unlike our other CI, this is one that we expect to fail frequently, and so we don't run it against
# every commit and PR in the repository. Rather, we run it on a schedule, and failures lead to an
# issue being created or updated.
# Portions of this adapted from https://github.com/pydata/xarray/blob/main/.github/workflows/upstream-dev-ci.yaml
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
on:
schedule:
- cron: "0 12 * * *" # Daily at 12:00 UTC
workflow_dispatch: # allows triggering the workflow run manually
pull_request: # Automatically trigger on pull requests affecting this file
branches:
- main
paths:
- '**workflows/upstream-nightly.yml'
jobs:
upstream-dev:
runs-on: ROCM-Ubuntu
permissions:
contents: read
checks: write # for upload-artifact
defaults:
run:
shell: bash -l {0}
strategy:
fail-fast: false
matrix:
python-version: ["3.13"]
outputs:
artifacts_availability: ${{ steps.status.outputs.ARTIFACTS_AVAILABLE }}
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
with:
python-version: ${{ matrix.python-version }}
- name: Install JAX test requirements
run: |
pip install -r build/test-requirements.txt
pip install pytest-reportlog
- name: Install numpy & scipy development versions
run: |
pip install \
-i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple \
--no-deps \
--pre \
--upgrade \
numpy \
scipy
python -c "import numpy; print(f'{numpy.__version__=}')"
python -c "import scipy; print(f'{scipy.__version__=}')"
- name: Install JAX
run: |
pip install .[ci]
- name: Run tests
if: success()
id: status
env:
JAX_NUM_GENERATED_CASES: 1
JAX_ENABLE_X64: true
JAX_ENABLE_CHECKS: true
JAX_SKIP_SLOW_TESTS: true
PY_COLORS: 1
run: |
echo "JAX_NUM_GENERATED_CASES=$JAX_NUM_GENERATED_CASES"
echo "JAX_ENABLE_X64=$JAX_ENABLE_X64"
echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
echo "JAX_SKIP_SLOW_TESTS=$JAX_SKIP_SLOW_TESTS"
pytest -n auto --tb=short -rf --maxfail=20 \
--report-log output-${{ matrix.python-version }}-log.jsonl \
tests \
|| (
echo 'ARTIFACTS_AVAILABLE=true' >> $GITHUB_OUTPUT && false
)
- name: Upload artifacts
if: |
failure()
&& steps.status.outcome == 'failure'
&& github.event_name == 'schedule'
&& github.repository == 'jax-ml/jax'
uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
with:
name: output-${{ matrix.python-version }}-log.jsonl
path: output-${{ matrix.python-version }}-log.jsonl
retention-days: 5
report:
name: report
needs: upstream-dev
permissions:
contents: read
issues: write
if: |
failure()
&& github.event_name == 'schedule'
&& needs.upstream-dev.outputs.artifacts_availability == 'true'
runs-on: ubuntu-latest
defaults:
run:
shell: bash
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
with:
python-version: "3.x"
- uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
with:
path: /tmp/workspace/logs
- name: install requirements
run: |
python -m pip install pytest
- name: Move all log files into a single directory
run: |
rsync -a /tmp/workspace/logs/output-*/ ./logs
ls -R ./logs
cat logs/*.jsonl > pytest-logs.txt
python .github/workflows/parse_logs.py pytest-logs.txt --outfile=parsed-logs.txt
- name: Report failures
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const fs = require('fs');
const parsed_logs = fs.readFileSync('parsed-logs.txt', 'utf8');
const title = "⚠️ Nightly upstream-dev CI failed ⚠️"
const workflow_url = `https://github.com/${process.env.GITHUB_REPOSITORY}/actions/runs/${process.env.GITHUB_RUN_ID}`
const issue_body = `[Workflow Run URL](${workflow_url})\n${parsed_logs}`
// Run GraphQL query against GitHub API to find the most recent open issue used for reporting failures
const query = `query($owner:String!, $name:String!, $creator:String!, $label:String!){
repository(owner: $owner, name: $name) {
issues(first: 1, states: OPEN, filterBy: {createdBy: $creator, labels: [$label]}, orderBy: {field: CREATED_AT, direction: DESC}) {
edges {
node {
body
id
number
}
}
}
}
}`;
const variables = {
owner: context.repo.owner,
name: context.repo.repo,
label: 'CI',
creator: "github-actions[bot]"
}
const result = await github.graphql(query, variables)
// If no issue is open, create a new issue,
// else update the body of the existing issue.
if (result.repository.issues.edges.length === 0) {
github.rest.issues.create({
owner: variables.owner,
repo: variables.name,
body: issue_body,
title: title,
labels: [variables.label]
})
} else {
github.rest.issues.update({
owner: variables.owner,
repo: variables.name,
issue_number: result.repository.issues.edges[0].node.number,
body: issue_body
})
}