mirror of
https://github.com/ROCm/jax.git
synced 2025-04-14 19:06:07 +00:00
022da913e6
* Account test totals correctly for dashboards * Add blurb to the dev guide on skipping tests * Remove extra newline * Default to 0 if "skipped" isn't found Co-authored-by: Mathew Odden <1471252+mrodden@users.noreply.github.com> --------- Co-authored-by: Mathew Odden <1471252+mrodden@users.noreply.github.com>
70 lines
2.1 KiB
YAML
70 lines
2.1 KiB
YAML
name: ROCm GPU CI
|
|
|
|
on:
|
|
# Trigger the workflow on push or pull request,
|
|
# but only for the rocm-main branch
|
|
push:
|
|
branches:
|
|
- rocm-main
|
|
- 'rocm-jaxlib-v*'
|
|
pull_request:
|
|
branches:
|
|
- rocm-main
|
|
- 'rocm-jaxlib-v*'
|
|
|
|
concurrency:
|
|
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
|
|
cancel-in-progress: true
|
|
|
|
jobs:
|
|
build-jax-in-docker: # strategy and matrix come here
|
|
runs-on: mi-250
|
|
env:
|
|
BASE_IMAGE: "ubuntu:22.04"
|
|
TEST_IMAGE: ubuntu-jax-${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
|
|
PYTHON_VERSION: "3.10"
|
|
ROCM_VERSION: "6.2.4"
|
|
WORKSPACE_DIR: workdir_${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
|
|
steps:
|
|
- name: Clean up old runs
|
|
run: |
|
|
ls
|
|
# Make sure that we own all of the files so that we have permissions to delete them
|
|
docker run -v "./:/jax" ubuntu /bin/bash -c "chown -R $UID /jax/workdir_* || true"
|
|
# Remove any old work directories from this machine
|
|
rm -rf workdir_*
|
|
ls
|
|
- name: Print system info
|
|
run: |
|
|
whoami
|
|
printenv
|
|
df -h
|
|
rocm-smi
|
|
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
|
|
with:
|
|
path: ${{ env.WORKSPACE_DIR }}
|
|
- name: Build JAX
|
|
run: |
|
|
pushd $WORKSPACE_DIR
|
|
python3 build/rocm/ci_build \
|
|
--rocm-version $ROCM_VERSION \
|
|
--base-docker $BASE_IMAGE \
|
|
--python-versions $PYTHON_VERSION \
|
|
--compiler=clang \
|
|
dist_docker \
|
|
--image-tag $TEST_IMAGE
|
|
- name: Archive jax wheels
|
|
uses: actions/upload-artifact@v4
|
|
with:
|
|
name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION }}_id${{ github.run_id }}
|
|
path: ./dist/*.whl
|
|
- name: Run tests
|
|
env:
|
|
ROCM_TEST_INCLUDE_SKIPS: "1"
|
|
GPU_COUNT: "8"
|
|
GFX: "gfx90a"
|
|
run: |
|
|
cd $WORKSPACE_DIR
|
|
python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd "pytest tests/core_test.py"
|
|
|