Upstream ci download jax rocm wheels#592
Draft
mminutoli wants to merge 6 commits intoupstream_rocm_cifrom
Draft
Conversation
mminutoli
force-pushed
the
upstream_ci_download_jax_rocm_wheels
branch
9 times, most recently
from
e046e80 to
3787aea
Compare
mminutoli
force-pushed
the
upstream_ci_download_jax_rocm_wheels
branch
2 times, most recently
from
00bddc8 to
c2198ba
Compare
Add standard JAX CI components to run_pytest_rocm.sh: bash header with license, environment setup (source default.env and install_wheels_locally.sh), diagnostic output (pip freeze, device check, rocm-smi), and proper pytest invocation using $JAXCI_PYTHON. Remove CSV/HTML reporting and log redirection for initial implementation, marked with TODOs for future consideration.
Extend wheel detection to support ROCm packages from both PyPI (jax-rocm*-pjrt, jax-rocm*-plugin) and GitHub releases (jax_rocm*_pjrt, jax_rocm*_plugin, jaxlib-*+rocm*). This enables the script to install ROCm wheels alongside existing CUDA wheel support.
Create GitHub Actions composite action to download JAX and ROCm wheels. Supports downloading from PyPI (jaxlib and jax-rocm plugins) or GitHub releases (custom jaxlib with ROCm plugins), with ROCm version mapping and platform-specific filtering. Uses gh CLI for GitHub release downloads and pip for PyPI packages.
Create reusable workflow for ROCm pytest execution, following pytest_cuda.yml pattern. Runs in ROCm team's GHCR containers with GPU device access, supports configurable JAX versions (default 0.8.0), and uses custom jaxlib from GitHub releases or PyPI. Installs GitHub CLI for release downloads and executes run_pytest_rocm.sh for test execution.
Integrates the pytest_rocm.yml workflow into the main CI pipeline for nightly and release testing. The workflow tests ROCm functionality across multiple GPU configurations (1/4/8 GPUs), Python versions (3.11, 3.12), and Ubuntu versions (22, 24) with ROCm 7.1.1. Also adds documentation for ROCm pytest porting plan.
mminutoli
force-pushed
the
upstream_ci_download_jax_rocm_wheels
branch
from
c2198ba to
0e10202
Compare
mminutoli
force-pushed
the
upstream_rocm_ci
branch
3 times, most recently
from
e8ab9de to
76386e8
Compare
mminutoli
force-pushed
the
upstream_rocm_ci
branch
3 times, most recently
from
04927a5 to
f9d2e1b
Compare
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist