add jax core test

.github/workflows/build_linux_jax_wheels.yml

@@ -23,6 +23,17 @@ on:
       tar_url:
         description: URL to TheRock tarball to build against
         type: string
+      cloudfront_url:
+        description: CloudFront URL pointing to Python index
+        required: true
+        type: string
+      cloudfront_staging_url:
+        description: CloudFront base URL pointing to staging Python index
+        required: true
+        type: string
+      ref:
+        description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
+        type: string
   workflow_dispatch:
     inputs:
       amdgpu_family:
@@ -57,6 +68,21 @@ on:
       tar_url:
         description: URL to TheRock tarball to build against
         type: string
+      cloudfront_url:
+        description: CloudFront URL pointing to Python index
+        required: true
+        type: string
+      cloudfront_staging_url:
+        description: CloudFront base URL pointing to staging Python index
+        required: true
+        type: string
+      ref:
+        description: "Branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow."
+        type: string
+      jax_ref:
+        description: Branch/tag/sha of rocm/rocm-jax to build
+        type: string
+        default: "rocm-jaxlib-v0.7.1"
 
 permissions:
   id-token: write
@@ -72,16 +98,21 @@ jobs:
     env:
       PACKAGE_DIST_DIR: ${{ github.workspace }}/jax/jax_rocm_plugin/wheelhouse
       S3_BUCKET_PY: "therock-${{ inputs.release_type }}-python"
+    outputs:
+      jax_whl_list: ${{ steps.jax_wheel_list.outputs.jax_whl_list }}
     steps:
-      - name: Checkout TheRock
+      - name: Checkout
         uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        with:
+          repository: ${{ inputs.repository || github.repository }}
+          ref: ${{ inputs.ref || '' }}
 
       - name: Checkout JAX
         uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
         with:
           path: jax
           repository: rocm/rocm-jax
-          ref: ${{ matrix.jax_ref }}
+          ref: ${{ inputs.jax_ref || matrix.jax_ref }}
 
       - name: Configure Git Identity
         run: |
@@ -117,6 +148,13 @@ jobs:
           aws-region: us-east-2
           role-to-assume: arn:aws:iam::692859939525:role/therock-${{ inputs.release_type }}-releases
 
+      - name: jax wheel list
+        id: jax_wheel_list
+        run: |
+          export jax_whl_list="$(python3 ./build_tools/mapping_built_jax_wheels.py --dir "${{ env.PACKAGE_DIST_DIR }}")"
+          echo "jax_whl_list=${jax_whl_list}" >> "$GITHUB_OUTPUT"
+          echo "${jax_whl_list}" >> "$GITHUB_STEP_SUMMARY"
+
       - name: Upload wheels to S3
         if: ${{ github.repository_owner == 'ROCm' }}
         run: |
@@ -130,3 +168,23 @@ jobs:
           source .venv/bin/activate
           pip3 install boto3 packaging
           python3 ./build_tools/third_party/s3_management/manage.py ${{ inputs.s3_subdir }}/${{ inputs.amdgpu_family }}
+
+  test_jax_wheels:
+    name: Test JAX wheels | ${{ inputs.amdgpu_family }}
+    needs: [build_jax_wheels]
+    if: ${{ github.repository_owner == 'ROCm' }}
+    permissions:
+      id-token: write
+      contents: read
+      packages: write
+    uses: ./.github/workflows/test_linux_jax_wheels.yml
+    with:
+      amdgpu_family: ${{ inputs.amdgpu_family }}
+      release_type: ${{ inputs.release_type }}
+      package_index_url: ${{ inputs.cloudfront_staging_url }}
+      rocm_version: ${{ inputs.rocm_version }}
+      tar_url: ${{ inputs.tar_url }}
+      python_versions: ${{ inputs.python_versions }}
+      jax_ref: rocm-jaxlib-v0.7.1
+      jax_whl_list: ${{ needs.build_jax_wheels.outputs.jax_whl_list }}
+      test_runs_on: "linux-mi325-1gpu-ossci-rocm"

.github/workflows/release_portable_linux_packages.yml

@@ -303,8 +303,12 @@ jobs:
               "python_versions": "3.12",
               "release_type": "${{ env.RELEASE_TYPE }}",
               "s3_subdir": "${{ env.S3_STAGING_SUBDIR }}",
+              "s3_staging_subdir": "${{ env.S3_STAGING_SUBDIR }}",
               "rocm_version": "${{ needs.setup_metadata.outputs.version }}",
-              "tar_url": "${{ steps.url-encode-tar.outputs.tar_url }}"
+              "tar_url": "${{ steps.url-encode-tar.outputs.tar_url }}",
+              "cloudfront_url": "${{ needs.setup_metadata.outputs.cloudfront_url }}",
+              "cloudfront_staging_url": "${{ needs.setup_metadata.outputs.cloudfront_staging_url }}",
+              "ref": "${{ inputs.ref || '' }}"
             }
 
       - name: Save cache

