introduce rbe env var to decide when to use rbe flags

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -159,6 +159,14 @@ jobs:
 
     steps:
       - uses: actions/checkout@v3
+      - name: Enable RBE on platforms where its supported
+        run: |
+          os=$(uname -s | awk '{print tolower($0)}')
+          arch=$(uname -m)
+
+          if [[ ($os == "linux" || $os  =~ "msys_nt" ) && $arch == "x86_64" ]]; then
+            echo "JAXCI_BUILD_ARTIFACT_WITH_RBE=1" >> $GITHUB_ENV
+          fi
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

ci/build_artifacts.sh

@@ -40,7 +40,7 @@ arch=$(uname -m)
 
 # Adjust the values when running on Windows x86 to match the config in
 # .bazelrc
-if [[ $os =~ "msys_nt" ]] && [[ $arch == "x86_64" ]]; then
+if [[ $os =~ "msys_nt"  && $arch == "x86_64" ]]; then
   os="windows"
   arch="amd64"
 fi
@@ -54,26 +54,17 @@ if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
 
     # For bazel builds, use the "rbe_" config for Linux x86/Windows and "ci_" for other platforms
     bazelrc_config="${os}_${arch}"
-    if ( [[ "$os" == "linux" ]] && [[ "$arch" == "x86_64" ]] ) || [[ "$os" == "windows" ]]; then
-      bazelrc_config="rbe_$bazelrc_config"
+    if [[ "$JAXCI_BUILD_ARTIFACT_WITH_RBE" == 1 ]]; then
+      bazelrc_config="rbe_${bazelrc_config}"
     else
-      bazelrc_config="ci_$bazelrc_config"
+      bazelrc_config="ci_${bazelrc_config}"
     fi
 
-    # Build the jaxlib CPU artifact
-    if [[ "$artifact" == "jaxlib" ]]; then
-      python build/build.py build --wheels="jaxlib" --bazel_options=--config="$bazelrc_config" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+    if [[ ("$artifact" == "jax-cuda-plugin") || ("$artifact" == "jax-cuda-pjrt") ]]; then
+      bazelrc_config="${bazelrc_config}_cuda"
     fi
 
-    # Build the jax-cuda-plugin artifact
-    if [[ "$artifact" == "jax-cuda-plugin" ]]; then
-      python build/build.py build --wheels="jax-cuda-plugin" --bazel_options=--config="${bazelrc_config}_cuda" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
-    fi
-
-    # Build the jax-cuda-pjrt artifact
-    if [[ "$artifact" == "jax-cuda-pjrt" ]]; then
-      python build/build.py build --wheels="jax-cuda-pjrt" --bazel_options=--config="${bazelrc_config}_cuda" --verbose
-    fi
+     python build/build.py build --wheels="$artifact" --bazel_options=--config="$bazelrc_config" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
 
     # If building `jaxlib` or `jax-cuda-plugin` or `jax-cuda-pjrt` for Linux, we
     # run `auditwheel show` to verify manylinux compliance.

ci/envs/default.env

@@ -36,6 +36,9 @@ export JAXCI_XLA_COMMIT=${JAXCI_XLA_COMMIT:-}
 # Controls the location where the artifacts are stored.
 export JAXCI_OUTPUT_DIR="$(pwd)/dist"
 
+# When enabled, artifacts will be built with RBE.
+export JAXCI_BUILD_ARTIFACT_WITH_RBE=${JAXCI_BUILD_ARTIFACT_WITH_RBE:-0}
+
 # #############################################################################
 # Docker specific environment variables.
 # #############################################################################