[DRAFT] Add rocprofiler-sdk tests

.github/workflows/build_linux_jax_wheels.yml

@@ -42,10 +42,6 @@ on:
       ref:
         description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
         type: string
-      jax_ref:
-        description: "rocm-jax repository ref/branch to check out"
-        type: string
-        default: "rocm-jaxlib-v0.8.0-fixdevtar"
   workflow_dispatch:
     inputs:
       amdgpu_family:
@@ -97,14 +93,10 @@ on:
         description: CloudFront base URL pointing to staging Python index
         type: string
         default: "https://rocm.devreleases.amd.com/v2-staging"
-      ref:
-        description: TheRock repository ref/branch to check out
-        type: string
-        default: ""
       jax_ref:
         description: rocm-jax repository ref/branch to check out
         type: string
-        default: "rocm-jaxlib-v0.8.0-fixdevtar"
+        default: rocm-jaxlib-v0.8.0
 
 permissions:
   id-token: write
@@ -116,19 +108,15 @@ jobs:
   build_jax_wheels:
     strategy:
       matrix:
-        jax_ref: [rocm-jaxlib-v0.8.0-fixdevtar]
+        jax_ref: [rocm-jaxlib-v0.8.0]
     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/jax_rocm_plugin/wheelhouse
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
     outputs:
       cp_version: ${{ env.cp_version }}
-      jax_version: ${{ steps.write_jax_versions.outputs.jax_version }}
-      jaxlib_version: ${{ steps.write_jax_versions.outputs.jaxlib_version }}
-      jax_plugin_version: ${{ steps.write_jax_versions.outputs.jax_plugin_version }}
-      jax_pjrt_version: ${{ steps.write_jax_versions.outputs.jax_pjrt_version }}
-      jax_ref: ${{ matrix.jax_ref }}
+      jax_version: ${{ steps.extract_jax_version.outputs.jax_version }}
     steps:
       - name: Checkout TheRock
         uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
@@ -168,19 +156,25 @@ jobs:
             --therock-path="${{ inputs.tar_url }}" \
             dist_wheels
 
-      - name: Extract JAX versions from built wheels
-        id: write_jax_versions
+      - name: Extract JAX version
+        id: extract_jax_version
         run: |
-          python ./build_tools/github_actions/write_jax_versions.py \
-            --dist-dir ${{ env.PACKAGE_DIST_DIR }}
+            # Extract JAX version from requirements.txt (e.g., "jax==0.8.0")
+            # Remove all whitespace from requirements.txt to simplify parsing
+            # Search for lines starting with "jax==" or "jaxlib==" followed by version (excluding comments)
+            # Extract the version number by splitting on '=' and taking the 3rd field
+            # [^#]+ matches one or more characters that are NOT '#', ensuring we stop before any inline comments
+            JAX_VERSION=$(tr -d ' ' < jax/build/requirements.txt \
+            | grep -E '^(jax|jaxlib)==[^#]+' | head -n1 | cut -d'=' -f3)
+            echo "jax_version=$JAX_VERSION" >> "$GITHUB_OUTPUT"
 
       - name: Install AWS CLI
         if: always()
         run: bash ./dockerfiles/install_awscli.sh
 
       - name: Configure AWS Credentials
         if: always()
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708 # v5.1.1
         with:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}
@@ -232,15 +226,14 @@ jobs:
     with:
       amdgpu_family: ${{ inputs.amdgpu_family }}
       release_type: ${{ inputs.release_type }}
+      s3_subdir: ${{ inputs.s3_subdir }}
       package_index_url: ${{ inputs.cloudfront_staging_url }}
       rocm_version: ${{ inputs.rocm_version }}
       tar_url: ${{ inputs.tar_url }}
       python_version: ${{ inputs.python_version }}
       repository: ${{ inputs.repository || github.repository }}
       ref: ${{ inputs.ref || '' }}