.github/workflows/test_linux_jax_wheels.yml

@@ -0,0 +1,167 @@
+name: Test JAX Wheels
+on:
+  workflow_call:
+    inputs:
+      amdgpu_family:
+        description: GPU family to test
+        required: true
+        type: string
+      release_type:
+        description: The type of release
+        required: true
+        type: string
+      rocm_version:
+        description: ROCm version to use in Docker
+        required: true
+        type: string
+      tar_url:
+        description: URL to TheRock tarball to build against
+        required: true
+        type: string
+      test_runs_on:
+        required: true
+        type: string
+        default: "linux-mi325-1gpu-ossci-rocm"
+      python_versions:
+        description: Python version string used to compute CP tag
+        required: true
+        type: string
+      jax_ref:
+        description: Branch/tag/sha of rocm/rocm-jax to build
+        required: false
+        type: string
+        default: "rocm-jaxlib-v0.7.1"
+      package_index_url:
+        required: true
+        type: string
+      jax_whl_list:
+        description: list of built jax wheel filenames to fetch & test
+        required: true
+        type: string
+
+  workflow_dispatch:
+    inputs:
+      amdgpu_family:
+        type: choice
+        options:
+          - gfx110X-dgpu
+          - gfx1151
+          - gfx120X-all
+          - gfx94X-dcgpu
+          - gfx950-dcgpu
+        default: gfx94X-dcgpu
+      release_type:
+        type: choice
+        options: [dev, nightly, prerelease]
+        default: dev
+      rocm_version:
+        type: string
+        required: true
+      tar_url:
+        type: string
+        required: true
+      test_runs_on:
+        description: Runner label to use. The selected runner should have a GPU supported by amdgpu_family
+        required: true
+        type: string
+        default: "linux-mi325-1gpu-ossci-rocm"
+      python_versions:
+        type: string
+        required: true
+        default: "3.12"
+      jax_ref:
+        type: string
+        default: "rocm-jaxlib-v0.7.1"
+      package_index_url:
+        description: Base Python package index URL to test, typically nightly/dev URL with a "v2" or "v2-staging" subdir (without a GPU family subdir)
+        required: true
+        type: string
+        default: "https://rocm.nightlies.amd.com/v2"
+      jax_whl_list:
+        description: list of built jax wheel filenames to fetch & test
+        required: true
+        type: string
+
+permissions:
+  id-token: write
+  contents: read
+  packages: none
+
+jobs:
+  test_jax_wheels:
+    name: Test | ${{ inputs.amdgpu_family }}
+    runs-on: ${{ inputs.test_runs_on }}
+    env:
+      WHEELHOUSE_DIR: ${{ github.workspace }}/jax/wheelhouse
+    steps:
+      - name: Checkout TheRock
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        with:
+          repository: ${{ inputs.repository || github.repository }}
+          ref: ${{ inputs.ref || '' }}
+
+      - name: Checkout rocm-jax (plugin + build scripts)
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        with:
+          path: jax
+          repository: rocm/rocm-jax
+          ref: ${{ inputs.jax_ref }}
+
+      - name: Setup Python
+        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
+        with:
+          python-version: "3.12"
+
+      - name: Prepare wheelhouse and download JAX wheels
+        shell: bash
+        working-directory: jax
+        run: |
+          python3 -m pip install --upgrade pip
+          python3 -m pip install requests
+          python3 ../build_tools/fetch_wheels.py \
+            --cloudfront_url ${{ inputs.package_index_url }} \
+            --amdgpu_family ${{ inputs.amdgpu_family }} \
+            --dir "wheelhouse" \
+            --list_whls '${{ inputs.jax_whl_list }}'
+
+      - name: Compute ROCM_VERSION_SHORT
+        working-directory: jax
+        env:
+            ROCM_VERSION: ${{ inputs.rocm_version }}
+        run: |
+            # Extract major.minor.patch from ROCM_VERSION
+            # sed command captures first three numeric components separated by dots and reconstructs them as major.minor.patch
+            ROCM_VERSION_SHORT="$(echo "${ROCM_VERSION}" | sed -E 's/^([0-9]+)\.([0-9]+)\.([0-9]+).*/\1.\2.\3/')"
+            echo "ROCM_VERSION_SHORT=${ROCM_VERSION_SHORT}" >> "$GITHUB_ENV"
+            echo "Using semantic ROCm version: ${ROCM_VERSION_SHORT}"
+
+      - name: Build JAX Docker image
+        working-directory: jax
+        env:
+          ROCM_VERSION: ${{ inputs.rocm_version }}
+        run: |
+            python3 build/ci_build \
+            --rocm-version="${ROCM_VERSION_SHORT}" \
+            --therock-path="${{ inputs.tar_url }}" \
+            build_dockers \
+            -f ubu24
+            # Assign JAX_IMAGE to the expected image name produced by the build script
+            JAX_IMAGE="jax-ubu24.rocm$(echo "${ROCM_VERSION_SHORT}" | tr -d '.'):latest"
+            echo "JAX_IMAGE=${JAX_IMAGE}" >> "$GITHUB_ENV"
+            echo "Built JAX Docker image: ${JAX_IMAGE}"
+
+      - name: Checkout JAX tests repo (for extended tests)
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        with:
+          repository: rocm/jax
+          ref: ${{ inputs.jax_ref }}
+          path: jax/jax_tests
+
+      - name: Run unit tests
+        working-directory: jax
+        env:
+          ROCM_VERSION: ${{ inputs.rocm_version }}
+        run: |
+            python3 build/ci_build test \
+              "${JAX_IMAGE}" \
+              --test-cmd "pytest jax_tests/tests/core_test.py"

