Fix jax plugin build (#1226)

## 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>

Author: charleshofer

.github/workflows/build_linux_jax_wheels.yml

@@ -26,12 +26,18 @@ on:
   workflow_dispatch:
     inputs:
       amdgpu_family:
-        required: true
-        type: string
+        type: choice
+        options:
+          - gfx110X-dgpu
+          - gfx1151
+          - gfx120X-all
+          - gfx94X-dcgpu
+          - gfx950-dcgpu
+        default: gfx94X-dcgpu
       python_versions:
         required: true
         type: string
-        default:
+        default: "3.12"
       release_type:
         description: The type of release to build ("nightly", or "dev")
         required: true
@@ -60,7 +66,7 @@ jobs:
     name: Build Linux JAX Wheels | ${{ inputs.amdgpu_family }} | Python ${{ inputs.python_version }}
     runs-on: ${{ github.repository_owner == 'ROCm' && 'azure-linux-scale-rocm' || 'ubuntu-24.04' }}
     env:
-      PACKAGE_DIST_DIR: ${{ github.workspace }}/jax_rocm_plugin/wheelhouse
+      PACKAGE_DIST_DIR: ${{ github.workspace }}/jax/jax_rocm_plugin/wheelhouse
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
     steps:
       - name: Checkout TheRock
@@ -71,23 +77,35 @@ jobs:
         with:
           path: jax
           repository: rocm/rocm-jax
-          ref: ${{ inputs.jax_ref }}
+          ref: ${{ matrix.jax_ref }}
 
       - name: Configure Git Identity
         run: |
           git config --global user.name "therockbot"
           git config --global user.email "therockbot@amd.com"
 
+      - name: "Setting up Python"
+        uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
+        with:
+          python-version: ${{ inputs.python_versions }}
+
       - name: Build JAX Wheels
+        env:
+          ROCM_VERSION: ${{ inputs.rocm_version }}
         run: |
-          pushd rocm-jax
+          ls -lah
+          pushd jax
           python3 build/ci_build \
             --compiler=clang \
             --python-versions="${{ inputs.python_versions }}" \
-            --rocm-version="${{ inputs.rocm_version }}" \
+            --rocm-version="${ROCM_VERSION:0:5}" \
             --therock-path="${{ inputs.tar_url }}" \
             dist_wheels
 
+      - name: Install AWS CLI
+        if: always()
+        run: bash ./dockerfiles/install_awscli.sh
+
       - name: Configure AWS Credentials
         if: always()
         uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4.3.1
@@ -104,5 +122,7 @@ 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 }}
+          python3 -m venv .venv
+          source .venv/bin/activate
+          pip3 install boto3 packaging
+          python3 ./build_tools/third_party/s3_management/manage.py ${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}

.github/workflows/release_portable_linux_packages.yml

@@ -252,6 +252,24 @@ jobs:
               "ref": "${{ inputs.ref || '' }}"
             }
 
+      - name: URL-encode .tar URL
+        id: url-encode-tar
+        run: python -c "from urllib.parse import quote; print('tar_url=https://therock-${{ env.RELEASE_TYPE }}-tarball.s3.amazonaws.com/' + quote('therock-dist-linux-${{ matrix.target_bundle.amdgpu_family }}${{ inputs.package_suffix }}-${{ needs.setup_metadata.outputs.version }}.tar.gz'))" >> ${GITHUB_OUTPUT}
+
+      - name: Trigger build JAX wheels
+        if: ${{ github.repository_owner == 'ROCm' }}
+        uses: benc-uk/workflow-dispatch@e2e5e9a103e331dad343f381a29e654aea3cf8fc # v1.2.4
+        with:
+          workflow: build_linux_jax_wheels.yml
+          inputs: |
+            { "amdgpu_family": "${{ matrix.target_bundle.amdgpu_family }}",
+              "python_versions": "3.12",
+              "release_type": "${{ env.RELEASE_TYPE }}",
+              "s3_subdir": "${{ env.S3_STAGING_SUBDIR }}",
+              "rocm_version": "${{ needs.setup_metadata.outputs.version }}",
+              "tar_url": "${{ steps.url-encode-tar.outputs.tar_url }}"
+            }
+
       - name: Save cache
         uses: actions/cache/save@0400d5f644dc74513175e3cd8d07132dd4860809 # v4.2.4
         if: ${{ !cancelled() }}