Resolve latest MaxText Transformer Engine wheel from ROCm MaxText releases

Author: psanal35

.github/workflows/benchmark_rocm.yml

@@ -1,14 +1,14 @@
 # CI - Benchmark ROCm
 #
-# This workflow runs ROCm benchmarks in ROCm team's GHCR containers.
+# This workflow runs the ROCm benchmarks in ROCm team's GHCR containers.
 # It can be triggered manually via workflow_dispatch or called by other workflows
 # via workflow_call.
 #
 # It consists of the following job:
-# run-benchmark:
+# run-benchmarks:
 #    - Runs in ROCm team's container (ghcr.io/rocm/jax-base-ubu24-rocm*:latest)
 #    - Downloads the JAX and jaxlib wheels from GCS, and ROCm plugins from S3.
-#    - Executes the target benchmark runner script at `targets/<target>/run.sh`.
+#    - Executes the target benchmark scripts at `targets/<target>/run.sh`.
 name: CI - Benchmark ROCm
 on:
   workflow_dispatch:
@@ -135,7 +135,7 @@ env:
   UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
 
 jobs:
-  run-benchmark:
+  run-benchmarks:
     defaults:
       run:
         # Set the shell to bash as GitHub actions run with /bin/sh by default
@@ -144,7 +144,7 @@ jobs:
     continue-on-error: true
     # Run in ROCm team's GHCR container with GPU access
     container:
-      image: ghcr.io/rocm/jax-base-ubu24.rocm722:2d65281b00de2bcafc811247d563b6e5e7c887af
+      image: ghcr.io/rocm/jax-base-ubu24.${{ inputs.rocm-tag }}:latest  # zizmor: ignore[unpinned-images]
       credentials:
         username: ${{ github.actor }}
         password: ${{ secrets.GITHUB_TOKEN }}
@@ -187,7 +187,7 @@ jobs:
           TARGET: ${{ inputs.target }}
           WORKLOAD: ${{ inputs.workload }}
         timeout-minutes: 120
-        run: ./ci/benchmark_targets/${TARGET}/run_${TARGET}_rocm.sh" --workload "${WORKLOAD}"
+        run: ./ci/benchmark_targets/${TARGET}/run_${TARGET}_rocm.sh --workload "${WORKLOAD}"
       - name: Upload GitHub artifacts
         if: always()
         continue-on-error: true

.github/workflows/wheel_benchmarks_nightly_release.yml

@@ -1,3 +1,9 @@
+# CI - Wheel Tests (Nightly/Release)
+#
+# This workflow is used to benchmark the JAX wheels that were built by internal CI jobs.
+#
+# 1. run-benchmark-rocm: Calls the `benchmark_rocm.yml` workflow which downloads the JAX wheels and
+#                     ROCm plugins from S3 and runs ROCm benchmarks.
 name: CI - Wheel Benchmarks (Nightly/Release)
 
 on:
@@ -55,6 +61,4 @@ jobs:
       jaxlib-version: "head"
       skip-download-jaxlib-and-plugins-from-gcs: ${{inputs.skip-download-jaxlib-and-plugins-from-gcs}}
       gcs_download_uri: ${{ inputs.gcs_download_uri }}
-      s3_download_uri: ${{ inputs.s3_download_uri }}
       use-te: "1"
-      te-wheel-url: "https://github.com/ROCm/maxtext/releases/download/te-rocm-wheels-2026-05-04-86438dc3d04e/transformer_engine-2.12.0.dev0+86438dc3-1.mi355-cp312-cp312-linux_x86_64.whl"

ci/benchmark_targets/maxtext_rocm/run_maxtext_rocm.sh

@@ -43,11 +43,27 @@ source ./ci/utilities/install_wheels_locally.sh
 
 [[ -f "${ENV_FILE}" ]] && source "${ENV_FILE}"
 
-[[ "${USE_TE:-0}" == "1" ]] && {
-  : "${TE_WHEEL_URL:?TE_WHEEL_URL must be set when USE_TE=1}"
-  echo "[setup] installing Transformer Engine"
+if [[ "${USE_TE:-0}" == "1" ]]; then
+  echo "[setup] resolving latest Transformer Engine wheel"
+  PY_TAG="cp$(echo "${JAXCI_HERMETIC_PYTHON_VERSION:-3.12}" | tr -d '.')"
+  TE_WHEEL_URL="$(
+    curl -fsSL https://api.github.com/repos/ROCm/maxtext/releases \
+      | grep "browser_download_url:" \
+      | grep "te-rocm-wheels-" \
+      | grep "${PY_TAG}" \
+      | grep linux_x86_64.whl \
+      | head -n1 \
+      | cut -d '"' -f4
+  )"
+
+  [[ -n "${TE_WHEEL_URL}" ]] || {
+    echo "failed to resolve Transformer Engine wheel for ${PY_TAG}" >&2
+    exit 1
+  }
+  
+  echo "[setup] installing Transformer Engine from ${TE_WHEEL_URL}"
   "${PYTHON_BIN}" -m pip install --no-deps "${TE_WHEEL_URL}"
-}
+fi
 
 export PY_COLORS=1
 export NCCL_DEBUG=WARN