[Testing do not merge] Add ROCm wheel build and test pipeline to continuous CI#723
Draft
alekstheod wants to merge 58 commits intoamd-mainfrom
Draft
[Testing do not merge] Add ROCm wheel build and test pipeline to continuous CI#723alekstheod wants to merge 58 commits intoamd-mainfrom
alekstheod wants to merge 58 commits intoamd-mainfrom
Conversation
alekstheod
commented
Mar 4, 2026
alekstheod
commented
- Add jax-rocm-plugin and jax-rocm-pjrt to allowed artifacts in build_artifacts.sh with ROCm version flag passthrough.
- Create build_rocm_artifacts.yml reusable workflow that builds ROCm wheels in a ROCm container and uploads them to S3 via OIDC.
- Extend wheel_tests_continuous.yml with build-rocm-artifacts, run-pytest-rocm, and run-bazel-test-rocm jobs.
…tignore (#563) When jaxlib was built in debug more, an assertion in LLVM code that lazy-loads VHLO dialect could fire, since the code path could execute in a multi-threaded environment, and LLVM dialect repositories aren't thread safe to modify. This patch applies the same changes that upstream makes to fix this: jax-ml@48c8762 (this includes disabling a call to `jax_mlir_ext.enter_multi_threaded_execution(context)` in `mlir.py`. Presumably, the whole functionality related to `enter_multi_threaded_execution()` multithreaded checks isn't ready yet, and it was prematurely rolled into the production code. Manual testing
(forgot this skip in the previous PR)
Co-authored-by: Daniel Suo <danielsuo@gmail.com> Co-authored-by: Jake VanderPlas <jakevdp@google.com>
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
4670ac2 to
f2b145c
Compare
alekstheod
changed the title
alekstheod
changed the title
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
2 times, most recently
from
173cd80 to
ae3fc9f
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
ae3fc9f to
4a37416
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
8b729bb to
250595b
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
80d8ed5 to
76ac8c0
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
f172dc1 to
3965854
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
ab6bb69 to
3ea5e36
Compare
…build_wheels_pipeline_rocm_internal_v2
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
3ea5e36 to
a4958bb
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
6289d0f to
bc61e1e
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
bc61e1e to
9a3bb3b
Compare
alekstheod
force-pushed
the
implement_periodic_build_wheels_pipeline_rocm_internal_v2
branch
from
714ba20 to
b4f4b37
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
10 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.