Switch to use rbe workers (#405)

* Switch to use rbe workers

* Introduce rocm rbe pools

* First check for multigpu tag

* Fix buildifier issue

* Use valid platform name

* Fix docker url

* Run multigpu tests locally

* Build mutigpu tests locally

* Ignore numa related leaks

* Limit the supported archs

* Build locally run remotely

* Enable remote and disk_cache for jax tests

* Execute iota_test locally

* Mark failing tests as local

* Ignore iota_test as it is flaky and large

* Make dot operation local and exclude flaky test

* Bump to rocm7 container

* Disable failing test on tsan

* Switch rbe image

* Fixing hermetic build

* Force tsan builds to run locally

* Set proper multigpu tag for all_reduce_test

* Run flaky on rbe tests locally

* Make collective ops test local

Author: alekstheod

build_tools/rocm/BUILD

@@ -34,21 +34,33 @@ filegroup(
 )
 
 genrule(
-    name = "san_wrapper_script",
-    srcs = [":sanitizer_ignore_lists"],
-    outs = ["san_wrapper.sh"],
+    name = "exclusive_wrapper_script",
+    outs = ["exclusive_wrapper.sh"],
     cmd = """
       echo '#!/bin/bash' > $@
-      echo 'exec "$$@"' >> $@
+      echo 'exec {lock_fd}>/var/lock/gpulock || exit 1' >> $@
+      echo 'flock "$$lock_fd"' >> $@
+      echo '"$$@"' >> $@
+      echo 'return_code=$$?' >> $@
+      echo 'flock -u "$$lock_fd"' >> $@
+      echo 'exit $$return_code' >> $@
       chmod +x $@
     """,
 )
 
+# this wrapper ensures the test target
+# take into account any changes in the ignore list files
+sh_binary(
+    name = "exclusive_local_wrapper",
+    srcs = [":exclusive_wrapper_script"],
+    visibility = ["//visibility:public"],
+)
+
 # this wrapper ensures the test target
 # take into account any changes in the ignore list files
 sh_binary(
     name = "sanitizer_wrapper",
-    srcs = [":san_wrapper_script"],
+    srcs = [":exclusive_wrapper_script"],
     data = [":sanitizer_ignore_lists"],
     visibility = ["//visibility:public"],
 )

build_tools/rocm/lsan_ignore_list.txt

@@ -3,3 +3,4 @@ leak:libstdc++.so
 leak:libamdhip64.so
 leak:libhiprtc.so
 leak:librccl.so
+leak:hwloc_bitmap_alloc

build_tools/rocm/platform/linux_x64/BUILD

@@ -11,7 +11,8 @@ platform(
         "@bazel_tools//tools/cpp:clang",
     ],
     exec_properties = {
-        "container-image": "rocm/tensorflow-build@sha256:7cd444ac48657fee2f5087fbda7766266704d3f8fb2299f681952ae4eabed060",
+        # rocm/tensorflow-build:2.18-jammy-python3.9-rocm7.0.2
+        "container-image": "docker://rocm/tensorflow-build@sha256:a2672ff2510b369b4a5f034272a518dc93c2e492894e3befaeef19649632ccaa",
         "OSFamily": "Linux",
     },
 )
@@ -24,7 +25,8 @@ platform(
         "@bazel_tools//tools/cpp:clang",
     ],
     exec_properties = {
-        "container-image": "rocm/tensorflow-build@sha256:7cd444ac48657fee2f5087fbda7766266704d3f8fb2299f681952ae4eabed060",
+        # rocm/tensorflow-build:2.18-jammy-python3.9-rocm7.0.2
+        "container-image": "docker://rocm/tensorflow-build@sha256:a2672ff2510b369b4a5f034272a518dc93c2e492894e3befaeef19649632ccaa",
         "OSFamily": "Linux",
         "Pool": "linux_x64_gpu",
     },

build_tools/rocm/rocm_xla.bazelrc