-      jax_ref: ${{ needs.build_jax_wheels.outputs.jax_ref }}
-      jax_version: ${{ needs.build_jax_wheels.outputs.jax_version }}
-      jaxlib_version: ${{ needs.build_jax_wheels.outputs.jaxlib_version }}
+      jax_ref: ${{ inputs.jax_ref }}
       test_runs_on: ${{ needs.generate_target_to_run.outputs.test_runs_on }}
 
   upload_jax_wheels:
@@ -250,9 +243,8 @@ jobs:
     runs-on: ubuntu-24.04
     env:
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
-      JAXLIB_VERSION: "${{ needs.build_jax_wheels.outputs.jaxlib_version }}"
-      JAX_PLUGIN_VERSION: "${{ needs.build_jax_wheels.outputs.jax_plugin_version }}"
-      JAX_PJRT_VERSION: "${{ needs.build_jax_wheels.outputs.jax_pjrt_version }}"
+      JAX_VERSION: "${{ needs.build_jax_wheels.outputs.jax_version }}"
+      ROCM_VERSION: "${{ inputs.rocm_version }}"
       CP_VERSION: "${{ needs.build_jax_wheels.outputs.cp_version }}"
 
     steps:
@@ -264,7 +256,7 @@ jobs:
 
       - name: Configure AWS Credentials
         if: always()
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708 # v5.1.1
         with:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}
@@ -281,17 +273,14 @@ jobs:
         if: ${{ env.upload == 'true' }}
         run: |
           echo "Copying exact tested wheels to release S3 bucket..."
-          echo "  JAXLIB_VERSION=${JAXLIB_VERSION}"
-          echo "  JAX_PLUGIN_VERSION=${JAX_PLUGIN_VERSION}"
-          echo "  JAX_PJRT_VERSION=${JAX_PJRT_VERSION}"
           aws s3 cp \
             s3://${S3_BUCKET_PY}/${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}/ \
             s3://${S3_BUCKET_PY}/${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}/ \
             --recursive \
             --exclude "*" \
-            --include "jaxlib-${JAXLIB_VERSION}-${CP_VERSION}-manylinux_2_27_x86_64.whl" \
-            --include "jax_rocm7_plugin-${JAX_PLUGIN_VERSION}-${CP_VERSION}-manylinux_2_28_x86_64.whl" \
-            --include "jax_rocm7_pjrt-${JAX_PJRT_VERSION}-py3-none-manylinux_2_28_x86_64.whl"
+            --include "jaxlib-${JAX_VERSION}+rocm${ROCM_VERSION}-${CP_VERSION}-manylinux_2_27_x86_64.whl" \
+            --include "jax_rocm7_plugin-${JAX_VERSION}+rocm${ROCM_VERSION}-${CP_VERSION}-manylinux_2_28_x86_64.whl" \
+            --include "jax_rocm7_pjrt-${JAX_VERSION}+rocm${ROCM_VERSION}-py3-none-manylinux_2_28_x86_64.whl"
 
       - name: (Re-)Generate Python package release index
         if: ${{ env.upload == 'true' }}

.github/workflows/build_native_linux_packages.yml

@@ -69,7 +69,7 @@ jobs:
       fail-fast: false
     runs-on: ${{ github.repository_owner == 'ROCm' && 'azure-linux-scale-rocm' || 'ubuntu-24.04' }}
     env:
-      BUILD_IMAGE: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:d6ae5712a9c7e8b88281d021e907b312cd8a26295b95690baef3e8dde4805858
+      BUILD_IMAGE: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:db2b63f938941dde2abc80b734e64b45b9995a282896d513a0f3525d4591d6cb
       ARTIFACT_RUN_ID: ${{ inputs.artifact_run_id || github.run_id }}
       PACKAGE_SUFFIX: ${{ inputs.package_suffix != '' && inputs.package_suffix || '' }}
       OUTPUT_DIR: ${{ github.workspace }}/output
@@ -120,7 +120,7 @@ jobs:
         run: bash ./dockerfiles/install_awscli.sh
 
       - name: Configure AWS Credentials for non-forked repos
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708 # v5.1.1
         with:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}

