Add workflow for JAX wheel builds by charleshofer · Pull Request #1033 · ROCm/TheRock

charleshofer · 2025-07-15T15:48:25Z

Adds a workflow for building JAX with release tarballs

charleshofer · 2025-07-15T16:24:17Z

Need to merge #1033 before this one

marbre · 2025-07-15T17:21:28Z

Need to merge #1033 before this one

This one is #1033 :)

charleshofer · 2025-07-15T18:37:39Z

Sorry. Need to merge ROCm/rocm-jax#58 first.

gabeweisz · 2025-07-17T19:11:34Z

@ScottTodd - any feedback on this?

ScottTodd · 2025-07-17T19:33:25Z

@@ -0,0 +1,117 @@
+name: Build Linux PyTorch Wheels


Update workflow name to match the file name

Suggested change

name: Build Linux PyTorch Wheels

name: Build Linux JAX Wheels

Also, are these builds "portable", e.g. by building under a manylinux docker container? I recently renamed some workflows to show that more clearly: #1023. It doesn't look like it from a read of https://github.com/ROCm/rocm-jax/blob/master/build/ci_build, so are these Linux packages only runnable on Ubuntu (or whichever distro is used for the build)?

Yes, these are built under a manylinux wheel. For a number of reasons, our build scripts are a little oniony. But we create a manylinux image with ROCm here, and then do the build in that: https://github.com/ROCm/rocm-jax/blob/master/jax_rocm_plugin/build/rocm/ci_build#L57

ScottTodd · 2025-07-17T19:38:15Z

+      - name: Checkout TheRock
+        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          path: TheRock


I'd prefer to omit path for the workflow's own repository and instead add path as needed for other repositories like rocm/rocm-jax above.

I flipped this around. I just want to point out that the only thing we use from checked-out TheRock is the ./build_tools/third_party/s3_management/manage.py. We use the tarball to actually do the build.

marbre · 2025-07-17T20:58:14Z

+      - name: Build JAX Wheels
+        run: |
+          cd rocm-jax
+          python3 build/ci_build \


It might make sense to move this script over to TheRock, similar to what we have for PyTorch. Wdyt?

No. This is the script that all CI teams use to build JAX. TheRock aren't the only users. DevOps team uses this when they build stuff against regular ROCm, and we use this in our CI on the JAX team.

It might make more sense then to checkout TheRock in rocm-jax, similar to what is done in rocm-libraries?

You guys wanted the opposite above, so I changed it. I can change it back if you want?

When this workflow checks out TheRock, all it uses is ./build_tools/third_party/s3_management/manage.py. The actual build uses the tarball that you give it.

You guys wanted the opposite above, so I changed it. I can change it back if you want?

I have not been involved in all prior discussions and might miss context here.

If it was agreed on earlier that this might be okay as is.

stellaraccident · 2025-07-25T18:47:56Z

How close are we to getting this landed? Any details we can defer until we get it in and some runs under our belt?

charleshofer · 2025-07-28T16:14:54Z

How close are we to getting this landed? Any details we can defer until we get it in and some runs under our belt?

I've addressed everyone's comments on this, I think. Just waiting on an approval.

ScottTodd · 2025-07-28T16:19:38Z

Ah, please click "re-request review" (if you can, otherwise just ping with @username) when you are ready for another round of review please. That helps PRs get added back to review dashboards like https://github.com/pulls/review-requested

charleshofer · 2025-07-28T16:24:15Z

Yeah, I'm not seeing the "re-request review" option where it normally is. But it's ready for another round @marbre.

ScottTodd · 2025-07-28T16:27:05Z

I clicked it for you:

It looked closer to this before:

ScottTodd · 2025-07-28T16:31:42Z

@marbre any concerns with using the same S3 bucket, subdir, etc. for JAX wheels?

We may need to add more packages to our index pages in the https://github.com/ROCm/TheRock/tree/main/build_tools/third_party/s3_management scripts.

If the build process outputs (or depends on) wheels other than JAX, is there a chance that they could conflict with what the PyTorch build produces?

@marbre any concerns with using the same S3 bucket, subdir, etc. for JAX wheels?

I am still undecided. We will move creating regular releases (and writing to the nightly releases bucket) to another repo but this does does not necessarily conflict with the proposed changes. What is a bit unfortunate is the fact that the build script is not in TheRock and that it gets called in a release workflow. Furthermore, changes to the script upstream it will not trigger a CI run in TheRock.

Furthermore, with #1110, we plan to land a gating mechanism soon. Untested wheels (and I consider the JAX wheels untested for now) will go to a staging subdir in the bucket to allow testing. If tests pass, they get copied to the final location. This is something we might want to have for JAX wheels as well or at least want it reworked after #1110 was landed.

We may need to add more packages to our index pages in the https://github.com/ROCm/TheRock/tree/main/build_tools/third_party/s3_management scripts.

Yes, without adding the package itself to manage.py, the index will not include the JAX wheel. Furthermore, it should be checked if additional packages must be added to update_dependencies.py.

If the build process outputs (or depends on) wheels other than JAX, is there a chance that they could conflict with what the PyTorch build produces?

Yes, running the workflow to see what it produces and what it pushes to the dev is indeed something we should check.

What is a bit unfortunate is the fact that the build script is not in TheRock and that it gets called in a release workflow. Furthermore, changes to the script upstream it will not trigger a CI run in TheRock.