@@ -5,9 +5,9 @@ build:rocm_dev --remote_cache="https://wardite.cluster.engflow.com"
 
 build:rocm_rbe --bes_backend="grpcs://wardite.cluster.engflow.com"
 build:rocm_rbe --bes_results_url="https://wardite.cluster.engflow.com/invocation/"
-build:rocm_rbe --host_platform="//build_tools/rocm/platform/linux_x64:linux_x64_gpu"
-build:rocm_rbe --extra_execution_platforms="//build_tools/rocm/platform/linux_x64:linux_x64_gpu"
-build:rocm_rbe --platforms="//build_tools/rocm/platform/linux_x64:linux_x64_gpu"
+build:rocm_rbe --host_platform="//build_tools/rocm/platform/linux_x64:linux_x64"
+build:rocm_rbe --extra_execution_platforms="//build_tools/rocm/platform/linux_x64:linux_x64"
+build:rocm_rbe --platforms="//build_tools/rocm/platform/linux_x64:linux_x64"
 build:rocm_rbe --bes_timeout=600s
 build:rocm_rbe --tls_client_certificate="/tf/certificates/ci-cert.crt"
 build:rocm_rbe --tls_client_key="/tf/certificates/ci-cert.key"
@@ -19,7 +19,7 @@ build:rocm_rbe --remote_timeout=3600
 build:rocm_rbe --remote_download_minimal
 build:rocm_rbe --remote_upload_local_results
 
-test:rocm_rbe --strategy=TestRunner=local
+test:rocm_rbe --strategy=TestRunner=remote,local
 
 build:asan --strip=never
 build:asan --copt -fsanitize=address
@@ -62,7 +62,9 @@ build:xla_sgpu -- \
 -//xla/pjrt/distributed:topology_util_test \
 -//xla/pjrt/distributed:client_server_test \
 -//xla/service/gpu/tests:dynamic_shared_memory_test_amdgpu_any \
--//xla/service/gpu/tests:gpu_cub_sort_test_amdgpu_any
+-//xla/service/gpu/tests:gpu_cub_sort_test_amdgpu_any \
+-//xla/tests:iota_test_amdgpu_any \
+-//xla/tests:reduce_window_test_amdgpu_any # TODO: return when it is not flaky!
 
 test:xla_mgpu -- \
 //xla/tests:collective_ops_e2e_test \

build_tools/rocm/run_jax_ut.sh

@@ -16,6 +16,10 @@ python build/build.py build \
 # TODO: run the tests when they are green
 bazel build \
     --config=rocm \
+    --disk_cache=/tf/disk_cache/jaxlib-v0.7.1 \
+    --remote_cache="grpcs://wardite.cluster.engflow.com" \
+    --tls_client_certificate="/tf/certificates/ci-cert.crt" \
+    --tls_client_key="/tf/certificates/ci-cert.key" \
     --build_tag_filters=cpu,gpu,-tpu,-config-cuda-only \
     --test_tag_filters=cpu,gpu,-tpu,-config-cuda-only \
     --action_env=TF_ROCM_AMDGPU_TARGETS=gfx908,gfx90a,gfx942 \

build_tools/rocm/run_xla.sh

@@ -27,17 +27,17 @@ N_BUILD_JOBS=$(grep -c ^processor /proc/cpuinfo)
 rocm-smi -i
 STATUS=$?
 if [ $STATUS -ne 0 ]; then TF_GPU_COUNT=1; else
-   TF_GPU_COUNT=$(rocm-smi -i|grep 'Device ID' |grep 'GPU' |wc -l)
+    TF_GPU_COUNT=$(rocm-smi -i | grep 'Device ID' | grep 'GPU' | wc -l)
 fi
 TF_TESTS_PER_GPU=1
 N_TEST_JOBS=$(expr ${TF_GPU_COUNT} \* ${TF_TESTS_PER_GPU})
