Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion .github/workflows/bazel_rocm.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}" \
Expand Down
32 changes: 32 additions & 0 deletions .github/workflows/bazel_rocm_presubmit.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
13 changes: 10 additions & 3 deletions .github/workflows/wheel_tests_continuous.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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:
Expand All @@ -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 }})"
Expand All @@ -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
13 changes: 10 additions & 3 deletions .github/workflows/wheel_tests_nightly_release.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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:
Expand All @@ -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 }})"
Expand All @@ -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
Expand Down
17 changes: 17 additions & 0 deletions build/rocm/ci_blocking_test_targets.txt
Original file line number Diff line number Diff line change
@@ -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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

//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
49 changes: 49 additions & 0 deletions build/rocm/ci_test_targets.txt
Original file line number Diff line number Diff line change
@@ -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
13 changes: 13 additions & 0 deletions build/rocm/rocm.bazelrc
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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
Expand Down
70 changes: 13 additions & 57 deletions ci/run_bazel_test_rocm_rbe.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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
Expand Down
1 change: 0 additions & 1 deletion jaxlib/rocm/rocm_rpath.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down