.github/workflows/build_portable_linux_artifacts.yml

@@ -66,7 +66,7 @@ jobs:
     permissions:
       id-token: write
     container:
-      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:d6ae5712a9c7e8b88281d021e907b312cd8a26295b95690baef3e8dde4805858
+      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:db2b63f938941dde2abc80b734e64b45b9995a282896d513a0f3525d4591d6cb
       # --cap-add=SYS_PTRACE : to enable ptrace insided the build container for tsan builds
       # --security-opt seccomp=unconfined : to disable the system call filtering for tsan builds
       options: -v /runner/config:/home/awsconfig/ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined
@@ -164,7 +164,7 @@ jobs:
 
       - name: Configure AWS Credentials for non-forked repos
         if: ${{ always() && !github.event.pull_request.head.repo.fork }}
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708 # v5.1.1
         with:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-ci

.github/workflows/build_portable_linux_python_packages.yml

@@ -27,30 +27,20 @@ on:
         type: string
       package_version:
         type: string
-    outputs:
-      package_find_links_url:
-        description: URL for pip --find-links to install built packages
-        value: ${{ jobs.build_rocm_wheels.outputs.package_find_links_url }}
 
 permissions:
   contents: read
 
 run-name: Build portable Linux Python Packages (${{ inputs.artifact_group }}, ${{ inputs.package_version }})
 
 jobs:
-  build_rocm_wheels:
+  build:
     name: Build Python | ${{ inputs.artifact_group }}
     # Note: GitHub-hosted runners run out of disk space for some gpu families
     runs-on: ${{ github.repository_owner == 'ROCm' && 'azure-linux-scale-rocm' || 'ubuntu-24.04' }}
-    permissions:
-      id-token: write
     container:
-      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:d6ae5712a9c7e8b88281d021e907b312cd8a26295b95690baef3e8dde4805858
-      options: -v /runner/config:/home/awsconfig/
-    outputs:
-      package_find_links_url: ${{ steps.upload.outputs.package_find_links_url }}
+      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:db2b63f938941dde2abc80b734e64b45b9995a282896d513a0f3525d4591d6cb
     env:
-      AWS_SHARED_CREDENTIALS_FILE: /home/awsconfig/credentials.ini
       ARTIFACT_RUN_ID: "${{ inputs.artifact_run_id != '' && inputs.artifact_run_id || github.run_id }}"
       ARTIFACTS_DIR: "${{ github.workspace }}/artifacts"
       PACKAGES_DIR: "${{ github.workspace }}/packages"
@@ -63,13 +53,15 @@ jobs:
       - name: Install Python requirements
         run: pip install -r requirements.txt
 
+      # Note: we could fetch "all" artifacts if we wanted to include more files
       - name: Fetch artifacts
         run: |
           python ./build_tools/fetch_artifacts.py \
             --run-github-repo="${{ inputs.artifact_github_repo }}" \
             --run-id="${{ env.ARTIFACT_RUN_ID }}" \
             --artifact-group="${{ inputs.artifact_group }}" \
-            --output-dir="${{ env.ARTIFACTS_DIR }}"
+            --output-dir="${{ env.ARTIFACTS_DIR }}" \
+            _dev_ _lib_ _run_
 
       - name: Build Python packages
         run: |
@@ -78,20 +70,17 @@ jobs:
             --dest-dir="${{ env.PACKAGES_DIR }}" \
             --version="${{ inputs.package_version }}"
 
-      - name: Configure AWS Credentials for non-forked repos
-        if: ${{ !github.event.pull_request.head.repo.fork }}
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
-        with:
-          aws-region: us-east-2
-          role-to-assume: arn:aws:iam::692859939525:role/therock-ci
-
-      # NOTE: we use `github.run_id` and NOT `env.ARTIFACT_RUN_ID` here!
-      # This ensures that if they are different we _download_ artifacts from the
-      # input run's subdirectory and _upload_ to our current run's subdirectory.
-      - name: Upload Python packages
-        id: upload
+      - name: Inspect Python packages
         run: |