Well it's the same for PyTorch. The scripts we have in TheRock are wrapping https://github.com/pytorch/pytorch/blob/main/setup.py. We're bootstrapping here, the bulk of support belongs upstream.

So, do I need to do anything here to switch the bucket or anything? Or is this good as-is for now?

I did add the packages that JAX needs to the list though. I think I've got everything in the right place.

ScottTodd · 2025-07-28T16:33:40Z

Have you run this workflow to test it? Can you share the logs?

If you wanted to test in this repository with the self-hosted runners, you could add a placeholder workflow first like with #828. This looks far enough along to land as-is, test, and then send new PRs for improvements.

@charleshofer can you try with the self-hosted runner here and post the log. Would really like to land this..

I've just run things through act to make sure the syntax was okay. I think it'd be easier to land this one rather than create another PR.

It's hard to review a change a workflow file without being able to see any logs of the workflow running, especially for a new workflow. I can make an attempt to parse the code and walk through it by eye, but there are details that I can't foresee from just the code, like how long each step takes, if there are permissions issues, if there are syntax errors, etc.

Will remember for future 🫡

ScottTodd · 2025-07-31T19:44:21Z

This input doesn't appear to be used yet. Safe to remove?

Yeah, I'll remove this

ScottTodd · 2025-07-31T19:46:45Z

It's hard to review a change a workflow file without being able to see any logs of the workflow running, especially for a new workflow. I can make an attempt to parse the code and walk through it by eye, but there are details that I can't foresee from just the code, like how long each step takes, if there are permissions issues, if there are syntax errors, etc.

ScottTodd · 2025-07-31T19:49:08Z

Oh interesting, these package names encode the ROCm version? That isn't part of a version identifier or part of the index URL that the packages come from? As we rework packaging for PyTorch too, let's keep an eye on this.

Co-authored-by: Marius Brehler <marius.brehler@amd.com>

## Motivation Fixes problems with the JAX CI workflow that was created in #1033 ## Technical Details Fixes minor problems that make the workflow crash, and ensures some command-line tools that we need are installed. ## Test Plan Make sure the workflow works when run with Actions: https://github.com/ROCm/TheRock/actions/workflows/build_linux_jax_wheels.yml --------- Co-authored-by: Scott Todd <scott.todd0@gmail.com>

charleshofer requested a review from gabeweisz July 15, 2025 15:48

github-project-automation Bot added this to TheRock Triage Jul 15, 2025

github-project-automation Bot moved this to TODO in TheRock Triage Jul 15, 2025

ScottTodd reviewed Jul 17, 2025

View reviewed changes

marbre reviewed Jul 17, 2025

View reviewed changes

Comment thread .github/workflows/build_linux_jax_wheels.yml Outdated

Comment thread .github/workflows/build_linux_jax_wheels.yml Outdated

charleshofer marked this pull request as ready for review July 25, 2025 18:15

charleshofer force-pushed the add-jax-ci branch from 4ae119b to c2eb9d8 Compare July 25, 2025 18:15

charleshofer force-pushed the add-jax-ci branch from 49724d2 to 31c6599 Compare July 28, 2025 16:12

ScottTodd requested review from ScottTodd and marbre July 28, 2025 16:18

ScottTodd reviewed Jul 28, 2025

View reviewed changes

charleshofer force-pushed the add-jax-ci branch from c845b0b to 9161b73 Compare July 31, 2025 14:36

charleshofer requested review from ScottTodd and jayfurmanek July 31, 2025 17:15

charleshofer force-pushed the add-jax-ci branch from 5e571f4 to 04a402f Compare July 31, 2025 19:00

ScottTodd approved these changes Jul 31, 2025

View reviewed changes

marbre reviewed Jul 31, 2025

View reviewed changes

Comment thread build_tools/third_party/s3_management/update_dependencies.py Outdated

charleshofer added 3 commits August 1, 2025 16:31

Add workflow for JAX wheel builds

a6cc99a

Address review comments

d9cb083

Fix TheRock path for rebuilding Pytyhon package index

a2d5484

charleshofer and others added 4 commits August 1, 2025 16:31

Remove newline at end

07c66fd

Add packages to management scripts

ac31b24

Update build_tools/third_party/s3_management/update_dependencies.py

9a9f068

Co-authored-by: Marius Brehler <marius.brehler@amd.com>

Remove unused workflow params

3dcfa73

charleshofer force-pushed the add-jax-ci branch from f2bd6dd to 3dcfa73 Compare August 1, 2025 16:32

charleshofer merged commit 53a722f into main Aug 1, 2025
1 check passed

charleshofer deleted the add-jax-ci branch August 1, 2025 16:33

github-project-automation Bot moved this from TODO to Done in TheRock Triage Aug 1, 2025

charleshofer mentioned this pull request Aug 11, 2025

Fix jax plugin build #1226

Merged

	name: Build Linux PyTorch Wheels
	name: Build Linux JAX Wheels

Conversation

charleshofer commented Jul 15, 2025

Uh oh!

charleshofer commented Jul 15, 2025

Uh oh!

marbre commented Jul 15, 2025

Uh oh!

charleshofer commented Jul 15, 2025

Uh oh!

gabeweisz commented Jul 17, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

stellaraccident commented Jul 25, 2025

Uh oh!

charleshofer commented Jul 28, 2025

Uh oh!

ScottTodd commented Jul 28, 2025

Uh oh!

charleshofer commented Jul 28, 2025

Uh oh!

ScottTodd commented Jul 28, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants