feat: (jax) Add Build JAX wheels in Multi-Arch CI by erman-gurses · Pull Request #6439 · ROCm/TheRock

erman-gurses · 2026-07-09T07:18:31Z

Progress #3878.

Summary

Adds JAX wheel build coverage to Multi-Arch CI.

This replaces the placeholder JAX CI workflow with a real manylinux JAX
wheel build flow, wires it into Linux Multi-Arch CI, and enables JAX
wheel builds by default alongside PyTorch.

Changes

Replaces the placeholder JAX CI reusable workflow with a real build flow:
- checks out TheRock, rocm-jax, and JAX
- builds the JAX manylinux image
- builds jax-rocm-plugin and jax-rocm-pjrt wheels
- extracts JAX wheel versions
- uploads built wheels to run-scoped CI Python artifacts
Adds build_jax to multi_arch_ci.yml:
- default: true
- push / pull_request: JAX builds run by default
- workflow_dispatch: JAX builds can be disabled manually
Wires build_jax_wheels into multi_arch_ci_linux.yml:
- depends on build_python_packages
- consumes build_python_packages.outputs.package_find_links_url
- uses jax_build_matrix from the generated Linux build config
Adds a CI-specific JAX matrix helper:
- CI uses only manylinux JAX refs
- release workflows keep the full JAX release matrix

CI vs release matrix behavior

The release JAX matrix still includes:

rocm-jaxlib-v0.9.1
rocm-jaxlib-v0.10.0
rocm-jaxlib-v0.10.2

The CI JAX matrix filters out the legacy/native rocm-jaxlib-v0.9.1
entry because the new CI workflow is manylinux-only.

Default CI currently generates:

rocm-jaxlib-v0.10.0
Python 3.12

This produces 1 JAX CI build cell.

Testing

Ran local Python checks:

python -m compileall build_tools/github_actions/configure_multi_arch_ci.py \
  build_tools/github_actions/configure_jax_release_matrix.py

CI Test: https://github.com/ROCm/TheRock/actions/runs/29396032944/job/87311830950?pr=6439

Follow up

Wire JAX wheel tests into Multi-Arch CI after the build lane is stable.
feat: (jax) Add Build JAX wheels in Multi-Arch CI #6439 (comment)

TODO

#6388

therock-pr-bot · 2026-07-09T07:20:17Z

✅ All Policy Checks Passed

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

🎉 All policy checks passed!

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

therock-pr-bot · 2026-07-09T07:20:18Z

🎉 All checks passed! This PR is ready for review.

…LOWS_CI_FILENAMES

marbre

Please see compare with PyTorch where we limit to only build for Python 3.12. We shouldn't run the full Python matrix on every CI build.

erman-gurses · 2026-07-09T19:05:26Z

Please see compare with PyTorch where we limit to only build for Python 3.12. We shouldn't run the full Python matrix on every CI build.

Thanks - that makes sense - just updated the code

marbre

Is it worth to build rocm-jaxlib-v0.10.0 and rocm-jaxlib-v0.10.2?

Furthermore, the CI description needs to be updated with regards to the Python versions:

Default CI currently generates:

* rocm-jaxlib-v0.10.0
* rocm-jaxlib-v0.10.2
* Python 3.11, 3.12, 3.13, and 3.14

marbre · 2026-07-09T21:09:02Z

+          docker build \
+            -t "${MANYLINUX_IMAGE_TAG}" \
+            --file=Dockerfile.jax-manylinux_2_28-therock \
+            --build-arg=THEROCK_INDEX_URL="${ROCM_PACKAGE_FIND_LINKS_URL}" \
+            --build-arg=THEROCK_VERSION="${THEROCK_VERSION}" \
+            --build-arg=GFX_ARCH="${GFX_ARCH}" \
+            --progress=plain \
+            .


As this workflow is intended to run on every CI run we should avoid building a docker container every time. This overhead might be acceptable as an intermediate solution for release builds but I think it is not for CI builds.

I see the point but I am not controlling that part - If you want I can keep only rocm-jaxlib-v0.10.0 until other solution comes from JAX Team or should we wait?

