Change to building a single artifact per execution

Author: nitins17

ci/build_artifacts.sh

@@ -14,10 +14,9 @@
 # limitations under the License.
 # ==============================================================================
 # Build JAX artifacts. 
-# Usage: ./ci/build_artifacts.sh "<comma-separated artifact values>"
+# Usage: ./ci/build_artifacts.sh "<artifact>"
 # Supported artifact values are: jax, jaxlib, jax-cuda-plugin, jax-cuda-pjrt
 # E.g: ./ci/build_artifacts.sh "jax" or ./ci/build_artifacts.sh "jaxlib"
-# Multiple artifacts builds are permitted. E.g: ./ci/build_artifacts.sh "jax,jaxlib"
 #
 # -e: abort script if one command fails
 # -u: error if undefined variable used
@@ -46,53 +45,51 @@ 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
+# Adjust the values when running on Windows x86 to match the config 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"
+# For bazel builds, use the "rbe_" config for Linux x86/Windows and "ci_" for other platforms
+ if [[ "$artifact" != "jax" ]]; then
+  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
 fi
 
-for artifact in "${artifacts[@]}"; do
-
-  if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
-    # Build the jax artifact
-    if [[ "$artifact" == "jax" ]]; then
-      python -m build --outdir $JAXCI_OUTPUT_DIR
-    fi
-
-    # Build the jaxlib CPU artifact
-    if [[ "$artifact" == "jaxlib" ]]; then
-      python build/build.py build --wheels="jaxlib" --bazel_build_options=--config=$bazelrc_config --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
-    fi
+if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
+  # Build the jax artifact
+  if [[ "$artifact" == "jax" ]]; then
+    python -m build --outdir $JAXCI_OUTPUT_DIR
+  fi
 
-    # Build the jax-cuda-plugin artifact
-    if [[ "$artifact" == "jax-cuda-plugin" ]]; then
-      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 jaxlib CPU artifact
+  if [[ "$artifact" == "jaxlib" ]]; then
+    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-pjrt artifact
-    if [[ "$artifact" == "jax-cuda-pjrt" ]]; then
-      python build/build.py build --wheels="jax-cuda-pjrt" --bazel_build_options=--config="${bazelrc_config}_cuda" --verbose
-    fi
+  # Build the jax-cuda-plugin artifact
+  if [[ "$artifact" == "jax-cuda-plugin" ]]; then
+    python build/build.py build --wheels="jax-cuda-plugin" --bazel_build_options=--config="${bazelrc_config}_cuda" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+  fi
 
-    # If building `jaxlib` or `jax-cuda-plugin` or `jax-cuda-pjrt` for Linux, we
-    # run `auditwheel show` to verify manylinux compliance.
-    if  [[ "$os" == "linux" ]] && [[ "$artifact" != "jax" ]]; then
-      ./ci/utilities/run_auditwheel.sh
-    fi
+  # Build the jax-cuda-pjrt artifact
+  if [[ "$artifact" == "jax-cuda-pjrt" ]]; then
+    python build/build.py build --wheels="jax-cuda-pjrt" --bazel_build_options=--config="${bazelrc_config}_cuda" --verbose
+  fi
 
-  else
-    echo "Error: Invalid artifact '$artifact'. Allowed values are: ${allowed_artifacts[@]}"
-    exit 1
+  # If building `jaxlib` or `jax-cuda-plugin` or `jax-cuda-pjrt` for Linux, we
+  # run `auditwheel show` to verify manylinux compliance.
+  if  [[ "$os" == "linux" ]] && [[ "$artifact" != "jax" ]]; then
+    ./ci/utilities/run_auditwheel.sh
   fi
 
-done
+else
+  echo "Error: Invalid artifact '$artifact'. Allowed values are: ${allowed_artifacts[@]}"
+  exit 1
+fi