-          python build_tools/github_actions/upload_python_packages.py \
-            --input-packages-dir="${{ env.PACKAGES_DIR }}" \
-            --artifact-group="${{ inputs.artifact_group }}" \
-            --run-id="${{ github.run_id }}"
+          ls -la "${{ env.PACKAGES_DIR }}"
+
+      # TODO(#1559): Sanity check (Linux can't find the directories, maybe Docker issues?)
+
+      # - name: Sanity check Python packages
+      #   run: |
+      #     piprepo build "${{ env.PACKAGES_DIR }}/dist"
+      #     pip install rocm[devel]==${{ inputs.package_version }} \
+      #       --extra-index-url "${{ env.PACKAGES_DIR }}/dist/simple/"
+      #     rocm-sdk test
+
+      # TODO(#1559): upload packages to artifacts S3 bucket and/or a dedicated Python packages bucket

.github/workflows/build_portable_linux_pytorch_wheels.yml

@@ -103,11 +103,10 @@ jobs:
     name: Build | ${{ inputs.amdgpu_family }} | py ${{ inputs.python_version }} | torch ${{ inputs.pytorch_git_ref }}
     runs-on: ${{ github.repository_owner == 'ROCm' && 'azure-linux-scale-rocm' || 'ubuntu-24.04' }}
     container:
-      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:d6ae5712a9c7e8b88281d021e907b312cd8a26295b95690baef3e8dde4805858
+      image: ghcr.io/rocm/therock_build_manylinux_x86_64@sha256:db2b63f938941dde2abc80b734e64b45b9995a282896d513a0f3525d4591d6cb
     env:
       OUTPUT_DIR: ${{ github.workspace }}/output
       PACKAGE_DIST_DIR: ${{ github.workspace }}/output/packages/dist
-      RELEASE_TYPE: ${{ inputs.release_type }}
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
       optional_build_prod_arguments: ""
     outputs:
@@ -127,7 +126,6 @@ jobs:
       # packages, for example via `pip install torch==${TORCH_VERSION}`.
       torch_version: ${{ steps.build-pytorch-wheels.outputs.torch_version }}
       torchaudio_version: ${{ steps.build-pytorch-wheels.outputs.torchaudio_version }}
-      apex_version: ${{ steps.build-pytorch-wheels.outputs.apex_version }}
       torchvision_version: ${{ steps.build-pytorch-wheels.outputs.torchvision_version }}
       triton_version: ${{ steps.build-pytorch-wheels.outputs.triton_version }}
     steps:
@@ -162,7 +160,6 @@ jobs:
         run: |
           ./external-builds/pytorch/pytorch_torch_repo.py checkout --repo-hashtag nightly
           ./external-builds/pytorch/pytorch_audio_repo.py checkout --repo-hashtag nightly
-          ./external-builds/pytorch/pytorch_apex_repo.py checkout --repo-hashtag master
           ./external-builds/pytorch/pytorch_vision_repo.py checkout --repo-hashtag nightly
           ./external-builds/pytorch/pytorch_triton_repo.py checkout
 
@@ -172,12 +169,12 @@ jobs:
         run: |
           ./external-builds/pytorch/pytorch_torch_repo.py checkout --gitrepo-origin https://github.com/ROCm/pytorch.git --repo-hashtag ${{ inputs.pytorch_git_ref }}
           ./external-builds/pytorch/pytorch_audio_repo.py checkout --require-related-commit
-          ./external-builds/pytorch/pytorch_apex_repo.py checkout --require-related-commit
           ./external-builds/pytorch/pytorch_vision_repo.py checkout --require-related-commit
           ./external-builds/pytorch/pytorch_triton_repo.py checkout
 
-      # Note: determine_version.py sets optional_build_prod_arguments in
-      # GITHUB_ENV, which includes --rocm-sdk-version and --version-suffix.
+      - name: Create pip cache directory
+        run: mkdir -p /tmp/pipcache
+
       - name: Determine optional arguments passed to `build_prod_wheels.py`
         if: ${{ inputs.rocm_version }}
         run: |