-amdgpuname=(`rocminfo | grep gfx | head -n 1`)
+amdgpuname=($(rocminfo | grep gfx | head -n 1))
 AMD_GPU_GFX_ID=${amdgpuname[1]}
 echo ""
 echo "Bazel will use ${N_BUILD_JOBS} concurrent build job(s) and ${N_TEST_JOBS} concurrent test job(s) for gpu ${AMD_GPU_GFX_ID}."
 echo ""
 
-export PYTHON_BIN_PATH=`which python3`
+export PYTHON_BIN_PATH=$(which python3)
 export TF_NEED_ROCM=1
 export ROCM_PATH="/opt/rocm"
 
@@ -99,12 +99,13 @@ BAZEL_DISK_CACHE_SIZE=100G
 BAZEL_DISK_CACHE_DIR="/tf/disk_cache/rocm-jaxlib-v0.7.1"
 mkdir -p ${BAZEL_DISK_CACHE_DIR}
 if [ ! -d /tf/pkg ]; then
-	mkdir -p /tf/pkg
+    mkdir -p /tf/pkg
 fi
 
 SCRIPT_DIR=$(realpath $(dirname $0))
 TAG_FILTERS=$($SCRIPT_DIR/rocm_tag_filters.sh),-multigpu,-multi_gpu_h100,requires-gpu-amd,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial
 
+RBE_OPTIONS=()
 SANITIZER_ARGS=()
 if [[ $1 == "asan" ]]; then
     SANITIZER_ARGS+=("--config=asan")
@@ -121,6 +122,12 @@ elif [[ $1 == "tsan" ]]; then
         HostExecuteStartThunkTest*
         HostExecuteDoneThunkTest*
     )
+
+    #  tsan tests appear to be flaky in rbe due to the heavy load
+    #  force them to run locally
+    RBE_OPTIONS+=(
+         --strategy=TestRunner=local
+    )
     shift
 fi
 
@@ -139,16 +146,18 @@ bazel --bazelrc=build_tools/rocm/rocm_xla.bazelrc test \
     --flaky_test_attempts=3 \
     --keep_going \
     --local_test_jobs=${N_TEST_JOBS} \
+    --repo_env=TF_ROCM_AMDGPU_TARGETS=gfx908,gfx90a,gfx942,gfx1100 \
     --test_env=TF_TESTS_PER_GPU=$TF_TESTS_PER_GPU \
     --action_env=XLA_FLAGS="--xla_gpu_enable_llvm_module_compilation_parallelism=true --xla_gpu_force_compilation_parallelism=16" \
     --run_under=//build_tools/ci:parallel_gpu_execute \
     --test_env=MIOPEN_FIND_ENFORCE=5 \
     --test_env=MIOPEN_FIND_MODE=1 \
     --test_filter=-$(IFS=: ; echo "${EXCLUDED_TESTS[*]}") \
     "${SANITIZER_ARGS[@]}" \
-    "$@"
+    "$@" \
+    "${RBE_OPTIONS[@]}"
 
 # clean up bazel disk_cache
 bazel shutdown \
-  --disk_cache=${BAZEL_DISK_CACHE_DIR} \
-  --experimental_disk_cache_gc_max_size=${BAZEL_DISK_CACHE_SIZE}
+    --disk_cache=${BAZEL_DISK_CACHE_DIR} \
+    --experimental_disk_cache_gc_max_size=${BAZEL_DISK_CACHE_SIZE}

build_tools/rocm/run_xla_multi_gpu.sh

@@ -76,6 +76,7 @@ EXCLUDED_TESTS=(
 SCRIPT_DIR=$(realpath $(dirname $0))
 TAG_FILTERS="$($SCRIPT_DIR/rocm_tag_filters.sh)"
 
+RBE_OPTIONS=()
 SANITIZER_ARGS=()
 if [[ $1 == "asan" ]]; then
     SANITIZER_ARGS+=("--run_under=//build_tools/rocm:sanitizer_wrapper")
@@ -99,6 +100,12 @@ elif [[ $1 == "tsan" ]]; then
         AsyncMemcpyCollectiveOps*
         RaggedAllToAllTest*
     )
