Conversation
Upstream CI has stopped publishing jax/jaxlib nightly wheels to GCS (gs://jax-nightly-artifacts/latest), leaving the ROCm nightly test legs running against missing/stale base wheels. On the ROCm fork, build the base jax and jaxlib wheels on the same amd-do-linux.jax.cpu runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, so the ROCm test legs consume freshly built wheels instead of the broken GCS bucket. Upstream (jax-ml) behavior is unchanged. - ci/build_rocm_artifacts.sh: allow "jax"/"jaxlib" artifacts; skip auditwheel for the pure-python jax wheel. - wheel_tests_nightly_release.yml: new fork-only build-rocm-jax-jaxlib job; wire it into the ROCm test legs and select S3 as the jax/jaxlib source on the fork. - pytest_rocm.yml / bazel_rocm.yml: new optional jax-jaxlib-source input (default gcs), forwarded to the download action. - download-jax-rocm-wheels: when source is s3, fetch jax/jaxlib from the CloudFront listing alongside the ROCm wheels instead of GCS.
This adds 8028c6c78ea42b7184be31bcbf9d1fe40809bde1 on top of upstream's pin.
Ensure GitHub Actions workspace paths are available inside ROCm container jobs by mounting /home/runner/_work in pytest_rocm and bazel_rocm workflows. Co-authored-by: Cursor <cursoragent@cursor.com>
2 tasks
mminutoli
force-pushed
the
rocm-jaxlib-v0.11.1-alpha
branch
9 times, most recently
from
August 17, 2026 18:27
4f9521c to
a94e1bb
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
3 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.
Summary
Backport of #840 to
rocm-jaxlib-v0.11.1-alpha.Adds
/home/runner/_work:/home/runner/_workvolume mounts topytest_rocm.ymlandbazel_rocm.ymlso gfx950.8 container jobs can initialize without runner communication failures.Validation
Test plan
CI - Wheel Tests (Nightly/Release)on release branch