Fix jax plugin build by charleshofer · Pull Request #1226 · ROCm/TheRock

charleshofer · 2025-08-11T15:56:42Z

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

marbre

Drive-by: I noticed jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl and jax_rocm7_plugin-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl got uploaded to v2/gfx942 in the dev bucket but they must go into v2/gfx94X-dcgpu instead.

charleshofer · 2025-08-19T15:32:15Z

Drive-by: I noticed jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl and jax_rocm7_plugin-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl got uploaded to v2/gfx942 in the dev bucket but they must go into v2/gfx94X-dcgpu instead.

Gotcha. Looks like the amdgpu_family gets used in two different ways in the workflow. Once as an input to the JAX plugin build, and another time as an input to the upload script. I assume it'd be bad to hardcode the upload to v2/gfx94X-dcgpu, right? Because we might want to upload to gfx95X. The JAX build requires specific GFX numbers. It can't take something like gfx94X. Maybe we just make two inputs to the workflow? One for the build and one for the bucket?

marbre · 2025-08-19T23:24:07Z

Drive-by: I noticed jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl and jax_rocm7_plugin-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl got uploaded to v2/gfx942 in the dev bucket but they must go into v2/gfx94X-dcgpu instead.

Gotcha. Looks like the amdgpu_family gets used in two different ways in the workflow. Once as an input to the JAX plugin build, and another time as an input to the upload script. I assume it'd be bad to hardcode the upload to v2/gfx94X-dcgpu, right? Because we might want to upload to gfx95X. The JAX build requires specific GFX numbers. It can't take something like gfx94X. Maybe we just make two inputs to the workflow? One for the build and one for the bucket?

Other workflows incorporate scripts that use https://github.com/ROCm/TheRock/blob/main/build_tools/github_actions/amdgpu_family_matrix.py. @geomin12 has more context here but normally we avoid to have to inputs.

geomin12 · 2025-08-19T23:53:18Z

Drive-by: I noticed jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl and jax_rocm7_plugin-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl got uploaded to v2/gfx942 in the dev bucket but they must go into v2/gfx94X-dcgpu instead.

Gotcha. Looks like the amdgpu_family gets used in two different ways in the workflow. Once as an input to the JAX plugin build, and another time as an input to the upload script. I assume it'd be bad to hardcode the upload to v2/gfx94X-dcgpu, right? Because we might want to upload to gfx95X. The JAX build requires specific GFX numbers. It can't take something like gfx94X. Maybe we just make two inputs to the workflow? One for the build and one for the bucket?

hi! yes, for inputs, we figured that the abbreviation would be easier for workflow_dispatch triggers. in general, we typically have a Setup step/job that determines which targets to run, where we get these from amdgpu_family_matrix.py. We use the gfx94X as a key, then get the corresponding full family name (gfx94X -> gfx94X-dcgpu)

I would recommend creating a setup job/step, use GitHub outputs (setup.yml as an example) to pass this information to the jax plugin segment, so it can be a bit more dynamic

amd-chrissosa · 2025-08-26T07:05:08Z

Any progress needed here? I understand Jax was critical to enable a specific set of users

charleshofer · 2025-08-28T22:38:02Z

Do you guys want me to set this up to run on a schedule? Nightly? Weekly? @gabeweisz

gabeweisz · 2025-08-28T23:36:09Z

Do you guys want me to set this up to run on a schedule? Nightly? Weekly? @gabeweisz

Nightly please

charleshofer · 2025-09-04T15:13:24Z

@@ -247,6 +247,20 @@ jobs:
              "rocm_version": "${{ needs.setup_metadata.outputs.version }}"
            }

+      - name: Trigger build JAX wheels


@marbre I think this is the right place to stick the trigger for the nightly build. Is there a way to get the URL of the latest TheRock .tar release though? I didn't see a simple way to do it through this workflow.

The base URL right now is always https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/ (no CloudFront distribution planned). ROCm version is known by ${{ needs.setup_metadata.outputs.version }} and GPU family via ${{ matrix.target_bundle.amdgpu_family }}. Thus the URL should be

https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}-${{ needs.setup_metadata.outputs.version }}.tar.gz

See https://github.com/ROCm/TheRock/blob/bb4372b9b502177915ade4d0d9f97397283212fd/.github/workflows/release_portable_linux_packages.yml#L117C53-L117C193 and maybe even better use

therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}${{ inputs.package_suffix }}-${{ needs.setup_metadata.outputs.version }}.tar.gz

for the filename to make it more robust.

If I understand, that's the tarball that will sit on the worker's local filesystem, correct? I'm not super familiar with that specific workflow dispatch action, but is it guaranteed to run the wheel build workflow on the same runner as the workflow that called it? If not, we should keep using the URL

Oh I was only referring to the where the variable is composed. The tarbal behind is uploaded to S3 here: https://github.com/ROCm/TheRock/blob/bb4372b9b502177915ade4d0d9f97397283212fd/.github/workflows/release_portable_linux_packages.yml#L210

