Build JAX artifacts only in CPU pools for RBE configurations.

PiperOrigin-RevId: 839788470

Author: ybaturina

.bazelrc

@@ -396,6 +396,9 @@ common:rbe --spawn_strategy=remote,worker,standalone,local
 common:rbe --remote_download_toplevel
 test:rbe --test_env=USER=anon
 
+common:rbe_cpu_pool --repo_env=REMOTE_GPU_TESTING=0
+common:rbe_gpu_pool --repo_env=REMOTE_GPU_TESTING=1
+
 # RBE configs for Linux x86
 # Set the remote worker pool
 common:rbe_linux_x86_64_base --remote_instance_name=projects/tensorflow-testing/instances/default_instance
@@ -416,7 +419,7 @@ common:rbe_linux_x86_64 --config=rbe_linux_x86_64_base
 common:rbe_linux_x86_64 --config=ci_linux_x86_64
 
 common:rbe_linux_x86_64_cuda_common --config=rbe_linux_x86_64_base
-common:rbe_linux_x86_64_cuda_common --repo_env=REMOTE_GPU_TESTING=1
+common:rbe_linux_x86_64_cuda_common --config=rbe_gpu_pool
 
 common:rbe_linux_x86_64_cuda12 --config=rbe_linux_x86_64_cuda_common
 common:rbe_linux_x86_64_cuda12 --config=ci_linux_x86_64_cuda12

ci/build_artifacts.sh

@@ -109,10 +109,11 @@ if [[ "${allowed_artifacts[@]}" =~ "${artifact}" ]]; then
   # Build the artifact.
   python build/build.py build --wheels="$artifact" \
     --bazel_options=--config="$bazelrc_config" $bazel_remote_cache \
+    --bazel_options=--config=rbe_cpu_pool \
     --bazel_startup_options="$bazel_startup_options" \
     --python_version=$JAXCI_HERMETIC_PYTHON_VERSION \
     $cuda_version_flag \
-    --verbose --detailed_timestamped_log --use_new_wheel_build_rule \
+    --verbose --detailed_timestamped_log \
     --output_path="$JAXCI_OUTPUT_DIR" \
     $artifact_tag_flags
 
@@ -121,10 +122,11 @@ if [[ "${allowed_artifacts[@]}" =~ "${artifact}" ]]; then
   if [[ "$JAXCI_ARTIFACT_TYPE" == "release" ]]; then
     python build/build.py build --wheels="$artifact" \
       --bazel_options=--config="$bazelrc_config" $bazel_remote_cache \
+      --bazel_options=--config=rbe_cpu_pool \
       --bazel_startup_options="$bazel_startup_options" \
       --python_version=$JAXCI_HERMETIC_PYTHON_VERSION \
       $cuda_version_flag \
-      --verbose --detailed_timestamped_log --use_new_wheel_build_rule \
+      --verbose --detailed_timestamped_log \
       --output_path="$JAXCI_OUTPUT_DIR" \
       $artifact_tag_flags --bazel_options=--repo_env=ML_WHEEL_VERSION_SUFFIX="$JAXCI_WHEEL_RC_VERSION"
   fi