Add new ml-build-rbe container to the configurations of RBE used with remote config.

Update .bazelrc file to use the new RBE config.

PiperOrigin-RevId: 706005323

Author: quoctruong

.bazelrc

@@ -533,9 +533,9 @@ build:rbe_linux_cpu --crosstool_top="@local_config_cuda//crosstool:toolchain"
 build:rbe_linux_cpu --extra_toolchains="@local_config_cuda//crosstool:toolchain-linux-x86_64"
 build:rbe_linux_cpu --repo_env=CC="/usr/lib/llvm-18/bin/clang"
 build:rbe_linux_cpu --repo_env=TF_SYSROOT="/dt9"
-build:rbe_linux_cpu --extra_execution_platforms="@sigbuild-r2.17-clang_config_platform//:platform"
-build:rbe_linux_cpu --host_platform="@sigbuild-r2.17-clang_config_platform//:platform"
-build:rbe_linux_cpu --platforms="@sigbuild-r2.17-clang_config_platform//:platform"
+build:rbe_linux_cpu --extra_execution_platforms="@ml_build_config_platform//:platform"
+build:rbe_linux_cpu --host_platform="@ml_build_config_platform//:platform"
+build:rbe_linux_cpu --platforms="@ml_build_config_platform//:platform"
 # This is needed for all Clang17 builds but must not be present in GCC builds.
 build:rbe_linux_cpu --copt=-Wno-error=unused-command-line-argument
 # This was added in clang-16 by https://reviews.llvm.org/D133574.

tools/toolchains/remote_config/configs.bzl

@@ -1,6 +1,6 @@
 """Configurations of RBE builds used with remote config."""
 
-load("//tools/toolchains/remote_config:rbe_config.bzl", "sigbuild_tf_configs", "tensorflow_local_config", "tensorflow_rbe_config", "tensorflow_rbe_win_config")
+load("//tools/toolchains/remote_config:rbe_config.bzl", "ml_build_rbe_config", "sigbuild_tf_configs", "tensorflow_local_config", "tensorflow_rbe_config", "tensorflow_rbe_win_config")
 
 def initialize_rbe_configs():
     tensorflow_local_config(
@@ -47,6 +47,11 @@ def initialize_rbe_configs():
         python_bin_path = "C:/Python37/python.exe",
     )
 
+    # The `ml-build-rbe` image is identical to the `ml-build` image except for the base image.
+    # The `ml-build`'s base image is a standard `ubuntu22.04` image.
+    # The `ml-build-rbe`'s base image is `nvidia/cuda:12.3.2-base-ubuntu22.04` which has nvidia driver installed.
+    ml_build_rbe_config("docker://us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-rbe@sha256:aaeb29799463729092c05f5ac8393113b3bb5d1ecf085f9f1f2016e3a1ece11c")
+
     # TF-Version-Specific SIG Build RBE Configs. The crosstool generated from these
     # configs are python-version-independent because they only care about the
     # tooling paths; the container mapping is useful only so that TF RBE users

tools/toolchains/remote_config/rbe_config.bzl

@@ -92,10 +92,24 @@ def _tensorflow_local_config(name):
         platform_constraint = "@%s_config_platform//:platform_constraint" % name,
     )
 
+def _ml_build_rbe_config(container_image):
+    exec_properties = {
+        "container-image": container_image,
+        "Pool": "default",
+    }
+
+    remote_platform_configure(
+        name = "ml_build_config_platform",
+        platform = "linux",
+        platform_exec_properties = exec_properties,
+    )
+
 tensorflow_rbe_config = _tensorflow_rbe_config
 tensorflow_rbe_win_config = _tensorflow_rbe_win_config
 tensorflow_local_config = _tensorflow_local_config
+ml_build_rbe_config = _ml_build_rbe_config
 
+# TODO(b/369382309): Remove this once ml_build_rbe_config is used everywhere.
 # Streamlined platform configuration for the SIG Build containers.
 # See //tensorflow/tools/tf_sig_build_dockerfiles
 # These containers do not support ROCm and all have CUDA.