Wanted to point out that the file name can have an the package_suffix in it if set. This does not apply to scheduled builds by default but can be the case for manually triggered builds.

Ahh, gotcha. Also, right now this points to

https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}${{ inputs.package_suffix }}-${{ needs.setup_metadata.outputs.version }}.tar.gz

Could we also have

https://therock-dev-tarball.s3.us-east-2.amazonaws.com/therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}${{ inputs.package_suffix }}-${{ needs.setup_metadata.outputs.version }}.tar.gz

for dev builds? (Notice the change in the first part of the URL)

I used to encode this as build/release type variable so that can create either or URL.

geomin12 · 2025-09-04T16:53:54Z

+      - name: Install AWS CLI
+        if: always()
+        run: |
+          curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
+          unzip awscliv2.zip
+          ./aws/install -i ./aws-cli-files -b ./aws-bin
+          sudo ./aws/install
+


if you use this container (

TheRock/.github/workflows/build_linux_packages.yml

Line 42 in 7dff9ee

image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:543ba2609de3571d2c64f3872e5f1af42fdfa90d074a7baccb1db120c9514be2

), it comes already installed with AWS! I would recommend replacing this and just using container :)

We cannot use the container that has AWS pre-installed. The build/ci_build script that does the JAX wheel build creates a docker image and does the actual build inside of that so that we're manylinux compliant.

in that case, we do have this script (https://github.com/ROCm/TheRock/blob/main/dockerfiles/install_awscli.sh) that does the same thing here! less duplication!

I'll give that a shot

geomin12 · 2025-09-04T16:58:34Z

@@ -104,5 +119,6 @@ jobs:
      - name: (Re-)Generate Python package release index
        if: ${{ github.repository_owner == 'ROCm' }}
        run: |
-          pip install boto3 packaging
-          python ./build_tools/third_party/s3_management/manage.py ${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}
+          sudo apt install python3-venv python3-pip -y


can be removed if container is used!

you can also use:

TheRock/.github/actions/setup_test_environment/action.yml

Lines 23 to 26 in 7dff9ee

- name: "Setting up Python"

uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0

with:

python-version: 3.11

to help with python versioning!

We can't use the container, but I'll switch it to using the actions step

marbre · 2025-09-04T17:07:55Z

@@ -247,6 +247,20 @@ jobs:
              "rocm_version": "${{ needs.setup_metadata.outputs.version }}"
            }

+      - name: Trigger build JAX wheels


The base URL right now is always https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/ (no CloudFront distribution planned). ROCm version is known by ${{ needs.setup_metadata.outputs.version }} and GPU family via ${{ matrix.target_bundle.amdgpu_family }}. Thus the URL should be

https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}-${{ needs.setup_metadata.outputs.version }}.tar.gz

marbre · 2025-09-04T17:08:13Z

+              "release_type": "${{ env.RELEASE_TYPE }}",
+              "s3_subdir": "${{ env.S3_SUBDIR }}",
+              "rocm_version": "${{ needs.setup_metadata.outputs.version }}",
+              "tag_url": "${{ ??? }}"


CI fails with invalid workflow on this.

geomin12 · 2025-09-04T19:45:11Z

+      - name: "Setting up Python"
+        uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
+        with:
+          python-version: 3.11


want to use the input here? inputs.python_versions

geomin12 · 2025-09-04T19:46:46Z

+              "release_type": "${{ env.RELEASE_TYPE }}",
+              "s3_subdir": "${{ env.S3_SUBDIR }}",
+              "rocm_version": "${{ needs.setup_metadata.outputs.version }}",
+              "tag_url": "https://therock-nightly-tarball.s3.us-east-2.amazonaws.com/therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}-${{ needs.setup_metadata.outputs.version }}.tar.gz"


nit: pretty sure you can remove the us-east-2 part

Yeah, should be possible as well.

geomin12

lgtm! although, can we test this by triggering release_portable_linux_packages using workflow_dispatch? just to make sure everything works?

charleshofer · 2025-09-04T23:30:17Z

I kicked off a nightly build with this branch: https://github.com/ROCm/TheRock/actions/runs/17479247009

Co-authored-by: Scott Todd <scott.todd0@gmail.com>

charleshofer · 2025-09-26T15:52:08Z

@marbre @ScottTodd I got this to work:
https://github.com/ROCm/TheRock/actions/runs/18017743344
https://github.com/ROCm/TheRock/actions/runs/18023230999

ScottTodd · 2025-09-26T16:10:12Z

      - name: Upload wheels to S3
        if: ${{ github.repository_owner == 'ROCm' }}
        run: |
          aws s3 cp ${{ env.PACKAGE_DIST_DIR }}/ s3://${{ env.S3_BUCKET_PY }}/${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}/ \
            --recursive --exclude "*" --include "*.whl"


How close are these wheels ready to be advertised? If not ready yet, should we upload them to s3_staging_subdir instead of s3_subdir? We started uploading pytorch wheels first to "staging" then we copy them out once tests complete.

Once ready, we should advertise JAX support here:

