fix(jax): match jax_rocm plugin wheels by ROCm major version - #6956
Conversation
The JAX version-detection step hardcoded jax_rocm7_plugin and jax_rocm7_pjrt. JAX 0.10.0 emits jax_rocm10_* wheels, so the build failed with 'Did not find jax_rocm7_plugin wheel'. Match the wheels with a jax_rocm*_plugin / jax_rocm*_pjrt glob so detection tracks the ROCm major version instead of a fixed one. Signed-off-by: Maisam Arif <Maisam.Arif@amd.com>
❌ PR Check — Action Required
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
therock-pr-bot
Bot
added
the
Not ready to Review
|
🚫 Please fix the failed policies before requesting reviews. The following policy checks failed:
The |
|
will this solution work for jax v0.9.1.. |
|
Resolves fails in #6960 |
There was a problem hiding this comment.
The jax_rocm1_* naming only happens on the 0.9.1 path. That branch builds through ROCm/rocm-jax's wrapper, which passes the truncated value straight through:
build/build.py build --wheels=jax-rocm-plugin --rocm_path=/opt/rocm-10.0.0a20260729 --rocm_version=1 ...
The 0.10.0 path never passes --rocm_version at all and it derives the namespace inside the JAX build from ROCM_PATH via rocm-sdk path --root, which is why it correctly produces jax_rocm10_. Since #6708 leaves the release set at 0.10.0 through 0.11.0 and deletes the native flow, the only code path that can emit jax_rocm1_ goes away with it.
|
This fix is necessary but not sufficient, test_linux_jax_wheels_partial.yml:240 and test_linux_jax_wheels.yml:253 still pip install "jax_rocm7_plugin==...", plus the allowlists in download_python_packages.py, s3_management/manage.py, and promote_packages.py, so the build will go green and the next stage will fail. Second, the small robustness point: get_wheel_version takes wheel_paths[0] in filesystem order when the glob matches more than one wheel, so sorting or validating the major against the ROCm version would make it deterministic. |
Issue: #6958 ## Motivation The JAX plugin/pjrt wheels embed the ROCm major version in their package name (`jax_rocm7_plugin` for ROCm 7, `jax_rocm10_plugin` for ROCm 10). Now that `version.json` is at 10.0.0, the builds emit `jax_rocm10_*`, but everything downstream of the build still hardcodes the ROCm 7 spelling: the test job can't resolve the package, and the promotion/index allowlists silently drop the wheels. #6956 fixes detection at build time. This fixes the consumers. ## Changes - `test_linux_jax_wheels.yml`, `test_linux_jax_wheels_partial.yml`: derive the package name from `rocm_version` instead of installing a literal `jax_rocm7_plugin`. Fails fast if `rocm_version` is empty. - `download_python_packages.py`, `promote_packages.py`: match the wheels with a `^jax_rocm\d+_(plugin|pjrt)$` pattern instead of listing them once per ROCm release. - `s3_management/manage.py`: allow the wheels into the package index by prefix, alongside the existing `rocm_sdk_device_` / `amd_torch_device` cases. ## Test Plan New `build_tools/packaging/tests/download_python_packages_test.py` 7 passed. Package recognition, before vs after: | wheel | before | after | | --- | --- | --- | | `jax_rocm7_plugin-0.9.2+rocm7.15.0-…` | `promote`, allowed | `promote`, allowed | | `jax_rocm10_plugin-0.10.0+rocm10.0.0-…` | `unknown`, **not** allowed | `promote`, allowed | Name derivation, against real version strings: ```console rocm_version='10.0.0' -> jax_rocm10_plugin rocm_version='10.0.0a20260729' -> jax_rocm10_plugin rocm_version='7.15.0' -> jax_rocm7_plugin rocm_version='7.15.0.dev0+abc' -> jax_rocm7_plugin rocm_version='' -> error (fail fast) ``` ## Test Results ```console $ python3 -m pytest build_tools/packaging/tests/download_python_packages_test.py -v platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0 rootdir: /home/ggudukba/JAX/th-jaxpkg/build_tools configfile: pyproject.toml collected 7 items CategorizePackageTest::test_jax_rocm10_wheels_are_promoted PASSED [ 14%] CategorizePackageTest::test_jax_rocm7_wheels_are_promoted PASSED [ 28%] CategorizePackageTest::test_jaxlib_is_still_promoted PASSED [ 42%] CategorizePackageTest::test_unrelated_jax_rocm_name_is_unknown PASSED [ 57%] IsAllowedMultiArchPackageTest::test_jax_rocm10_wheels_are_allowed PASSED [ 71%] IsAllowedMultiArchPackageTest::test_jax_rocm7_wheels_are_allowed PASSED [ 85%] IsAllowedMultiArchPackageTest::test_unknown_package_is_not_allowed PASSED [100%] ===================== 7 passed, 8 subtests passed in 0.06s ===================== ``` ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/CONTRIBUTING.md.
3 participants
Summary
The JAX multi-arch build fails at the version-detection step with:
build_tools/github_actions/write_jax_versions.pyhardcodesjax_rocm7_pluginandjax_rocm7_pjrt, but the JAX 0.10.0 build now emitsjax_rocm10_plugin/jax_rocm10_pjrt. The ROCm major version embedded in the wheel name tracks the ROCm release, so a fixed7no longer matches.Fix
Detect the wheels with a
jax_rocm*_plugin/jax_rocm*_pjrtglob so the step tracks whatever ROCm major version the build produces (rocm7, rocm10, and future), instead of a hardcoded major. Docstring and error messages updated to match.Scope
Limited to the failing build-time detection logic. The other
jax_rocm7_*references live in the promotion/download allowlists (download_python_packages.py,promote_packages.py,s3_management/manage.py), which are separate jobs and intentionally explicit; they are left out of this change.Test
pre-commit run --files build_tools/github_actions/write_jax_versions.pypasses (black, etc.).jax_rocm*_plugin-*.whlmatchesjax_rocm10_plugin-0.10.0+devrocm10.0.0-...whlwhile the oldjax_rocm7_plugin-*.whlmatches nothing (reproducing the failure).CI
This is a CI script change; the JAX build job exercises it.