-
Notifications
You must be signed in to change notification settings - Fork 5
GH-90: Add workflow files for test CI jobs #4
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| name: CI | ||
|
|
||
| on: | ||
| push: | ||
| branches: | ||
| - master | ||
| - 'release/*' | ||
| pull_request: | ||
| branches: | ||
| - master | ||
| - 'release/*' | ||
|
|
||
| concurrency: | ||
| group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} | ||
| cancel-in-progress: true | ||
|
|
||
| env: | ||
| PYTHON_VERSIONS: "3.10.13,3.12.8" | ||
|
|
||
| jobs: | ||
| call-build-wheels: | ||
| strategy: | ||
| matrix: | ||
| rocm-version: ["6.3.3"] | ||
| uses: jax/rocm-jax/.github/workflows/reuseable/build-wheels.yml@master | ||
| with: | ||
| python-versions: ${{ env.PYTHON_VERSIONS }} | ||
| rocm-version: ${{ matrix.rocm-version }} | ||
| call-build-docker: | ||
| strategy: | ||
| matrix: | ||
| python-version: ["3.10.13,3.12.8"] | ||
| rocm-version: ["6.3.3"] | ||
| uses: jax/rocm-jax/.github/workflows/reusable/build-docker.yml@master | ||
| with: | ||
| rocm-version: ${{ matrix.python-version }} | ||
| run-python-unit-tests: | ||
| strategy: | ||
| matrix: | ||
| rocm-version: ["6.3.3"] | ||
| steps: | ||
| - name: Checkout plugin repo | ||
| uses: actions/checkout@v4 | ||
| - name: Checkout JAX repo | ||
| uses: actions/chekcout@v4 | ||
| with: | ||
| repository: jax-ml/jax | ||
| path: jax | ||
| - name: Run tests | ||
| env: | ||
| GPU_COUNT: "8" | ||
| GFX: "gfx90a" | ||
| ROCM_VERSION: ${{ matrix.rocm-version }} | ||
| run: | | ||
| python3 build/ci_build test "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" --test-cmd "pytest tests/core_test.py tests/linalg_test.py" | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| name: Nightly | ||
|
|
||
| on: | ||
| schedule: | ||
| - cron: "0 1 * * *" | ||
|
|
||
| concurrency: | ||
| cancel-in-progress: true | ||
|
|
||
| env: | ||
| PYTHON_VERSIONS: "3.10.13,3.12.8" | ||
|
|
||
| jobs: | ||
| call-build-wheels: | ||
| strategy: | ||
| matrix: | ||
| rocm-version: ["6.3.3"] | ||
| uses: jax/rocm-jax/.github/workflows/reuseable/build-wheels.yml@master | ||
| with: | ||
| python-versions: ${{ env.PYTHON_VERSIONS }} | ||
| rocm-version: ${{ matrix.rocm-version }} | ||
| call-build-docker: | ||
| strategy: | ||
| matrix: | ||
| python-version: ["3.10.13,3.12.8"] | ||
| rocm-version: ["6.3.3"] | ||
| uses: jax/rocm-jax/.github/workflows/reusable/build-docker.yml@master | ||
| with: | ||
| rocm-version: ${{ matrix.python-version }} | ||
| run-python-unit-tests: | ||
| strategy: | ||
| matrix: | ||
| rocm-version: ["6.3.3"] | ||
| steps: | ||
| - name: Checkout plugin repo | ||
| uses: actions/checkout@v4 | ||
| - name: Checkout JAX repo | ||
| uses: actions/chekcout@v4 | ||
| with: | ||
| repository: jax-ml/jax | ||
| path: jax | ||
| - name: Run tests | ||
| env: | ||
| GPU_COUNT: "8" | ||
| GFX: "gfx90a" | ||
| ROCM_VERSION: ${{ matrix.rocm-version }} | ||
| run: | | ||
| python3 build/ci_build test "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" --test-cmd "pytest tests" | ||
|
|
||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| name: Build docker | ||
|
|
||
| on: | ||
| workflow_call: | ||
| inputs: | ||
| rocm-version: | ||
| required: true | ||
| type: string | ||
| artifact-prefix: | ||
| default: 'plugin_wheels' | ||
|
|
||
| jobs: | ||
| build-docker: | ||
| runs-on: mi-250 | ||
| 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 "./:/rocm-jax" ubuntu /bin/bash -c "chown -R $UID /rocm-jax/* || true" | ||
| # Remove any old work directories from this machine | ||
| rm -rf * || true | ||
| ls | ||
| - name: Print system info | ||
| run: | | ||
| whoami | ||
| printenv | ||
| df -h | ||
| - uses: actions/checkout@v4 | ||
| - name: Download wheel artifacts | ||
| uses: actions/download-artifacts@v4 | ||
| name: ${{ inputs.artifact-prefix }}_r${{ inputs.rocm-version }} | ||
| path: ./wheelhouse | ||
| - name: Build JAX docker image | ||
| run: | | ||
| python3 build/ci_build \ | ||
| --rocm-version ${{ inputs.rocm-version }} | ||
| build_dockers | ||
| - name: Authenticate to GitHub Container Registry | ||
| run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin | ||
| - name: Push docker images | ||
| env: | ||
| ROCM_VERSION: ${{ inputs.rocm-version }} | ||
| run: | | ||
| docker tag "jax-ubu22.rocm${ROCM_VERSION//.}" "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" | ||
| docker tag "jax-ubu24.rocm${ROCM_VERSION//.}" "ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" | ||
| docker push "ghcr.io/rocm/jax-ubu22.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" | ||
| docker push "ghcr.io/rocm/jax-ubu24.rocm${ROCM_VERSION//.}:${GITHUB_REF_NAME}" | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,44 @@ | ||
| name: Build wheels | ||
|
|
||
| on: | ||
| workflow_call: | ||
| inputs: | ||
| python-versions: | ||
| required: true | ||
| type: string | ||
| rocm-version: | ||
| required: true | ||
| type: string | ||
| artifact-prefix: | ||
| default: 'plugin_wheels' | ||
|
|
||
| build-plugin-docker: | ||
| runs-on: mi-250 | ||
| 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 "./:/rocm-jax" ubuntu /bin/bash -c "chown -R $UID /rocm-jax/* || true" | ||
| # Remove any old work directories from this machine | ||
| rm -rf * || true | ||
| ls | ||
| - name: Print system info | ||
| run: | | ||
| whoami | ||
| printenv | ||
| df -h | ||
| - uses: actions/checkout@v4 | ||
| - name: Build plugin wheels | ||
| run: | | ||
| python3 build/ci_build \ | ||
| --compiler clang \ | ||
| --python-versions ${{ inputs.python-versions }} \ | ||
| --rocm-version ${{ inputs.rocm-version }} \ | ||
| dist_wheels | ||
| - name: Archive plugin wheels | ||
| uses: actions/upload-artifact@v4 | ||
| with: | ||
| name: ${{ inputs.artifact-prefix }}_r${{ env.ROCM_VERSION }} | ||
| path: ./wheelhouse/*.whl | ||
|
|
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What version of jax does it checkout? rocm-main?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Tip of upstream main. We need to hammer out how we handle tests though. It'd be nice if we were just able to submit broken tests (where it's the test case's fault) to upstream if they merge them quickly. And we talked about doing a skip list somewhere, but I haven't actually made a go at implementing that yet. One option is to just put stuff in the skip list until upstream fixes it. We could also rely on rocm-main in rocm/jax, but then we need to keep up with our sync stuff.