https://github.com/ROCm/TheRock?tab=readme-ov-file#features (change "JAX ... builds are in the works" language)

https://github.com/ROCm/TheRock?tab=readme-ov-file#nightly-release-status (with a badge for workflow status)

Add a section to https://github.com/ROCm/TheRock/blob/main/RELEASES.md for JAX like the pytorch sections

Maybe add dedicated docs for how to build JAX (or link to a repo with docs). Could maybe add a page like https://github.com/ROCm/TheRock/blob/main/external-builds/pytorch/README.md for JAX, with just a README but no code/scripts?

These are not ready. I don't have the automated testing in place yet to make sure that they're good. Will switch to using the staging directory.

Thanks. You might want to put a TODO in there to

Plumb through both the staging and "prod" URL

First upload to staging, then run tests, then copy to "prod" if tests passed (matching what we do for torch)

github-project-automation Bot added this to TheRock Triage Aug 11, 2025

github-project-automation Bot moved this to TODO in TheRock Triage Aug 11, 2025

marbre reviewed Aug 12, 2025

View reviewed changes

charleshofer force-pushed the fix-jax-plugin-build branch from 5e6ffc9 to 61b1565 Compare August 19, 2025 23:02

charleshofer force-pushed the fix-jax-plugin-build branch from 702870b to 8943ecd Compare August 28, 2025 21:14

charleshofer marked this pull request as ready for review August 28, 2025 21:40

charleshofer force-pushed the fix-jax-plugin-build branch from 96dce9d to 73095e2 Compare September 4, 2025 14:59

charleshofer commented Sep 4, 2025

View reviewed changes

charleshofer requested review from ScottTodd, geomin12 and marbre September 4, 2025 16:50

geomin12 reviewed Sep 4, 2025

View reviewed changes

Comment thread .github/workflows/build_linux_jax_wheels.yml

geomin12 reviewed Sep 4, 2025

View reviewed changes

Comment thread .github/workflows/build_linux_jax_wheels.yml

marbre reviewed Sep 4, 2025

View reviewed changes

charleshofer requested review from geomin12 and marbre September 4, 2025 19:09

geomin12 reviewed Sep 4, 2025

View reviewed changes

charleshofer force-pushed the fix-jax-plugin-build branch 2 times, most recently from beec084 to 107c6bf Compare September 4, 2025 21:13

charleshofer and others added 15 commits September 17, 2025 20:50

Use python3 instead of python

49853c7

Add setup step

be4584e

Fix amdgpu_family

7393636

Fix needs

116b748

Add JAX build to nightly

cbe5af3

Fix nightly tarball path and address review comments:

097aa8d

Fix python version

556be66

Add suffix for tarball URL

1e7f0fb

Use AWS install script

f0b81a6

Fix typo in release workflow

8a633c2

Remove region from URL

eee39f0

Fix ROCm version string for JAX build

d44c3f2

Update .github/workflows/release_portable_linux_packages.yml

256a470

Co-authored-by: Scott Todd <scott.todd0@gmail.com>

Encode URL for tar file

bc1911a

Fix tar URL assignment

d95caf7

charleshofer force-pushed the fix-jax-plugin-build branch from fa9b009 to d95caf7 Compare September 17, 2025 20:50

charleshofer added 5 commits September 18, 2025 16:18

Fix url quoting

6896c59

Fix outputs JSON

b1fbd2c

Remove quotes

a6e2f3e

Remove extra comma

8aeb56a

Fix tar url

2b13c45

charleshofer requested a review from ScottTodd September 26, 2025 15:52

ScottTodd reviewed Sep 26, 2025

View reviewed changes

Use staging directory

666bfd9

ScottTodd approved these changes Sep 26, 2025

View reviewed changes

charleshofer merged commit 341d943 into main Sep 29, 2025
29 of 32 checks passed

charleshofer deleted the fix-jax-plugin-build branch September 29, 2025 08:15

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

	- name: "Setting up Python"
	uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
	with:
	python-version: 3.11

Conversation

charleshofer commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Uh oh!

marbre left a comment

Choose a reason for hiding this comment

Uh oh!

charleshofer commented Aug 19, 2025

Uh oh!

marbre commented Aug 19, 2025

Uh oh!

geomin12 commented Aug 19, 2025

Uh oh!

amd-chrissosa commented Aug 26, 2025

Uh oh!

charleshofer commented Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gabeweisz commented Aug 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!

marbre Sep 4, 2025 • 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!

charleshofer Sep 4, 2025 • edited by marbre Loading Uh oh! There was an error while loading. Please reload this page.

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!

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!

geomin12 left a comment

Choose a reason for hiding this comment

Uh oh!

charleshofer commented Sep 4, 2025

Uh oh!

charleshofer commented Sep 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

Uh oh!

charleshofer commented Aug 11, 2025 •

edited

Loading

charleshofer commented Aug 28, 2025 •

edited

Loading

marbre Sep 4, 2025 •

edited

Loading

charleshofer Sep 4, 2025 •

edited by marbre

Loading

charleshofer commented Sep 26, 2025 •

edited

Loading