+
+    #  tsan tests appear to be flaky in rbe due to the heavy load
+    #  force them to run locally
+    RBE_OPTIONS+=(
+         --strategy=TestRunner=local
+    )
     shift
 fi
 
@@ -116,14 +123,15 @@ bazel --bazelrc=build_tools/rocm/rocm_xla.bazelrc test \
     --test_output=errors \
     --flaky_test_attempts=3 \
     --keep_going \
-    --test_strategy=exclusive \
+    --run_under=//build_tools/rocm:exclusive_local_wrapper \
+    --repo_env=TF_ROCM_AMDGPU_TARGETS=gfx908,gfx90a,gfx942,gfx1100 \
     --action_env=XLA_FLAGS=--xla_gpu_force_compilation_parallelism=16 \
     --action_env=XLA_FLAGS=--xla_gpu_enable_llvm_module_compilation_parallelism=true \
     --action_env=NCCL_MAX_NCHANNELS=1 \
     --test_filter=-$(IFS=: ; echo "${EXCLUDED_TESTS[*]}") \
     "${SANITIZER_ARGS[@]}" \
     "$@" \
-    --strategy=TestRunner=local # execute multigpu tests locally as there is no gpu exclusive protection on rbe
+    "${RBE_OPTIONS[@]}"
 
 # clean up bazel disk_cache
 bazel shutdown \

third_party/gpus/rocm/BUILD.tpl

