Move around logic for when we are building jaxlib/cuda artifacts

Author: nitins17

ci/build_artifacts.sh

@@ -25,14 +25,7 @@
 # -o allexport: export all functions and variables to be available to subscripts
 set -exu -o history -o allexport
 
-# Store the comma-separated string in a variable
-artifacts="$1"
-
-# Replace commas with spaces
-artifacts=$(echo "$artifacts" | sed 's/,/ /g')
-
-# Create an array from the space-separated string
-artifacts=($artifacts)
+artifact="$1"
 
 # Source default JAXCI environment variables.
 source ci/envs/default.env
@@ -52,42 +45,40 @@ if [[ $os =~ "msys_nt" ]] && [[ $arch == "x86_64" ]]; then
   arch="amd64"
 fi
 
-# 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
-
 if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
   # Build the jax artifact
   if [[ "$artifact" == "jax" ]]; then
     python -m build --outdir $JAXCI_OUTPUT_DIR
-  fi
+  else
+    # 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"
+    else
+      bazelrc_config="ci_$bazelrc_config"
+    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 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-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 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 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-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
 
-  # 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
+    # If building `jaxlib` or `jax-cuda-plugin` or `jax-cuda-pjrt` for Linux, we
+    # run `auditwheel show` to verify manylinux compliance.
+    if  [[ "$os" == "linux" ]]; then
+      ./ci/utilities/run_auditwheel.sh
+    fi
+  fi 
 
 else
   echo "Error: Invalid artifact '$artifact'. Allowed values are: ${allowed_artifacts[@]}"