@@ -192,6 +189,7 @@ jobs:
           ./external-builds/pytorch/build_prod_wheels.py \
             build \
             --install-rocm \
+            --pip-cache-dir /tmp/pipcache \
             --index-url "${{ inputs.cloudfront_url }}/${{ inputs.amdgpu_family }}/" \
             --clean \
             --output-dir ${{ env.PACKAGE_DIST_DIR }} ${{ env.optional_build_prod_arguments }}
@@ -201,20 +199,9 @@ jobs:
         run: |
           python external-builds/pytorch/sanity_check_wheel.py ${{ env.PACKAGE_DIST_DIR }}/
 
-      - name: Generate PyTorch manifest
-        run: |
-          python build_tools/github_actions/generate_pytorch_manifest.py \
-            --manifest-dir "${{ env.PACKAGE_DIST_DIR }}/manifests" \
-            --python-version "${{ inputs.python_version }}" \
-            --pytorch-git-ref "${{ inputs.pytorch_git_ref }}" \
-            --pytorch-dir "external-builds/pytorch/pytorch" \
-            --pytorch-audio-dir "external-builds/pytorch/pytorch_audio" \
-            --pytorch-vision-dir "external-builds/pytorch/pytorch_vision" \
-            --triton-dir "external-builds/pytorch/triton"
-
       - name: Configure AWS Credentials
         if: always()
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708 # v5.1.1
         with:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}
@@ -225,16 +212,6 @@ jobs:
           aws s3 cp ${{ env.PACKAGE_DIST_DIR }}/ s3://${{ env.S3_BUCKET_PY }}/${{ inputs.s3_staging_subdir }}/${{ inputs.amdgpu_family }}/ \
             --recursive --exclude "*" --include "*.whl"
 
-      - name: Upload PyTorch manifest to S3 artifacts
-        if: ${{ github.repository_owner == 'ROCm' }}
-        run: |
-          python build_tools/github_actions/upload_pytorch_manifest.py \
-            --dist-dir "${{ env.PACKAGE_DIST_DIR }}" \
-            --run-id "${{ github.run_id }}" \
-            --amdgpu-family "${{ inputs.amdgpu_family }}" \
-            --python-version "${{ inputs.python_version }}" \
-            --pytorch-git-ref "${{ inputs.pytorch_git_ref }}"
-
       - name: (Re-)Generate Python package release index for staging
         if: ${{ github.repository_owner == 'ROCm' }}
         env:
@@ -290,7 +267,6 @@ jobs:
     env:
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
       CP_VERSION: "${{ needs.build_pytorch_wheels.outputs.cp_version }}"
-      APEX_VERSION: "${{ needs.build_pytorch_wheels.outputs.apex_version }}"
       TORCH_VERSION: "${{ needs.build_pytorch_wheels.outputs.torch_version }}"
       TORCHAUDIO_VERSION: "${{ needs.build_pytorch_wheels.outputs.torchaudio_version }}"
       TORCHVISION_VERSION: "${{ needs.build_pytorch_wheels.outputs.torchvision_version }}"
@@ -305,7 +281,7 @@ jobs:
 
       - name: Configure AWS Credentials
         if: always()
-        uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 # v6.0.0
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708 # v5.1.1
         with:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}
@@ -330,8 +306,7 @@ jobs:
             --include "torch-${TORCH_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
             --include "torchaudio-${TORCHAUDIO_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
             --include "torchvision-${TORCHVISION_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
-            --include "triton-${TRITON_VERSION}-${CP_VERSION}-linux_x86_64.whl" \
-            --include "apex-${APEX_VERSION}-${CP_VERSION}-linux_x86_64.whl"
+            --include "triton-${TRITON_VERSION}-${CP_VERSION}-linux_x86_64.whl"
 
       - name: (Re-)Generate Python package release index
         if: ${{ env.upload == 'true' }}