detect bazelrc config in the build artifacts script

Author: nitins17

build/build.py

@@ -327,8 +327,6 @@ async def main():
   add_global_arguments(build_artifact_parser)
 
   arch = platform.machine()
-  # Switch to lower case to match the case for the "ci_"/"rbe_" configs in the
-  # .bazelrc.
   os_name = platform.system().lower()
 
   args = parser.parse_args()

ci/build_artifacts.sh

@@ -41,9 +41,26 @@ source ci/envs/default.env
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
-os=$(uname -s | awk '{print tolower($0)}')
 allowed_artifacts=("jax" "jaxlib" "jax-cuda-plugin" "jax-cuda-pjrt")
 
+os=$(uname -s | awk '{print tolower($0)}')
+arch=$(uname -m)
+
+# Adjust the values when running on Windows x86 to match the case in
+# .bazelrc
+if [[ $os =~ "msys_nt" ]] && [[ $arch == "x86_64" ]]; then
+  os="windows"
+  arch="amd64"
+fi
+
+# Use "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"
+else
+  bazelrc_config="ci_$bazelrc_config"
+fi
+
 for artifact in "${artifacts[@]}"; do
 
   if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
@@ -54,17 +71,17 @@ for artifact in "${artifacts[@]}"; do
 
     # Build the jaxlib CPU artifact
     if [[ "$artifact" == "jaxlib" ]]; then
-      python build/build.py build_artifacts --wheel_list="jaxlib" --use_ci_bazelrc_flags --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+      python build/build.py build --wheels="jaxlib" --bazel_build_options=--config=$bazelrc_config --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
     fi
 
     # Build the jax-cuda-plugin artifact
     if [[ "$artifact" == "jax-cuda-plugin" ]]; then
-      python build/build.py build_artifacts --wheel_list="jax-cuda-plugin" --use_ci_bazelrc_flags --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+      python build/build.py build --wheels="jax-cuda-plugin" --bazel_build_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_artifacts --wheel_list="jax-cuda-pjrt" --use_ci_bazelrc_flags --verbose
+      python build/build.py build --wheels="jax-cuda-pjrt" --bazel_build_options=--config="${bazelrc_config}_cuda" --verbose
     fi
 
     # If building `jaxlib` or `jax-cuda-plugin` or `jax-cuda-pjrt` for Linux, we