[ROCm] Add blocking Bazel PR gate#806
Open
mminutoli wants to merge 10 commits into
Open
Conversation
mminutoli
force-pushed
the
rocm-blocking-pr-gate
branch
3 times, most recently
from
23d09af to
db866ec
Compare
Extends the wheel test continuous and nightly release workflows to run against TheRock 7.13.0 and TheRock latest builds in addition to ROCm 7.2.0.
Consolidate the ROCm wheel RUNPATH strategy in jaxlib/rocm/rocm_rpath.bzl and embed it at link time via _WHEEL_RPATHS instead of a post-build patchelf step (applied under rocm_path_type=link_only, which the release wheel config sets). A single wheel resolves ROCm libs across TheRock pip wheels (1- and 2-deep $ORIGIN layouts), TheRock tarballs under /opt/rocm-<ver>, and legacy /opt/rocm; the loader silently skips $ORIGIN entries whose directories don't exist. Also pin the hermetic build ROCm to the image's TheRock version via ROCM_DISTRO_URL so the wheel's ROCm SONAMEs match the runtime image.
…n /opt/rocm-less images
mminutoli
force-pushed
the
rocm-blocking-pr-gate
branch
2 times, most recently
from
6e0f791 to
29183d9
Compare
Adds a blocking-gate job to bazel_rocm_presubmit.yml that runs a curated subset of Bazel ROCm tests via RBE on every PR, gating on ROCm 7.2.0, TheRock 7.13, and TheRock latest. Wheels are pinned to a versioned S3 releases path (releases/<jax-version>/<rocm-version>/); update the s3-version matrix field when publishing new release wheels. Test target selection is now fully file-driven: - build/rocm/ci_test_targets.txt: full CI suite (targets + exclusions), used by run_bazel_test_rocm_rbe.sh as the default via --target_pattern_file - build/rocm/ci_blocking_test_targets.txt: curated gate subset, selected via the JAXCI_GATE_TARGETS_FILE env var (set by bazel_rocm.yml when gate-targets-file input is provided) This removes the hard-coded TESTS_TO_IGNORE array from the script and the conditional buffer_callback_test_gpu exclusion (now in ci_test_targets.txt). The standalone bazel_rocm_gate.yml and run_bazel_test_rocm_blocking.sh from the earlier prototype are removed.
mminutoli
force-pushed
the
rocm-blocking-pr-gate
branch
from
29183d9 to
c0ecae1
Compare
…, explicit target files - Replace ad-hoc target list in ci_blocking_test_targets.txt with the 10-test curated set selected by coverage/priority analysis: export_back_compat, sparse, lax, lax_numpy, pjit, export_harnesses, lax_autodiff, pallas, lax_numpy_operators, lax_control_flow. - Add rocm_gate Bazel config to rocm.bazelrc: routes to the multi-GPU RBE pool (linux_x64_multigpu) without matching single_gpu/multi_gpu in the script, so no tag filter is appended and the explicit target file controls selection entirely (including pjit_test_gpu which is tagged multiaccelerator). - Switch blocking-gate runner to linux-x86-64-2gpu-amd so that dynamic RBE execution can dispatch multiaccelerator tests locally without hitting a single-GPU limitation. - Add gpu-config input to bazel_rocm.yml (default: single_gpu); blocking gate passes rocm_gate. Thread GPU_CONFIG env var to the run step. - Make all bazel_rocm.yml callers explicit about which target file to use via gate-targets-file: presubmit run-tests, continuous, and nightly all pass build/rocm/ci_test_targets.txt; blocking gate passes ci_blocking_test_targets.txt.
mminutoli
changed the title
alekstheod
reviewed
Jun 22, 2026
| # Multi-GPU tests (tagged multiaccelerator) are supported: the gate runs with | ||
| # --config=multi_gpu so the RBE pool provides multi-GPU machines. | ||
|
|
||
| //tests:export_back_compat_test_gpu |
Collaborator
There was a problem hiding this comment.
I would suggest avoid using explicit targets in your gating CI. The reason is that these targets are getting renamed or removed. It leads to errors like target not found. It is always better to use either tag filters, skip test case internally or --test_filter argument to filter out failing tests. Consider it gives more granularity so more tests are executed.
mminutoli
force-pushed
the
amd-main
branch
3 times, most recently
from
3f9e8b4 to
f25f2b4
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
3 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.
Summary
Adds a blocking PR gate job (
blocking-gate) tobazel_rocm_presubmit.ymlthat runs a curated subset of Bazel tests on every PR using pre-built wheels pinned to a versioned S3 path.Gate design
build/rocm/ci_blocking_test_targets.txt— edit this file to tune the gate without touching workflow YAML.export_back_compat,sparse,lax,lax_numpy,pjit,export_harnesses_multi_platform,lax_autodiff,pallas,lax_numpy_operators,lax_control_flow.continue-on-error: true— shown but non-blocking).s3://jax-ci-amd/rocm-wheels/releases/<jax-version>/<rocm-version>/; update thes3-versionmatrix field when publishing new release wheels.linux-x86-64-2gpu-amd— 2 GPUs required because dynamic RBE execution can dispatchpjit_test_gpu(taggedmultiaccelerator) locally.rocm_gateBazel config (build/rocm/rocm.bazelrc)New config that routes to the multi-GPU RBE pool (
linux_x64_multigpu) without matchingsingle_gpuormulti_gpuin the script. This means no tag filter is appended — the explicit target file controls selection entirely, allowing both single- and multi-GPU tests to run.Script changes (
ci/run_bazel_test_rocm_rbe.sh)JAXCI_GATE_TARGETS_FILEenv var selects the target file; defaults tobuild/rocm/ci_test_targets.txtfor the full suite.gpu,-multiaccelerator/multiaccelerator) still applies based on--config=single_gpu/--config=multi_gpu;rocm_gateintentionally matches neither.Explicit target files for all callers
All
bazel_rocm.ymlcallers now explicitly passgate-targets-file:run-tests, continuous, nightly →build/rocm/ci_test_targets.txtblocking-gate→build/rocm/ci_blocking_test_targets.txtgpu-configinput (bazel_rocm.yml)New optional input (default:
single_gpu) that controls which Bazel GPU config is passed to the script. The blocking gate passesrocm_gate.After merging
Add
ROCm blocking gate (ROCm 7.2.0)as a required status check in branch protection settings.