diff --git a/.github/workflows/bazel_rocm.yml b/.github/workflows/bazel_rocm.yml index aa92dbfeed1b..6401ed830981 100644 --- a/.github/workflows/bazel_rocm.yml +++ b/.github/workflows/bazel_rocm.yml @@ -91,6 +91,14 @@ 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: '' + 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 @@ -191,9 +199,11 @@ 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 }} + 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 6a3a4bbbf5ec..7a37d6952a3c 100644 --- a/.github/workflows/bazel_rocm_presubmit.yml +++ b/.github/workflows/bazel_rocm_presubmit.yml @@ -50,5 +50,37 @@ 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 + + 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-2gpu-amd"] + rocm: + - {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 + 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 }}" + 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' }} clone_main_xla: 0 diff --git a/.github/workflows/wheel_tests_continuous.yml b/.github/workflows/wheel_tests_continuous.yml index 73e3bdfa385c..db9d95b3db59 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 }})" @@ -434,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 9c5dfb9ce04c..f3b297f67573 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 }})" @@ -264,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 new file mode 100644 index 000000000000..6d419d5464fd --- /dev/null +++ b/build/rocm/ci_blocking_test_targets.txt @@ -0,0 +1,17 @@ +# 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:export_back_compat_test_gpu +//tests:sparse_test_gpu +//tests:lax_test_gpu +//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/ci_test_targets.txt b/build/rocm/ci_test_targets.txt new file mode 100644 index 000000000000..d15dbe907b0e --- /dev/null +++ b/build/rocm/ci_test_targets.txt @@ -0,0 +1,49 @@ +//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 diff --git a/build/rocm/rocm.bazelrc b/build/rocm/rocm.bazelrc index 5d16d1726ad9..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" @@ -90,6 +98,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..2a6ac5c018c6 100755 --- a/ci/run_bazel_test_rocm_rbe.sh +++ b/ci/run_bazel_test_rocm_rbe.sh @@ -45,52 +45,13 @@ 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}" +# 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 @@ -100,10 +61,11 @@ for arg in "$@"; do 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 - ) + # 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 @@ -129,13 +91,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="${EFFECTIVE_TARGETS_FILE}" || bazel_retval=$? echo "::endgroup::" >&2 echo "::group::Cleanup" >&2 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