From dfb1a19595f7b10061338f73112906a4864f6651 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Thu, 21 May 2026 22:51:53 +0000 Subject: [PATCH 01/10] [ROCm] Add TheRock 7.13 and TheRock latest to CI test matrices 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. --- .github/workflows/wheel_tests_continuous.yml | 12 +++++++++--- .github/workflows/wheel_tests_nightly_release.yml | 12 +++++++++--- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 73e3bdfa385c..575d2cd68c60 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -358,7 +358,9 @@ jobs: artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", manylinux-image: ""}, + {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", manylinux-image: ""}, + {label: "TheRock 7.13", wheel-version: "7", release-version: "therock-7.13.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.13.0:latest"}, + {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, ] exclude: - artifact: "jax-rocm-pjrt" @@ -391,7 +393,9 @@ jobs: runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] python: ["3.12"] rocm: [ - {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "TheRock 7.13", wheel-version: "7", release-version: "therock-7.13.0", tag: "therock-7.13.0"}, + {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", tag: "therock-latest"}, ] name: "Pytest ROCm (${{ matrix.rocm.label }})" with: @@ -418,7 +422,9 @@ jobs: runner: ["linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] python: ["3.12"] rocm: [ - {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "TheRock 7.13", wheel-version: "7", release-version: "therock-7.13.0", tag: "therock-7.13.0"}, + {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", tag: "therock-latest"}, ] enable-x64: [0] name: "Bazel ROCm (${{ matrix.rocm.label }})" diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index 9c5dfb9ce04c..f2b43d140d7b 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -187,7 +187,9 @@ jobs: artifact: ["jax-rocm-plugin", "jax-rocm-pjrt"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", manylinux-image: ""}, + {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", manylinux-image: ""}, + {label: "TheRock 7.13", wheel-version: "7", release-version: "therock-7.13.0", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-7.13.0:latest"}, + {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", manylinux-image: "ghcr.io/rocm/jax-manylinux_2_28-therock-latest:latest"}, ] exclude: - artifact: "jax-rocm-pjrt" @@ -220,7 +222,9 @@ jobs: runner: ["linux-x86-64-1gpu-amd", "linux-x86-64-4gpu-amd", "linux-x86-64-8gpu-amd"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "TheRock 7.13", wheel-version: "7", release-version: "therock-7.13.0", tag: "therock-7.13.0"}, + {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", tag: "therock-latest"}, ] name: "Pytest ROCm (${{ matrix.rocm.label }})" with: @@ -247,7 +251,9 @@ jobs: runner: ["linux-x86-64-1gpu-amd"] python: ["3.12", "3.13", "3.14"] rocm: [ - {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "ROCm 7.2.0", wheel-version: "7", release-version: "7.2.0", tag: "rocm720"}, + {label: "TheRock 7.13", wheel-version: "7", release-version: "therock-7.13.0", tag: "therock-7.13.0"}, + {label: "TheRock latest", wheel-version: "7", release-version: "therock-latest", tag: "therock-latest"}, ] enable-x64: [0] name: "Bazel ROCm (${{ matrix.rocm.label }})" From d430e3bbedbdaf8a90dde8857e51cbe111c3c00e Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Fri, 5 Jun 2026 21:53:19 +0000 Subject: [PATCH 02/10] [ROCm] Embed TheRock-aware wheel RUNPATH at link time 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-, 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. --- jaxlib/rocm/rocm_rpath.bzl | 1 - 1 file changed, 1 deletion(-) diff --git a/jaxlib/rocm/rocm_rpath.bzl b/jaxlib/rocm/rocm_rpath.bzl index c8c0c1bf4fa3..0030a41e0b5f 100644 --- a/jaxlib/rocm/rocm_rpath.bzl +++ b/jaxlib/rocm/rocm_rpath.bzl @@ -65,7 +65,6 @@ def _rocm_wheel_rpaths(): where TheRock is in the container's system Python). The legacy `/opt/rocm` fallback is appended by the caller so it stays last. """ - # `_rocm_sdk_core/lib` holds the non-arch runtime libs; `_rocm_sdk_libraries # [...]/lib` holds the math libs + per-arch kernel data. site_libs covers the # multi-arch + per-family packages; core_libs is the cross-Python/absolute From b0ab89091dedccfb7329a3f861dbfeb5fe5bfc88 Mon Sep 17 00:00:00 2001 From: gulsumgudukbay Date: Mon, 8 Jun 2026 20:07:42 +0000 Subject: [PATCH 03/10] [ROCm] Put TheRock SDK bin on PATH via so rocminfo/rocm-smi resolve in /opt/rocm-less images From c0ecae1c71ab2ac51b84381447d78b5f13b5ab89 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Fri, 19 Jun 2026 00:52:35 +0000 Subject: [PATCH 04/10] [ROCm] Add blocking Bazel PR gate workflow 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///); 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. --- .github/workflows/bazel_rocm.yml | 5 ++ .github/workflows/bazel_rocm_presubmit.yml | 32 +++++++++++ build/rocm/ci_blocking_test_targets.txt | 15 +++++ build/rocm/ci_test_targets.txt | 51 +++++++++++++++++ build/rocm/rocm.bazelrc | 5 ++ ci/run_bazel_test_rocm_rbe.sh | 64 ++-------------------- 6 files changed, 113 insertions(+), 59 deletions(-) create mode 100644 build/rocm/ci_blocking_test_targets.txt create mode 100644 build/rocm/ci_test_targets.txt diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index aa92dbfeed1b..4bc23ec6fb4c 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -91,6 +91,10 @@ on: description: "Should latest XLA be used?" type: string required: true + gate-targets-file: + description: "Path to a Bazel target pattern file for the blocking gate. When set, overrides the default full test suite." + type: string + default: '' halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -191,6 +195,7 @@ jobs: PYTHON: ${{ inputs.python }} ENABLE_X64: ${{ inputs.enable-x64 }} RBE_IMAGE: ${{ steps.rbe-image.outputs.rbe_image }} + JAXCI_GATE_TARGETS_FILE: ${{ inputs.gate-targets-file }} run: | ./ci/run_bazel_test_rocm_rbe.sh \ --config=single_gpu \ diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 6a3a4bbbf5ec..904fd1d9f6da 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -52,3 +52,35 @@ jobs: build_jax: ${{ 'true' }} run_multiaccelerator_tests: ${{ 'false' }} clone_main_xla: 0 + + blocking-gate: + if: github.repository_owner == 'jax-ml' || github.repository_owner == 'ROCm' + uses: ./.github/workflows/bazel_rocm.yml + strategy: + fail-fast: false + matrix: + python: ["3.12"] + runner: ["linux-x86-64-1gpu-amd"] + rocm: + - {label: "ROCm 7.2.0", tag: "rocm720", release-version: "7.2.0", s3-version: "0.11.0/7.2.0", experimental: false} + - {label: "TheRock 7.13", tag: "therock-7.13.0", release-version: "therock-7.13.0", s3-version: "0.11.0/therock-7.13.0", experimental: true} + - {label: "TheRock latest", tag: "therock-latest", release-version: "therock-latest", s3-version: "0.11.0/therock-latest", experimental: true} +# Begin Presubmit Naming Check - name modification requires internal check to be updated + name: "ROCm blocking gate (${{ matrix.rocm.label }})" +# End Presubmit Naming Check github-rocm-blocking-presubmits + continue-on-error: ${{ matrix.rocm.experimental }} + with: + runner: ${{ matrix.runner }} + python: ${{ matrix.python }} + rocm-version: "7" + rocm-tag: ${{ matrix.rocm.tag }} + rocm-release-version: ${{ matrix.rocm.release-version }} + download-jax-from-gcs: 0 + jaxlib-version: "head" + build_jaxlib: "false" + build_jax: "true" + skip-download-jaxlib-and-plugins-from-gcs: 0 + s3_download_uri: "s3://jax-ci-amd/rocm-wheels/releases/${{ matrix.rocm.s3-version }}" + gate-targets-file: "build/rocm/ci_blocking_test_targets.txt" + run_multiaccelerator_tests: "false" + clone_main_xla: 0 diff --git a/build/rocm/ci_blocking_test_targets.txt b/build/rocm/ci_blocking_test_targets.txt new file mode 100644 index 000000000000..37f2bad7c0ae --- /dev/null +++ b/build/rocm/ci_blocking_test_targets.txt @@ -0,0 +1,15 @@ +# Blocking PR gate targets for amd-main. +# This list is intentionally small and conservative. +# Add or remove targets here; use - prefix to exclude specific tests within a suite. +# Keep only tests that are fast and reliably passing — flaky tests will block all PRs. + +//tests:backend_independent_tests +//tests:lax_numpy_indexing_test_gpu +//tests:lax_reference_test_gpu +//tests:lax_test_gpu +//tests:jax_numpy_linalg_test_gpu +//tests:nn_test_gpu +//tests:random_test_gpu +//tests:xmap_test_gpu +//tests:api_test_gpu +//jaxlib/tools:check_gpu_wheel_sources_test diff --git a/build/rocm/ci_test_targets.txt b/build/rocm/ci_test_targets.txt new file mode 100644 index 000000000000..83a9e1651b7b --- /dev/null +++ b/build/rocm/ci_test_targets.txt @@ -0,0 +1,51 @@ +//tests:gpu_tests +//tests:backend_independent_tests +//tests/pallas:gpu_tests +//tests/pallas:backend_independent_tests +//jaxlib/tools:check_gpu_wheel_sources_test +-//tests/pallas:pallas_test_gpu +-//tests/pallas:ops_test_gpu +-//tests/pallas:ops_test_mgpu_gpu +-//tests/pallas:pallas_shape_poly_test_gpu +-//tests/pallas:pallas_vmap_test_gpu +-//tests/pallas:triton_pallas_test_gpu +-//tests:export_harnesses_multi_platform_test_gpu +-//tests:jet_test_gpu +-//tests:lax_autodiff_test_gpu +-//tests:lax_numpy_setops_test_gpu +-//tests:lax_numpy_test_gpu +-//tests:lax_test_gpu +-//tests:linalg_test_gpu +-//tests:logging_test_gpu +-//tests:random_lax_test_gpu +-//tests:scipy_signal_test_gpu +-//tests:stax_test_gpu +-//tests:ode_test_gpu +-//tests:lobpcg_test_gpu +-//tests:scipy_stats_test_gpu +-//tests:nn_test_gpu +-//tests:lax_scipy_sparse_test_gpu +-//tests:lax_scipy_spectral_dac_test_gpu +-//tests:lax_scipy_special_functions_test_gpu +-//tests:cholesky_update_test_gpu +-//tests:api_test_gpu +-//tests:ann_test_gpu +-//tests:experimental_rnn_test_gpu +-//tests:lax_vmap_test_gpu +-//tests:qdwh_test_gpu +-//tests:scaled_dot_test_gpu +-//tests:scipy_spatial_test_gpu +-//tests:shape_poly_test_gpu +-//tests:sparsify_test_gpu +-//tests:lax_numpy_reducers_test_gpu +-//tests:scipy_optimize_test_gpu +#TODO(mgaonka-amd) re-enable these tests once this issue is fixed. +-//tests:lax_numpy_einsum_test_gpu +-//tests:lax_scipy_test_gpu +-//tests:sparse_test_gpu +-//tests:sparse_bcoo_bcsr_test_gpu +-//tests:svd_test_gpu +-//tests:eigh_test_gpu +-//tests:image_test_gpu +# Excluded when using pre-built plugin wheels (incompatible with py_import deps). +-//tests:buffer_callback_test_gpu diff --git a/build/rocm/rocm.bazelrc b/build/rocm/rocm.bazelrc index 5d16d1726ad9..9b18364a99f0 100644 --- a/build/rocm/rocm.bazelrc +++ b/build/rocm/rocm.bazelrc @@ -90,6 +90,11 @@ test:rocm_rbe_dynamic --experimental_local_execution_delay=1000 # Maximum number of test jobs to run locally in parallel (50% of available CPUs) test:rocm_rbe_dynamic --local_resources=cpu=HOST_CPUS*0.5 +# Target list for the blocking PR gate — reads build/rocm/ci_blocking_test_targets.txt. +# Cannot be combined with explicit target patterns on the command line. +# Used locally via: bazel ... test --config=rocm_ci_local_single_gpu --config=rocm_blocking_targets +test:rocm_blocking_targets --target_pattern_file=build/rocm/ci_blocking_test_targets.txt + # Used for rules_python bootstrap 1.8.0+ build --@rules_python//python/config_settings:bootstrap_impl=script --repo_env=RULES_PYTHON_ENABLE_PIPSTAR=0 test --@rules_python//python/config_settings:bootstrap_impl=script --repo_env=RULES_PYTHON_ENABLE_PIPSTAR=0 diff --git a/ci/run_bazel_test_rocm_rbe.sh b/ci/run_bazel_test_rocm_rbe.sh index 2db9bde4919c..519454223f7d 100755 --- a/ci/run_bazel_test_rocm_rbe.sh +++ b/ci/run_bazel_test_rocm_rbe.sh @@ -45,52 +45,10 @@ echo "Running RBE GPU tests..." TAG_FILTERS="jax_test_gpu,-config-cuda-only,-manual" -TESTS_TO_IGNORE=( - -//tests/pallas:pallas_test_gpu - -//tests/pallas:ops_test_gpu - -//tests/pallas:ops_test_mgpu_gpu - -//tests/pallas:pallas_shape_poly_test_gpu - -//tests/pallas:pallas_vmap_test_gpu - -//tests/pallas:triton_pallas_test_gpu - -//tests:export_harnesses_multi_platform_test_gpu - -//tests:jet_test_gpu - -//tests:lax_autodiff_test_gpu - -//tests:lax_numpy_setops_test_gpu - -//tests:lax_numpy_test_gpu - -//tests:lax_test_gpu - -//tests:linalg_test_gpu - -//tests:logging_test_gpu - -//tests:random_lax_test_gpu - -//tests:scipy_signal_test_gpu - -//tests:stax_test_gpu - -//tests:ode_test_gpu - -//tests:lobpcg_test_gpu - -//tests:scipy_stats_test_gpu - -//tests:nn_test_gpu - -//tests:lax_scipy_sparse_test_gpu - -//tests:lax_scipy_spectral_dac_test_gpu - -//tests:lax_scipy_special_functions_test_gpu - -//tests:cholesky_update_test_gpu - -//tests:api_test_gpu - -//tests:ann_test_gpu - -//tests:experimental_rnn_test_gpu - -//tests:lax_vmap_test_gpu - -//tests:qdwh_test_gpu - -//tests:scaled_dot_test_gpu - -//tests:scipy_spatial_test_gpu - -//tests:shape_poly_test_gpu - -//tests:sparsify_test_gpu - -//tests:lax_numpy_reducers_test_gpu - -//tests:scipy_optimize_test_gpu - #TODO(mgaonka-amd) re-enable these tests once this issue is fixed. - -//tests:lax_numpy_einsum_test_gpu - -//tests:lax_scipy_test_gpu - -//tests:sparse_test_gpu - -//tests:sparse_bcoo_bcsr_test_gpu - -//tests:svd_test_gpu - -//tests:eigh_test_gpu - -//tests:image_test_gpu -) +# JAXCI_GATE_TARGETS_FILE selects which Bazel target pattern file to use. +# Defaults to the full CI suite; set to build/rocm/ci_blocking_test_targets.txt +# for the PR blocking gate. +TARGETS_FILE="${JAXCI_GATE_TARGETS_FILE:-build/rocm/ci_test_targets.txt}" for arg in "$@"; do if [[ "$arg" == "--config=multi_gpu" ]]; then @@ -99,12 +57,6 @@ for arg in "$@"; do if [[ "$arg" == "--config=single_gpu" ]]; then TAG_FILTERS="${TAG_FILTERS},gpu,-multiaccelerator" fi - if [[ "$arg" == "--//jax:build_jaxlib=false" ]]; then - # tests to ignore for pre-built plugin wheels - TESTS_TO_IGNORE+=( - -//tests:buffer_callback_test_gpu - ) - fi done TEST_ARTIFACTS_DIR="test-artifacts" @@ -129,13 +81,7 @@ bazel --bazelrc=build/rocm/rocm.bazelrc test \ --color=yes \ $@ \ --spawn_strategy=local \ - -- \ - //tests:gpu_tests \ - //tests:backend_independent_tests \ - //tests/pallas:gpu_tests \ - //tests/pallas:backend_independent_tests \ - //jaxlib/tools:check_gpu_wheel_sources_test \ - "${TESTS_TO_IGNORE[@]}" || bazel_retval=$? + --target_pattern_file="${TARGETS_FILE}" || bazel_retval=$? echo "::endgroup::" >&2 echo "::group::Cleanup" >&2 From 3ed2fa4b60afc6911165540ae52e2274a09664bf Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Sat, 20 Jun 2026 02:23:04 +0000 Subject: [PATCH 05/10] [ROCm] Rework blocking PR gate: file-driven targets, rocm_gate config, 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. --- .github/workflows/bazel_rocm.yml | 7 ++++- .github/workflows/bazel_rocm_presubmit.yml | 4 ++- .github/workflows/wheel_tests_continuous.yml | 1 + .../workflows/wheel_tests_nightly_release.yml | 1 + build/rocm/ci_blocking_test_targets.txt | 28 ++++++++++--------- build/rocm/rocm.bazelrc | 8 ++++++ 6 files changed, 34 insertions(+), 15 deletions(-) diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index 4bc23ec6fb4c..6401ed830981 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -95,6 +95,10 @@ on: description: "Path to a Bazel target pattern file for the blocking gate. When set, overrides the default full test suite." type: string default: '' + gpu-config: + description: "Bazel GPU config to use for the test run (single_gpu or multi_gpu)." + type: string + default: 'single_gpu' halt-for-connection: description: 'Should this workflow run wait for a remote connection?' type: string @@ -196,9 +200,10 @@ jobs: ENABLE_X64: ${{ inputs.enable-x64 }} RBE_IMAGE: ${{ steps.rbe-image.outputs.rbe_image }} JAXCI_GATE_TARGETS_FILE: ${{ inputs.gate-targets-file }} + GPU_CONFIG: ${{ inputs.gpu-config }} run: | ./ci/run_bazel_test_rocm_rbe.sh \ - --config=single_gpu \ + --config="${GPU_CONFIG}" \ --local_test_jobs=4 \ --//jax:build_jaxlib="${BUILD_JAXLIB}" \ --//jax:build_jax="${BUILD_JAX}" \ diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 904fd1d9f6da..ec630ed7b2f5 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -50,6 +50,7 @@ jobs: enable-x64: ${{ matrix.enable-x64 }} build_jaxlib: ${{ (matrix.jaxlib-version == 'head' && 'wheel') || 'false' }} build_jax: ${{ 'true' }} + gate-targets-file: "build/rocm/ci_test_targets.txt" run_multiaccelerator_tests: ${{ 'false' }} clone_main_xla: 0 @@ -60,7 +61,7 @@ jobs: fail-fast: false matrix: python: ["3.12"] - runner: ["linux-x86-64-1gpu-amd"] + runner: ["linux-x86-64-2gpu-amd"] rocm: - {label: "ROCm 7.2.0", tag: "rocm720", release-version: "7.2.0", s3-version: "0.11.0/7.2.0", experimental: false} - {label: "TheRock 7.13", tag: "therock-7.13.0", release-version: "therock-7.13.0", s3-version: "0.11.0/therock-7.13.0", experimental: true} @@ -82,5 +83,6 @@ jobs: skip-download-jaxlib-and-plugins-from-gcs: 0 s3_download_uri: "s3://jax-ci-amd/rocm-wheels/releases/${{ matrix.rocm.s3-version }}" gate-targets-file: "build/rocm/ci_blocking_test_targets.txt" + gpu-config: "rocm_gate" run_multiaccelerator_tests: "false" clone_main_xla: 0 diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 575d2cd68c60..db9d95b3db59 100644 --- a/.github/workflows/wheel_tests_continuous.yml +++ b/.github/workflows/wheel_tests_continuous.yml @@ -440,5 +440,6 @@ jobs: build_jaxlib: ${{ 'false' }} build_jax: ${{ 'false' }} jaxlib-version: "head" + gate-targets-file: "build/rocm/ci_test_targets.txt" run_multiaccelerator_tests: ${{ 'false' }} clone_main_xla: 0 diff --git a/.github/workflows/wheel_tests_nightly_release.yml b/.github/workflows/wheel_tests_nightly_release.yml index f2b43d140d7b..f3b297f67573 100644 --- a/.github/workflows/wheel_tests_nightly_release.yml +++ b/.github/workflows/wheel_tests_nightly_release.yml @@ -270,6 +270,7 @@ jobs: build_jaxlib: ${{ 'false' }} build_jax: ${{ 'false' }} jaxlib-version: "head" + gate-targets-file: "build/rocm/ci_test_targets.txt" write_to_bazel_remote_cache: 1 run_multiaccelerator_tests: ${{ 'true' }} clone_main_xla: 0 diff --git a/build/rocm/ci_blocking_test_targets.txt b/build/rocm/ci_blocking_test_targets.txt index 37f2bad7c0ae..6d419d5464fd 100644 --- a/build/rocm/ci_blocking_test_targets.txt +++ b/build/rocm/ci_blocking_test_targets.txt @@ -1,15 +1,17 @@ -# Blocking PR gate targets for amd-main. -# This list is intentionally small and conservative. -# Add or remove targets here; use - prefix to exclude specific tests within a suite. -# Keep only tests that are fast and reliably passing — flaky tests will block all PRs. +# Blocking PR gate targets. +# Edit this file to tune which tests gate PRs; no workflow YAML changes needed. +# Use - prefix to exclude specific tests within a suite. +# Keep only tests that are fast and reliably passing — flaky tests block all PRs. +# Multi-GPU tests (tagged multiaccelerator) are supported: the gate runs with +# --config=multi_gpu so the RBE pool provides multi-GPU machines. -//tests:backend_independent_tests -//tests:lax_numpy_indexing_test_gpu -//tests:lax_reference_test_gpu +//tests:export_back_compat_test_gpu +//tests:sparse_test_gpu //tests:lax_test_gpu -//tests:jax_numpy_linalg_test_gpu -//tests:nn_test_gpu -//tests:random_test_gpu -//tests:xmap_test_gpu -//tests:api_test_gpu -//jaxlib/tools:check_gpu_wheel_sources_test +//tests:lax_numpy_test_gpu +//tests:pjit_test_gpu +//tests:export_harnesses_multi_platform_test_gpu +//tests:lax_autodiff_test_gpu +//tests/pallas:pallas_test_gpu +//tests:lax_numpy_operators_test_gpu +//tests:lax_control_flow_test_gpu diff --git a/build/rocm/rocm.bazelrc b/build/rocm/rocm.bazelrc index 9b18364a99f0..d840edd52fba 100644 --- a/build/rocm/rocm.bazelrc +++ b/build/rocm/rocm.bazelrc @@ -50,6 +50,14 @@ build:multi_gpu --repo_env=TF_ROCM_RBE_POOL="linux_x64_multigpu" build:multi_gpu --repo_env=TF_ROCM_RBE_SINGLE_GPU_POOL="linux_x64_multigpu" build:multi_gpu --repo_env=TF_ROCM_RBE_MULTI_GPU_POOL="linux_x64_multigpu" +# Config for the blocking PR gate: uses the multi-GPU RBE pool so both +# single- and multi-GPU tests in ci_blocking_test_targets.txt can run. +# Does NOT match single_gpu/multi_gpu in the script, so no tag filter is +# appended — the explicit target file controls selection entirely. +build:rocm_gate --repo_env=TF_ROCM_RBE_POOL="linux_x64_multigpu" +build:rocm_gate --repo_env=TF_ROCM_RBE_SINGLE_GPU_POOL="linux_x64_gpu" +build:rocm_gate --repo_env=TF_ROCM_RBE_MULTI_GPU_POOL="linux_x64_multigpu" + build:rocm_rbe --tls_client_certificate="/data/ci-cert.crt" build:rocm_rbe --tls_client_key="/data/ci-cert.key" From 606b991d2c46e16834c405e5bc3f09cd5eb0aede Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Sat, 20 Jun 2026 02:26:54 +0000 Subject: [PATCH 06/10] [ROCm] Pass halt-for-connection to blocking-gate on workflow_dispatch --- .github/workflows/bazel_rocm_presubmit.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index ec630ed7b2f5..93c645f91c45 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -82,6 +82,7 @@ jobs: build_jax: "true" skip-download-jaxlib-and-plugins-from-gcs: 0 s3_download_uri: "s3://jax-ci-amd/rocm-wheels/releases/${{ matrix.rocm.s3-version }}" + halt-for-connection: ${{ inputs.halt-for-connection }} gate-targets-file: "build/rocm/ci_blocking_test_targets.txt" gpu-config: "rocm_gate" run_multiaccelerator_tests: "false" From 45220b1c39e79952945b452ecf43a00c559e4ba0 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Sat, 20 Jun 2026 02:32:37 +0000 Subject: [PATCH 07/10] [ROCm] Drop continue-on-error from blocking-gate (unsupported with uses:) --- .github/workflows/bazel_rocm_presubmit.yml | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 93c645f91c45..25f6ff5edc6b 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -63,13 +63,12 @@ jobs: python: ["3.12"] runner: ["linux-x86-64-2gpu-amd"] rocm: - - {label: "ROCm 7.2.0", tag: "rocm720", release-version: "7.2.0", s3-version: "0.11.0/7.2.0", experimental: false} - - {label: "TheRock 7.13", tag: "therock-7.13.0", release-version: "therock-7.13.0", s3-version: "0.11.0/therock-7.13.0", experimental: true} - - {label: "TheRock latest", tag: "therock-latest", release-version: "therock-latest", s3-version: "0.11.0/therock-latest", experimental: true} + - {label: "ROCm 7.2.0", tag: "rocm720", release-version: "7.2.0", s3-version: "0.11.0/7.2.0"} + - {label: "TheRock 7.13", tag: "therock-7.13.0", release-version: "therock-7.13.0", s3-version: "0.11.0/therock-7.13.0"} + - {label: "TheRock latest", tag: "therock-latest", release-version: "therock-latest", s3-version: "0.11.0/therock-latest"} # Begin Presubmit Naming Check - name modification requires internal check to be updated name: "ROCm blocking gate (${{ matrix.rocm.label }})" # End Presubmit Naming Check github-rocm-blocking-presubmits - continue-on-error: ${{ matrix.rocm.experimental }} with: runner: ${{ matrix.runner }} python: ${{ matrix.python }} From c8503f470f2e93d623528d5022efccdb168839d2 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Sat, 20 Jun 2026 02:37:47 +0000 Subject: [PATCH 08/10] [ROCm] Remove TheRock from blocking gate (ROCm 7.2.0 only) --- .github/workflows/bazel_rocm_presubmit.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 25f6ff5edc6b..91f6e2625554 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -63,9 +63,7 @@ jobs: python: ["3.12"] runner: ["linux-x86-64-2gpu-amd"] rocm: - - {label: "ROCm 7.2.0", tag: "rocm720", release-version: "7.2.0", s3-version: "0.11.0/7.2.0"} - - {label: "TheRock 7.13", tag: "therock-7.13.0", release-version: "therock-7.13.0", s3-version: "0.11.0/therock-7.13.0"} - - {label: "TheRock latest", tag: "therock-latest", release-version: "therock-latest", s3-version: "0.11.0/therock-latest"} + - {label: "ROCm 7.2.0", tag: "rocm720", release-version: "7.2.0", s3-version: "0.11.0/7.2.0"} # Begin Presubmit Naming Check - name modification requires internal check to be updated name: "ROCm blocking gate (${{ matrix.rocm.label }})" # End Presubmit Naming Check github-rocm-blocking-presubmits From 99fbfce6621c2aaed81ec55ef4f4f3d9450c6f94 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Sat, 20 Jun 2026 03:01:40 +0000 Subject: [PATCH 09/10] [ROCm] Restore conditional buffer_callback_test_gpu exclusion for pre-built wheels --- build/rocm/ci_test_targets.txt | 2 -- ci/run_bazel_test_rocm_rbe.sh | 12 +++++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/build/rocm/ci_test_targets.txt b/build/rocm/ci_test_targets.txt index 83a9e1651b7b..d15dbe907b0e 100644 --- a/build/rocm/ci_test_targets.txt +++ b/build/rocm/ci_test_targets.txt @@ -47,5 +47,3 @@ -//tests:svd_test_gpu -//tests:eigh_test_gpu -//tests:image_test_gpu -# Excluded when using pre-built plugin wheels (incompatible with py_import deps). --//tests:buffer_callback_test_gpu diff --git a/ci/run_bazel_test_rocm_rbe.sh b/ci/run_bazel_test_rocm_rbe.sh index 519454223f7d..2a6ac5c018c6 100755 --- a/ci/run_bazel_test_rocm_rbe.sh +++ b/ci/run_bazel_test_rocm_rbe.sh @@ -49,6 +49,9 @@ TAG_FILTERS="jax_test_gpu,-config-cuda-only,-manual" # Defaults to the full CI suite; set to build/rocm/ci_blocking_test_targets.txt # for the PR blocking gate. TARGETS_FILE="${JAXCI_GATE_TARGETS_FILE:-build/rocm/ci_test_targets.txt}" +# Effective file passed to --target_pattern_file; may be replaced by a temp +# file below when additional dynamic exclusions are needed. +EFFECTIVE_TARGETS_FILE="${TARGETS_FILE}" for arg in "$@"; do if [[ "$arg" == "--config=multi_gpu" ]]; then @@ -57,6 +60,13 @@ for arg in "$@"; do if [[ "$arg" == "--config=single_gpu" ]]; then TAG_FILTERS="${TAG_FILTERS},gpu,-multiaccelerator" fi + if [[ "$arg" == "--//jax:build_jaxlib=false" ]]; then + # buffer_callback_test_gpu uses py_import deps that are incompatible + # with pre-built plugin wheels; exclude it in that mode. + EFFECTIVE_TARGETS_FILE=$(mktemp) + cat "${TARGETS_FILE}" > "${EFFECTIVE_TARGETS_FILE}" + echo "-//tests:buffer_callback_test_gpu" >> "${EFFECTIVE_TARGETS_FILE}" + fi done TEST_ARTIFACTS_DIR="test-artifacts" @@ -81,7 +91,7 @@ bazel --bazelrc=build/rocm/rocm.bazelrc test \ --color=yes \ $@ \ --spawn_strategy=local \ - --target_pattern_file="${TARGETS_FILE}" || bazel_retval=$? + --target_pattern_file="${EFFECTIVE_TARGETS_FILE}" || bazel_retval=$? echo "::endgroup::" >&2 echo "::group::Cleanup" >&2 From 7406c5d8b4d6a63730ab04773fc5724e331b0d21 Mon Sep 17 00:00:00 2001 From: Marco Minutoli Date: Sat, 20 Jun 2026 03:03:14 +0000 Subject: [PATCH 10/10] [ROCm] Fix bool/string type error in blocking-gate workflow inputs --- .github/workflows/bazel_rocm_presubmit.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/bazel_rocm_presubmit.yml b/.github/workflows/bazel_rocm_presubmit.yml index 91f6e2625554..7a37d6952a3c 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -75,12 +75,12 @@ jobs: rocm-release-version: ${{ matrix.rocm.release-version }} download-jax-from-gcs: 0 jaxlib-version: "head" - build_jaxlib: "false" - build_jax: "true" + build_jaxlib: ${{ 'false' }} + build_jax: ${{ 'true' }} skip-download-jaxlib-and-plugins-from-gcs: 0 s3_download_uri: "s3://jax-ci-amd/rocm-wheels/releases/${{ matrix.rocm.s3-version }}" halt-for-connection: ${{ inputs.halt-for-connection }} gate-targets-file: "build/rocm/ci_blocking_test_targets.txt" gpu-config: "rocm_gate" - run_multiaccelerator_tests: "false" + run_multiaccelerator_tests: ${{ 'false' }} clone_main_xla: 0