I am aware you're not in control of it but I would rather delay this PR until it was fixed upstream. Would switching to 0.10.0 let us avoid re-building the Docker container every time? Normally I would prefer to stick with the latest patch version which in our case seems to be 0.10.2.

Anyway, let's wait for @ScottTodd to weight in. I would rather push back to merge this as is.

0.10.0 is the most stable one - 0.10.2 has some issues from the upstream. If we chose one, I think 0.10.0 is the best candidate.
Would switching to 0.10.0 let us avoid re-building the Docker container every time? Just an idea, It will reduce the pressure on the runners

Sounds good to me to wait for Scott's input.

Well if it still requires Docker my push back remains. As 0.10.2 is a patch release I would also rather go with that version instead of 0.10.0. Let's wait for Scott to weight in.

@ScottTodd , If I gate this for the rocm-systems and rocm-libraries bumps only, would it be ok? Since my main purpose is to detect how JAX is affected by those bump PRs.

I'll take a more detailed look tomorrow (only got through ~50% of my PR review queue today). Some quick thoughts:

At least this depends on build_python_packages and not tarball builds. If it depended on tarball builds that would bloat the storage size needed for each CI run.

We should compare the job time between JAX and PyTorch and look for opportunities to optimize both. I see ~1h of build time for both with opportunities (ccache/sccache and other tricks) to get that down to ~20 minutes.