build_tools/fetch_wheels.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""
+Download .whl files from a CloudFront URL amdgpu_family into a specified directory.
+
+USAGE:
+    python download_from_cloudfront.py \
+        --cloudfront_url ${PACKAGE_INDEX_URL} \
+        --amdgpu_family ${amdgpu_family} \
+        --dir ${WHEELHOUSE_DIR} \
+        --list_whls '["jax_pjrt.whl","jax_plugin.whl","jaxlib.whl"]'
+
+EXAMPLE:
+    python download_from_cloudfront.py \
+        --cloudfront_url https://rocm.nightlies.amd.com/v2-staging  \
+        --amdgpu_family gfx94X-dcgpu \
+        --dir wheelhouse \
+        --list_whls $jax_whl_list
+
+This will:
+    - Create wheelhouse if it does not exist.
+    - Download:
+        https://rocm.nightlies.amd.com/v2-staging/gfx94X-dcgpu/jax_pjrt.whl
+        https://rocm.nightlies.amd.com/v2-staging/gfx94X-dcgpu/jax_plugin.whl
+        https://rocm.nightlies.amd.com/v2-staging/gfx94X-dcgpu/jaxlib.whl
+    - Save all files into wheelhouse
+"""
+
+import argparse
+import ast
+import logging
+from pathlib import Path
+from urllib.parse import quote, urljoin
+import shutil
+import requests
+from github_actions.github_actions_utils import gha_append_step_summary
+
+LOG = logging.getLogger(__name__)
+
+
+def download_one(session, url: str, dest: Path, bufsize: int = 64 * 1024, timeout: int = 30):
+
+    # Stream-download url to dest using a temporary .part file.
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+
+    with session.get(url, stream=True, timeout=timeout) as r:
+        r.raise_for_status()
+        r.raw.decode_content = True
+        with tmp.open("wb") as fh:
+            shutil.copyfileobj(r.raw, fh, length=bufsize)
+
+    tmp.replace(dest)
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+
+    p = argparse.ArgumentParser("Simple wheel downloader")
+    p.add_argument("--cloudfront_url", required=True)
+    p.add_argument("--amdgpu_family", required=True)
+    p.add_argument("--dir", required=True, dest="wheelhouse_dir")
+    p.add_argument(
+        "--list_whls",
+        required=True,
+        help='Python list literal of filenames, e.g. \'["jax_pjrt.whl","jax_plugin.whl","jaxlib.whl"]\'',
+    )
+    args = p.parse_args()
+
+    try:
+        whl_list = ast.literal_eval(args.list_whls)
+        if not isinstance(whl_list, list):
+            raise ValueError
+    except Exception:
+        LOG.error("Invalid --list_whls; provide a Python list literal like '[\"jax_pjrt.whl\",\"jax_plugin.whl\",\"jaxlib.whl\"]'")
+        raise SystemExit(1)
+
+    base = args.cloudfront_url.rstrip("/") + "/"
+    family = args.amdgpu_family.strip("/")
+    out_dir = Path(args.wheelhouse_dir)
+
+    session = requests.Session()
+
+    failures = []
+    for name in whl_list:
+        encoded = quote(name, safe="")
+        url = urljoin(base, f"{family}/{encoded}")
+        dest = out_dir / name
+
+        LOG.info("Downloading %s -> %s", url, dest)
+        try:
+            download_one(session, url, dest)
+            msg = f"Downloaded {name} -> {out_dir}"
+            LOG.info(msg)
+            gha_append_step_summary(msg)
+        except Exception as e:
+            msg = f"Failed {name}: {e}"
+            LOG.error(msg)
+            gha_append_step_summary(msg)
+            failures.append(name)
+
+    if failures:
+        LOG.error("Failed downloads: %s", ", ".join(failures))
+        raise SystemExit(1)
+
+    LOG.info("All downloads complete.")
+
+
+if __name__ == "__main__":
+    main()