Add new CI scripts for building JAX artifacts

This commit introduces new CI scripts and environment files for building JAX artifacts. It makes use of the artifact envs inside the "ci/envs/build_artifacts" folder to control the build behavior. For e.g: for building jaxlib, we will need to run `./ci/build_artifacts.sh ./ci/envs/build_artifacts/jaxlib.env` from the JAX GitHub root.

PiperOrigin-RevId: 700104283

Author: nitins17

ci/build_artifacts.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+# Copyright 2024 The JAX Authors.
+##
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Build JAX artifacts.
+# 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"
+#
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -x: log all commands
+# -o history: record shell history
+# -o allexport: export all functions and variables to be available to subscripts
+set -exu -o history -o allexport
+
+artifact="$1"
+
+# Source default JAXCI environment variables.
+source ci/envs/default.env
+
+# Set up the build environment.
+source "ci/utilities/setup_build_environment.sh"
+
+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 config in
+# .bazelrc
+if [[ $os =~ "msys_nt"  && $arch == "x86_64" ]]; then
+  os="windows"
+  arch="amd64"
+fi
+
+if [[ "${allowed_artifacts[@]}" =~ "${artifact}" ]]; then
+
+  # Build the jax artifact
+  if [[ "$artifact" == "jax" ]]; then
+    python -m build --outdir $JAXCI_OUTPUT_DIR
+  else
+
+    # Figure out the bazelrc config to use. We will use one of the "rbe_"/"ci_"
+    # flags in the .bazelrc depending upon the platform we are building for.
+    bazelrc_config="${os}_${arch}"
+
+    # TODO(b/379903748): Add remote cache options for Linux and Windows.
+    if [[ "$JAXCI_BUILD_ARTIFACT_WITH_RBE" == 1 ]]; then
+      bazelrc_config="rbe_${bazelrc_config}"
+    else
+      bazelrc_config="ci_${bazelrc_config}"
+    fi
+
+    # Use the "_cuda" configs when building the CUDA artifacts.
+    if [[ ("$artifact" == "jax-cuda-plugin") || ("$artifact" == "jax-cuda-pjrt") ]]; then
+      bazelrc_config="${bazelrc_config}_cuda"
+    fi
+
+    # Build the artifact.
+    python build/build.py build --wheels="$artifact" --bazel_options=--config="$bazelrc_config" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+
+    # 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[@]}"
+  exit 1
+fi

ci/envs/default.env

@@ -34,4 +34,12 @@ export JAXCI_XLA_GIT_DIR=${JAXCI_XLA_GIT_DIR:-}
 export JAXCI_CLONE_MAIN_XLA=${JAXCI_CLONE_MAIN_XLA:-0}
 
 # Allows overriding the XLA commit that is used.
-export JAXCI_XLA_COMMIT=${JAXCI_XLA_COMMIT:-}
+export JAXCI_XLA_COMMIT=${JAXCI_XLA_COMMIT:-}
+
+# Controls the location where the artifacts are written to.
+export JAXCI_OUTPUT_DIR="$(pwd)/dist"
+
+# When enabled, artifacts will be built with RBE. Requires gcloud authentication
+# and only certain platforms support RBE. Therefore, this flag is enabled only
+# for CI builds where RBE is supported.
+export JAXCI_BUILD_ARTIFACT_WITH_RBE=${JAXCI_BUILD_ARTIFACT_WITH_RBE:-0}

ci/utilities/run_auditwheel.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Runs auditwheel to verify manylinux compatibility.
+
+# Get a list of all the wheels in the output directory. Only look for wheels
+# that need to be verified for manylinux compliance.
+WHEELS=$(find "$JAXCI_OUTPUT_DIR/" -type f \( -name "*jaxlib*whl" -o -name "*jax*cuda*pjrt*whl" -o -name "*jax*cuda*plugin*whl" \))
+
+if [[ -z "$WHEELS" ]]; then
+  echo "ERROR: No wheels found under $JAXCI_OUTPUT_DIR"
+  exit 1
+fi
+
+for wheel in $WHEELS; do
+    printf "\nRunning auditwheel on the following wheel:"
+    ls $wheel
+    OUTPUT_FULL=$(python -m auditwheel show $wheel)
+    # Remove the wheel name from the output to avoid false positives.
+    wheel_name=$(basename $wheel)
+    OUTPUT=${OUTPUT_FULL//${wheel_name}/}
+
+    # If a wheel is manylinux2014 compliant, `auditwheel show` will return the
+    # platform tag as manylinux_2_17. manylinux2014 is an alias for
+    # manylinux_2_17.
+    if echo "$OUTPUT" | grep -q "manylinux_2_17"; then
+        printf "\n$wheel_name is manylinux2014 compliant.\n"
+    else
+        echo "$OUTPUT_FULL"
+        printf "\n$wheel_name is NOT manylinux2014 compliant.\n"
+        exit 1
+    fi
+done