As for what actually happens during the build, I see that JAX is still building LLVM from source (we already build a copy in ROCm, ideally we wouldn't need another copy) and I see that we're still building a dockerfile every run (which pulls in uncontrolled dependencies).

If we build for releases already, we should at least have the ability to build on CI, even if only opt-in via a label or workflow_dispatch argument. I do like the idea of running these new jobs on at least bump PRs in TheRock, so we can catch if/when changes to rocm-systems and rocm-libraries break the JAX build (we've seen changes to RCCL break the PyTorch build).

So I think we can find a reasonable middle-ground for the CI system that gets us this much desired coverage, but there are still a few things to continue development on for the final shape that we'll all be happy with. Let's make sure that enough of that is tracked (and prioritized/assigned) so we don't let it slip for too long.

(will continue reviewing tomorrow)

Thanks, I opened an issue to track follow ups #6388, I will modify and fill out as they are determined.

Please add a TODO here with a link to an issue covering removing this dockerfile build. Our build containers should be frozen and updated infrequently, not build as part of every workflow run. This container in particular (https://github.com/ROCm/rocm-jax/blob/rocm-jax-infra/docker/manylinux/Dockerfile.jax-manylinux_2_28-therock) has multiple issues:

It builds LLVM from source (when we already build and distribute LLVM as part of ROCm)

It installs the AWS CLI from a domain we do not control

It likely has poor interactions with caches (docker layer caches are a poor substitute for ccache)

Dockerfiles setups like this are awkward for local development (including debugging, development, and reproducing what CI does)

We shouldn't need this at all, all we should need is to run in our existing https://github.com/ROCm/TheRock/blob/main/dockerfiles/build_manylinux_x86_64.Dockerfile image or a slightly modified image that also includes Bazel and any other build dependencies (for the entire job), install rocm into a venv, and then run the python build/build.py script.

ScottTodd

I'm okay with the extra time (~1h6m) that this will add to CI runs here in TheRock, but yes let's follow up later on some of the specifics to optimize both how long that takes and when it triggers (e.g. exclude jax builds for PRs that only modify pytorch code)

ScottTodd · 2026-07-14T18:20:49Z

+          docker build \
+            -t "${MANYLINUX_IMAGE_TAG}" \
+            --file=Dockerfile.jax-manylinux_2_28-therock \
+            --build-arg=THEROCK_INDEX_URL="${ROCM_PACKAGE_FIND_LINKS_URL}" \
+            --build-arg=THEROCK_VERSION="${THEROCK_VERSION}" \
+            --build-arg=GFX_ARCH="${GFX_ARCH}" \
+            --progress=plain \
+            .


Please add a TODO here with a link to an issue covering removing this dockerfile build. Our build containers should be frozen and updated infrequently, not build as part of every workflow run. This container in particular (https://github.com/ROCm/rocm-jax/blob/rocm-jax-infra/docker/manylinux/Dockerfile.jax-manylinux_2_28-therock) has multiple issues:

It builds LLVM from source (when we already build and distribute LLVM as part of ROCm)

It installs the AWS CLI from a domain we do not control

It likely has poor interactions with caches (docker layer caches are a poor substitute for ccache)

Dockerfiles setups like this are awkward for local development (including debugging, development, and reproducing what CI does)

We shouldn't need this at all, all we should need is to run in our existing https://github.com/ROCm/TheRock/blob/main/dockerfiles/build_manylinux_x86_64.Dockerfile image or a slightly modified image that also includes Bazel and any other build dependencies (for the entire job), install rocm into a venv, and then run the python build/build.py script.

feat: (jax) Add Build JAX wheels in Multi-Arch CI

e0a1f05

github-project-automation Bot added this to TheRock Triage Jul 9, 2026

github-project-automation Bot moved this to TODO in TheRock Triage Jul 9, 2026

therock-pr-bot Bot added the Not ready to Review PR has unresolved policy failures — reviews blocked label Jul 9, 2026

Add multi_arch_build_linux_jax_wheels_ci file name into _GITHUB_WORKF…

031550f

…LOWS_CI_FILENAMES

erman-gurses requested a review from WBobby July 9, 2026 07:27

Include JAX fields in Linux BuildConfig contract

4751f9e

therock-pr-bot Bot removed the Not ready to Review PR has unresolved policy failures — reviews blocked label Jul 9, 2026

Remove github.ref_name

f69fc01

marbre requested changes Jul 9, 2026

View reviewed changes

Limit Python matrix for only Python 3.12

c36e23f

erman-gurses requested a review from marbre July 9, 2026 19:05

marbre reviewed Jul 9, 2026

View reviewed changes

keep only rocm-jaxlib-v0.10.0 version

4bfebed

erman-gurses requested a review from marbre July 9, 2026 21:27

Adjust min matrix size to CI

3c1e222

marbre requested a review from ScottTodd July 10, 2026 10:20

ScottTodd reviewed Jul 14, 2026

View reviewed changes

erman-gurses and others added 5 commits July 15, 2026 04:03

Match PyTorch release-type matrix pattern

282210d

Refactor configure_jax_release_matrix.py

f492fc2

Add unit tests

a36e6fe

Set package dist dir in job env

69d11b0

Merge branch 'main' into users/erman-gurses/add-jax-to-ci-workflow

19837e1

erman-gurses added the ci:skip Skip all CI builds/tests for this PR label Jul 15, 2026

Solve conflict

33c156a

erman-gurses removed the ci:skip Skip all CI builds/tests for this PR label Jul 15, 2026

Add TODO 6388

24b0eeb

erman-gurses mentioned this pull request Jul 15, 2026

Replace per-run Dockerfile build with frozen CI image and optimize trigger policy #6388

Open

erman-gurses requested a review from ScottTodd July 15, 2026 06:34

Update TODO note

c4c61cb

ScottTodd approved these changes Jul 15, 2026

View reviewed changes

Uh oh!

Conversation

erman-gurses commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

CI vs release matrix behavior

Testing

Follow up

TODO

Uh oh!

therock-pr-bot Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Policy Checks Passed

Uh oh!

therock-pr-bot Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

marbre left a comment

Choose a reason for hiding this comment

Uh oh!

erman-gurses commented Jul 9, 2026

Uh oh!

marbre left a comment

Choose a reason for hiding this comment

Uh oh!

marbre Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

erman-gurses Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

marbre Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

erman-gurses Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

marbre Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

erman-gurses Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ScottTodd Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

erman-gurses Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ScottTodd Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

ScottTodd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ScottTodd Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

erman-gurses commented Jul 9, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 9, 2026 •

edited

Loading

therock-pr-bot Bot commented Jul 9, 2026 •

edited

Loading

erman-gurses Jul 9, 2026 •

edited

Loading

erman-gurses Jul 10, 2026 •

edited

Loading

erman-gurses Jul 13, 2026 •

edited

Loading

erman-gurses Jul 14, 2026 •

edited

Loading