Test PR CI trigger - #877
Open
charleshofer wants to merge 9 commits into
Open
charleshofer wants to merge 9 commits into
charleshofer wants to merge 9 commits into
Conversation
Drop 43 of the 45 test-target exclusions in build/rocm/ci_test_targets.txt, leaving only two: - buffer_callback_test_gpu, whose py_import deps are incompatible with pre-built plugin wheels (build_jaxlib=false). - pallas:ops_test_mgpu_gpu, a multi-accelerator test that is excluded from single_gpu runs anyway. Most of the dropped exclusions traced back to a complex-GEMM gap that no longer reproduces: C64 and C128 matmul return correct results on ROCm 7.14. The three teardown-hang exclusions (cholesky_update_test, pallas:pallas_vmap_test and pallas:gpu_paged_attention_test) each passed five consecutive runs with no hang. This brings the suite to 114 single-GPU GPU targets. Validated on MI300X (gfx942); the bazel CI pool is gfx950, so CI is the real gate.
Upstream CI has stopped publishing jax/jaxlib nightly wheels to GCS (gs://jax-nightly-artifacts/latest), leaving the ROCm nightly test legs running against missing/stale base wheels. On the ROCm fork, build the base jax and jaxlib wheels on the same amd-do-linux.jax.cpu runners as the ROCm plugin/pjrt wheels and upload them to the same S3 prefix, so the ROCm test legs consume freshly built wheels instead of the broken GCS bucket. Upstream (jax-ml) behavior is unchanged. - ci/build_rocm_artifacts.sh: allow "jax"/"jaxlib" artifacts; skip auditwheel for the pure-python jax wheel. - wheel_tests_nightly_release.yml: new fork-only build-rocm-jax-jaxlib job; wire it into the ROCm test legs and select S3 as the jax/jaxlib source on the fork. - pytest_rocm.yml / bazel_rocm.yml: new optional jax-jaxlib-source input (default gcs), forwarded to the download action. - download-jax-rocm-wheels: when source is s3, fetch jax/jaxlib from the CloudFront listing alongside the ROCm wheels instead of GCS. (cherry picked from commit 8654db6)
The nightly flow already builds the base jax/jaxlib wheels on the fork's own runners (see "[ROCm] nightly: build jax/jaxlib on fork runners"), because upstream stopped reliably publishing them to gs://jax-nightly-artifacts/latest. The continuous flow was left behind: build-jax-artifact and build-jaxlib-artifact are gated on the jax-ml owner, so on the fork the ROCm legs silently fell back to that same stale GCS bucket. Apply the nightly pattern to the continuous flow so both consume freshly built wheels. Upstream (jax-ml) behavior is unchanged. - wheel_tests_continuous.yml: new fork-only build-rocm-jax-jaxlib job uploading to the same wheel-tests-continuous S3 prefix as the ROCm plugin/pjrt wheels; wired into the needs of both ROCm test legs, which now pass jax-jaxlib-source=s3 on the fork. - select_continuous_wheel_test_jobs.py: register build_rocm_jax_jaxlib as a build job required by run_pytest_rocm and run_bazel_test_rocm, and exclude it from per-job connection halts like the other ROCm jobs. Only Python 3.12 is built, matching the Python versions the continuous ROCm test legs actually consume; unlike plugin/pjrt these wheels are not published, so building 3.13/3.14 every three hours would be wasted work.
build/rocm/ci_test_targets.txt has always listed //tests:backend_independent_tests and //tests/pallas:backend_independent_tests, but none of those targets ever ran. jaxlib/jax.bzl defines those suites as exactly the targets carrying no jax_test_<backend> tag, while run_bazel_test_rocm_rbe.sh passed --build_tag_filters/--test_tag_filters=jax_test_gpu,...,gpu for single_gpu. Bazel requires a target to match at least one positive tag, so the two suites were silently filtered down to nothing. Drop the positive tags from the single_gpu filter and keep only the exclusions. The target pattern file already scopes which suites run, so it is what should decide the set. This is also what the CUDA CI does -- ci/run_bazel_test_cuda_rbe.sh passes --test_tag_filters=-multiaccelerator with no positive tag -- which is why CUDA has been running these tests all along. The multi_gpu filter string is unchanged. Verified with bazel on gfx942 in ghcr.io/rocm/jax-base-ubu24.therock-7.14, using --nobuild to count selected targets: single_gpu, old filter: 112 test targets single_gpu, new filter: 147 test targets (+35) ci_blocking_test_targets.txt (PR gate): 1 target, unchanged The 35 newly enabled targets are all genuinely device-agnostic, e.g. config_test, tree_util_test, util_test, typing_test, xla_bridge_test, mesh_utils_test and pallas:pallas_cost_estimate_test. Nothing else is added: //jaxlib/tools:check_gpu_wheel_sources_test is tagged manual and stays excluded by -manual as before.
The exclusion claimed buffer_callback_test_gpu "uses py_import deps that are incompatible with pre-built plugin wheels (build_jaxlib=false)". That does not reproduce. The target was originally swept into the 45-target blanket TESTS_TO_IGNORE list added by ce621cf ("[ROCm] Include prebuilt rocm plugin pkgs as deps to tests"), and it was one of only two entries left behind when db2242f dropped the other 43 as stale; the stated reason was never verified against a real run. Verified on gfx942 in ghcr.io/rocm/jax-base-ubu24.therock-7.14, using the exact CI configuration -- --config=rocm_clang_hermetic --config=single_gpu with --//jax:build_jaxlib=false --//jax:build_jax=false and the prebuilt jax_rocm7_plugin/jax_rocm7_pjrt wheels staged in dist/ (the same wheels the CI download action stages, which is where the rocm_wheels repo rule picks them up): 51 cases: 23 pass, 28 skipped, 0 failures 5/5 consecutive runs pass, avg 3.5s, no teardown hang The 28 skips are @jtu.run_on_devices("cpu") cases that are CPU-only by design and skip identically on CUDA. The 23 that run are the 22 test_cuda_array_interface cases (11 dtypes x command_buffer_compatible) plus test_side_effect, so this adds real coverage of the __cuda_array_interface__ import path on ROCm. Brings the single-GPU suite to 148 targets. The blocking PR gate is unchanged at 1 target. Also corrects the comment on the remaining exclusion. ops_test_mgpu_gpu is a Mosaic GPU test, not a multi-accelerator one, and it carries no multiaccelerator tag -- verified by query -- so the claim that it is "excluded from single_gpu runs anyway" was wrong: the explicit exclusion is the only thing keeping it out.
Several fused attention tests were failing due to WAR hazard between VALU and GDS instructions which is fixed by llvm/llvm-project#203770. So reenable these tests on ROCm platform. (cherry picked from commit fd31e39)
The fork-only build-rocm-jax-jaxlib job was still pinned to TheRock 7.14 while every leg that consumes its output had moved to 10.0. Since release-version is interpolated into the shared S3 prefix, the pytest and bazel legs looked for jax/jaxlib under the 10.0.0 prefix while the build uploaded to the 7.14.0 one, and failed with "Could not find jaxlib wheel". This job block has no upstream counterpart, so it was not covered by the upstream switch of the pinned ROCm legs to 10.0.
magaonka-amd
force-pushed
the
amd-main
branch
2 times, most recently
from
September 2, 2026 21:33
360638e to
e156ed6
Compare
magaonka-amd
force-pushed
the
amd-main
branch
3 times, most recently
from
September 13, 2026 22:57
564a88e to
a53a18e
Compare
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.