@@ -140,7 +140,7 @@ cc_library(
     name = "rocm_rpath",
     linkopts = select({
         ":build_hermetic": [
-            "-Wl,-rpath,%{rocm_toolkit_path}/lib",
+            "-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib",
         ],
         ":multiple_rocm_paths": [
             "-Wl,-rpath=%{rocm_lib_paths}",
@@ -163,7 +163,7 @@ cc_library(
 
 cc_library(
     name = "rocm_hip",
-    srcs = glob(["%{rocm_root}/lib/libamdhip*.so"]),
+    srcs = glob(["%{rocm_root}/lib/libamdhip*.so*"]),
     hdrs = glob(["%{rocm_root}/include/hip/**"]),
     include_prefix = "rocm",
     includes = [
@@ -181,7 +181,10 @@ cc_library(
 # Used by jax_rocm_plugin to minimally link to hip runtime.
 cc_library(
     name = "hip_runtime",
-    srcs = glob(["%{rocm_root}/lib/libamdhip*.so"]),
+    srcs = glob([
+        "%{rocm_root}/lib/libamdhip*.so*",
+        "%{rocm_root}/lib/libamd_comgr.so*",
+    ]),
     hdrs = glob(["%{rocm_root}/include/hip/**"]),
     include_prefix = "rocm",
     includes = [
@@ -215,15 +218,16 @@ cc_library(
     ],
     # workaround to  bring tensile files to the same fs layout as expected in the lib
     # rocblas assumes that tensile files are located in ../roblas/libraries directory
-    linkopts = ["-Wl,-rpath,local_config_rocm/rocm/rocm_dis/lib"],
+    linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
     strip_include_prefix = "%{rocm_root}",
     visibility = ["//visibility:public"],
     deps = [":rocm_config"],
 )
 
 cc_library(
     name = "rocfft",
-    srcs = glob(["%{rocm_root}/lib/librocfft*.so*"]),
+    data = glob(["%{rocm_root}/lib/librocfft*.so*"]),
+    linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
     include_prefix = "rocm",
     includes = [
         "%{rocm_root}/include",
@@ -235,7 +239,8 @@ cc_library(
 
 cc_library(
     name = "hipfft",
-    srcs = glob(["%{rocm_root}/lib/libhipfft*.so*"]),
+    data = glob(["%{rocm_root}/lib/libhipfft*.so*"]),
+    linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
     include_prefix = "rocm",
     includes = [
         "%{rocm_root}/include",
@@ -266,7 +271,6 @@ miopen_libs = glob([
 
 cc_library(
     name = "miopen",
-    srcs = glob(["%{rocm_root}/lib/libMIOpen*.so*"]),
     hdrs = glob(["%{rocm_root}/include/miopen/**"]),
     data = select({
         ":build_hermetic": miopen_libs,
@@ -279,7 +283,7 @@ cc_library(
     ],
     # workaround to  bring miopen db files to the same fs layout as expected in the lib
     # rocblas assumes that miopen db files are located in ../share/miopen/db directory
-    linkopts = ["-Wl,-rpath,local_config_rocm/rocm/rocm_dis/lib"],
+    linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
     strip_include_prefix = "%{rocm_root}",
     visibility = ["//visibility:public"],
     deps = [":rocm_config"],
@@ -429,13 +433,13 @@ hipblas_libs = glob(["%{rocm_root}/lib/libhipblas.so*"])
 
 cc_library(
     name = "hipblas",
-    srcs = glob(["%{rocm_root}/lib/libhipblas.so*"]),
     hdrs = glob(["%{rocm_root}/include/hipblas/**"]),
     data = select({
         ":build_hermetic": hipblas_libs,
         ":multiple_rocm_paths": hipblas_libs,
         "//conditions:default": [],
     }),
+    linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
     include_prefix = "rocm",
     includes = [
         "%{rocm_root}/include/",
@@ -479,7 +483,7 @@ cc_library(
     ],
     # workaround to  bring tensile files to the same fs layout as expected in the lib
     # hibplatslt assumes that tensile files are located in ../hipblaslt/libraries directory
-    linkopts = ["-Wl,-rpath,local_config_rocm/rocm/rocm_dis/lib"],
+    linkopts = ["-Wl,-rpath,external/local_config_rocm/rocm/%{rocm_root}/lib"],
     strip_include_prefix = "%{rocm_root}",
     visibility = ["//visibility:public"],
     deps = [":rocm_config"],

xla/backends/gpu/runtime/BUILD

@@ -1812,6 +1812,8 @@ xla_test(
     backend_tags = {
         "gpu": [
             "multi_gpu_h100",
+            "multi_gpu",
+            "local",
             "no_oss",
         ],
     },

xla/tests/BUILD

@@ -1009,6 +1009,7 @@ xla_test(
     srcs = ["dot_operation_test.cc"],
     shard_count = 20,
     tags = [
+        "local",  # TODO: remove when remote execution is fixed
         "optonly",
     ],
     deps = [
@@ -1242,6 +1243,7 @@ xla_test(
     },
     shard_count = 20,
     tags = [
+        "local",
         "optonly",
     ],
     deps = [
@@ -1969,6 +1971,7 @@ cc_library(
     # This is set intentionally as to avoid the default behavior of the TSL
     # `cc_library` definition that is used in this file.
     compatible_with = [],
+    tags = ["local"],  # TODO: remove when remote execution is fixed
     deps = [
         ":client_library_test_base",
         ":hlo_test_base",
@@ -2622,6 +2625,7 @@ xla_test(
     backend_tags = {
         "gpu": [
             "multi_gpu",
+            "local",
             "no_oss",
         ],
         "cpu": [
@@ -2721,6 +2725,7 @@ xla_test(
     backend_tags = {
         "gpu": [
             "multi_gpu",
+            "local",
             "no_oss",
         ],
     },
@@ -3356,6 +3361,7 @@ xla_test(
     },
     shard_count = 50,
     tags = [
+        "local",  # TODO: remove when remote execution is fixed
         "test_migrated_to_hlo_runner_pjrt",
     ],
     deps = [