change how ci mode flags are parsed

Author: nitins17

build/build.py

@@ -110,7 +110,7 @@ def add_build_cuda_with_clang_argument(parser: argparse.ArgumentParser):
       action="store_true",
       help="""
         Should CUDA code be compiled using Clang? The default behavior is to
-        compile CUDA with NVCC. Ignored if --ci_mode is set, we always build
+        compile CUDA with NVCC. Ignored if --use_ci_bazelrc_flags is set, we always build
         CUDA with NVCC in CI builds.
         """,
   )
@@ -203,20 +203,11 @@ def add_global_arguments(parser: argparse.ArgumentParser):
 def add_artifact_subcommand_global_arguments(parser: argparse.ArgumentParser):
   """Adds all the global arguments that applies to the artifact subcommands."""
   parser.add_argument(
-      "--ci_mode",
+      "--use_ci_bazelrc_flags",
       action="store_true",
       help="""
         When set, the CLI will assume the build is being run in CI or CI like
-        environment and will use the "ci_" configs in the .bazelrc.
-        """,
-  )
-
-  parser.add_argument(
-      "--request_rbe",
-      action="store_true",
-      help="""
-        If set, the build will try to use RBE if it is supported for the
-        platform. Requires --ci_mode to be set.
+        environment and will use the "rbe_/ci_" configs in the .bazelrc.
         """,
   )
 
@@ -261,7 +252,7 @@ def add_artifact_subcommand_global_arguments(parser: argparse.ArgumentParser):
       type=str,
       default="",
       help="""
-        Path to the Clang binary to use. Ignored if --ci_mode is set as we use
+        Path to the Clang binary to use. Ignored if --use_ci_bazelrc_flags is set as we use
         a custom Clang toolchain in that case.
         """,
   )
@@ -440,16 +431,12 @@ async def main():
       wheel_cpus[args.target_cpu] if args.target_cpu is not None else arch
   )
 
-  # If running in CI mode, we use the "ci_"/"rbe_" configs in the .bazelrc.
+  # If running in CI, we use the "ci_"/"rbe_" configs in the .bazelrc.
   # These set a custom C++ Clang toolchain and the CUDA compiler to NVCC
-  # When not running in CI mode, we detect the path to Clang binary and pass it
+  # When not running in CI, we detect the path to Clang binary and pass it
   # to Bazel to use as the C++ compiler. NVCC is used as the CUDA compiler
   # unless the user explicitly sets --config=build_cuda_with_clang.
-  if args.ci_mode:
-    logging.debug(
-        "Running in CI mode. Run the CLI with --help for more details on what"
-        " this means."
-    )
+  if args.use_ci_bazelrc_flags:
     bazelrc_config = utils.get_bazelrc_config(
         os_name, arch, args.command, args.request_rbe
     )

build/tools/utils.py

@@ -203,7 +203,7 @@ def get_cuda_major_version():
   return None
 
 
-def get_bazelrc_config(os_name: str, arch: str, artifact: str, request_rbe: bool):
+def get_ci_bazelrc_config(os_name: str, arch: str, artifact: str):
   """Returns the bazelrc config for the given architecture and OS.
 
   Used in CI builds to retrieve either the "ci_"/"rbe_" configs from the
@@ -212,23 +212,13 @@ def get_bazelrc_config(os_name: str, arch: str, artifact: str, request_rbe: bool
 
   bazelrc_config = f"{os_name}_{arch}"
 
-  # If a build is requesting RBE, the CLI will use RBE if the host system supports
-  # it, otherwise it will use the "ci_" (non RBE) config.
-  if request_rbe:
-    if (os_name == "linux" and arch == "x86_64") or (
-        os_name == "windows" and arch == "amd64"
-    ):
-      bazelrc_config = "rbe_" + bazelrc_config
-    else:
-      logger.warning(
-          "RBE is not supported on %s_%s. Using the non RBE, ci_%s_%s, config"
-          " instead.",
-          os_name,
-          arch,
-          os_name,
-          arch,
-      )
-      bazelrc_config = "ci_" + bazelrc_config
+  # If building on Linux x86 or Windows, use the "rbe_" flags otherwise use
+  # the "ci_" (non-rbe) flags
+  if (os_name == "linux" and arch == "x86_64") or (
+      os_name == "windows" and arch == "amd64"
+  ):
+    
+    bazelrc_config = "rbe_" + bazelrc_config
   else:
     bazelrc_config = "ci_" + bazelrc_config
 

ci/build_artifacts.sh

@@ -25,17 +25,17 @@ fi
 
 # Build the jaxlib CPU artifact
 if [[ "$JAXCI_BUILD_JAXLIB" == 1 ]]; then
-  python build/build.py jaxlib --ci_mode --request_rbe --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+  python build/build.py jaxlib --use_ci_bazelrc_flags --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
 fi
 
 # Build the jax-cuda-plugin artifact
 if [[ "$JAXCI_BUILD_PLUGIN" == 1 ]]; then
-  python build/build.py jax-cuda-plugin --ci_mode --request_rbe --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+  python build/build.py jax-cuda-plugin --use_ci_bazelrc_flags --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
 fi
 
 # Build the jax-cuda-pjrt artifact
 if [[ "$JAXCI_BUILD_PJRT" == 1 ]]; then
-  python build/build.py jax-cuda-pjrt --ci_mode --request_rbe --verbose
+  python build/build.py jax-cuda-pjrt --use_ci_bazelrc_flags --verbose
 fi
 
 # After building `jaxlib`, `jaxcuda-plugin`, and `jax-cuda-pjrt`, we run