Make rbe image parametrized (#425)

Author: alekstheod

build_tools/rocm/platform/BUILD

@@ -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"
-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 --host_platform="@local_config_rocm//rocm:linux_x64"
+build:rocm_rbe --extra_execution_platforms="@local_config_rocm//rocm:linux_x64"
+build:rocm_rbe --platforms="@local_config_rocm//rocm: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"

build_tools/rocm/platform/linux_x64/BUILD

@@ -584,3 +584,17 @@ filegroup(
     srcs = glob(["%{rocm_root}/**"]),
     visibility = ["//visibility:public"],
 )
+
+platform(
+    name = "linux_x64",
+    constraint_values = [
+        "@platforms//os:linux",
+        "@platforms//cpu:x86_64",
+        "@bazel_tools//tools/cpp:clang",
+    ],
+
+    exec_properties = {
+        "container-image": "docker://%{rocm_rbe_docker_image}",
+        "OSFamily": "Linux",
+    },
+)

build_tools/rocm/rocm_xla.bazelrc

@@ -52,6 +52,10 @@ _DISTRIBUTION_PATH = "rocm/rocm_dist"
 _OS = "OS"
 _ROCM_VERSION = "ROCM_VERSION"
 
+_TF_ROCM_RBE_DOCKER_IMAGE = "TF_ROCM_RBE_DOCKER_IMAGE"
+# rocm/tensorflow-build:latest-jammy-python3.11-rocm7.0.2
+_DEFAULT_TF_ROCM_RBE_DOCKER_IMAGE = "rocm/tensorflow-build@sha256:a2672ff2510b369b4a5f034272a518dc93c2e492894e3befaeef19649632ccaa"
+
 _DEFAULT_ROCM_TOOLKIT_PATH = "/opt/rocm"
 _TF_ROCM_MULTIPLE_PATHS = "TF_ROCM_MULTIPLE_PATHS"
 _LLVM_PATH = "LLVM_PATH"
@@ -655,6 +659,7 @@ def _create_local_rocm_repository(repository_ctx):
     repository_dict = {
         "%{rocm_root}": rocm_toolkit_path,
         "%{rocm_toolkit_path}": str(repository_ctx.path(rocm_config.rocm_toolkit_path)),
+        "%{rocm_rbe_docker_image}": repository_ctx.os.environ.get(_TF_ROCM_RBE_DOCKER_IMAGE, _DEFAULT_TF_ROCM_RBE_DOCKER_IMAGE),
     }
 
     is_rocm_clang = _use_rocm_clang(repository_ctx)
@@ -873,6 +878,8 @@ _ENVIRONS = [
     _TF_ROCM_AMDGPU_TARGETS,
     _OS,
     _ROCM_VERSION,
+    _TF_ROCM_RBE_DOCKER_IMAGE,
+    _TF_ROCM_MULTIPLE_PATHS,
 ]
 
 remote_rocm_configure = repository_rule(