Release portable Linux JAX Wheels (gfx120X-all, nightly, 7.13.0a20260429) #1683
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Copyright Advanced Micro Devices, Inc. | |
| # SPDX-License-Identifier: MIT | |
| name: Release portable Linux JAX Wheels | |
| on: | |
| workflow_call: | |
| inputs: | |
| amdgpu_family: | |
| required: true | |
| type: string | |
| release_type: | |
| description: The type of release to build ("dev", "nightly", or "prerelease"). All developer-triggered jobs should use "dev"! | |
| type: string | |
| default: "dev" | |
| s3_subdir: | |
| description: S3 subdirectory, not including the GPU-family | |
| type: string | |
| default: "v2" | |
| s3_staging_subdir: | |
| description: Staging subdirectory to push the wheels for test | |
| type: string | |
| default: "v2-staging" | |
| 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 | |
| rocm_version: | |
| description: ROCm version to install (e.g. "7.10.0a20251124") | |
| type: string | |
| tar_url: | |
| description: "URL to TheRock tarball to build against (e.g. https://rocm.nightlies.amd.com/tarball/therock-dist-linux-gfx94X-dcgpu-7.10.0a20251124.tar.gz)" | |
| 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: | |
| type: choice | |
| options: | |
| - gfx101X-dgpu | |
| - gfx103X-all | |
| - gfx110X-all | |
| - gfx1150 | |
| - gfx1151 | |
| - gfx1152 | |
| - gfx1153 | |
| - gfx120X-all | |
| - gfx900 | |
| - gfx906 | |
| - gfx908 | |
| - gfx90a | |
| - gfx94X-dcgpu | |
| - gfx950-dcgpu | |
| default: gfx94X-dcgpu | |
| release_type: | |
| description: The type of release to build ("dev", "nightly", or "prerelease"). All developer-triggered jobs should use "dev"! | |
| type: string | |
| default: "dev" | |
| s3_subdir: | |
| description: S3 subdirectory, not including the GPU-family | |
| type: string | |
| default: "v2" | |
| s3_staging_subdir: | |
| description: "Staging subdirectory to push the wheels for test" | |
| type: string | |
| default: "v2-staging" | |
| cloudfront_url: | |
| description: CloudFront URL pointing to Python index | |
| type: string | |
| default: "https://rocm.devreleases.amd.com/v2" | |
| cloudfront_staging_url: | |
| description: CloudFront base URL pointing to staging Python index | |
| type: string | |
| default: "https://rocm.devreleases.amd.com/v2-staging" | |
| rocm_version: | |
| description: ROCm version to install (e.g. "7.10.0a20251124") | |
| type: string | |
| tar_url: | |
| description: "URL to TheRock tarball to build (e.g. https://rocm.nightlies.amd.com/tarball/therock-dist-linux-gfx94X-dcgpu-7.10.0a20251124.tar.gz)" | |
| type: string | |
| ref: | |
| description: "TheRock branch, tag or SHA to checkout. Defaults to the reference or SHA that triggered the workflow." | |
| type: string | |
| default: '' | |
| permissions: | |
| id-token: write | |
| contents: read | |
| packages: read | |
| run-name: Release portable Linux JAX Wheels (${{ inputs.amdgpu_family }}, ${{ inputs.release_type }}, ${{ inputs.rocm_version }}) | |
| jobs: | |
| release: | |
| name: Release | ${{ inputs.amdgpu_family }} | py ${{ matrix.python_version }} | jax ${{ matrix.jax_ref }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| python_version: ["3.11", "3.12", "3.13", "3.14"] | |
| jax_ref: ["rocm-jaxlib-v0.8.2", "rocm-jaxlib-v0.9.0", "rocm-jaxlib-v0.9.1"] | |
| include: | |
| - jax_ref: "rocm-jaxlib-v0.8.2" | |
| build_jaxlib: true | |
| - jax_ref: "rocm-jaxlib-v0.9.0" | |
| build_jaxlib: true | |
| - jax_ref: "rocm-jaxlib-v0.9.1" | |
| build_jaxlib: false | |
| uses: ./.github/workflows/build_linux_jax_wheels.yml | |
| with: | |
| amdgpu_family: ${{ inputs.amdgpu_family }} | |
| python_version: ${{ matrix.python_version }} | |
| release_type: ${{ inputs.release_type }} | |
| s3_subdir: ${{ inputs.s3_subdir }} | |
| s3_staging_subdir: ${{ inputs.s3_staging_subdir }} | |
| cloudfront_url: ${{ inputs.cloudfront_url }} | |
| cloudfront_staging_url: ${{ inputs.cloudfront_staging_url }} | |
| rocm_version: ${{ inputs.rocm_version }} | |
| tar_url: ${{ inputs.tar_url }} | |
| jax_ref: ${{ matrix.jax_ref }} | |
| build_jaxlib: ${{ matrix.build_jaxlib }} | |
| ref: ${{ inputs.ref }} |