update ci scripts

Author: nitins17

build/build.py

@@ -45,9 +45,9 @@
 EPILOG = """
 From the root directory of the JAX repository, run
   `python build/build.py build --wheels=<list of JAX wheels>` to build JAX
-  artifacts. 
-  
-  Multiple wheels can be built at the same time.
+  artifacts.
+
+  Multiple wheels can be built with a single invocation of the CLI.
   E.g. python build/build.py build --wheels=jaxlib,jax-cuda-plugin
 
   To update the requirements_lock.txt files, run
@@ -113,13 +113,13 @@ def add_global_arguments(parser: argparse.ArgumentParser):
   )
 
   bazel_group.add_argument(
-      "--bazel_build_options",
+      "--bazel_options",
       action="append",
       default=[],
       help="""
         Additional build options to pass to Bazel, can be specified multiple
         times to pass multiple options.
-        E.g. --bazel_build_options='--local_resources=HOST_CPUS'
+        E.g. --bazel_options='--local_resources=HOST_CPUS'
         """,
   )
 
@@ -369,11 +369,11 @@ async def main():
   # Requirements update subcommand execution
   if args.command == "requirements_update":
     requirements_command = copy.deepcopy(bazel_command_base)
-    if args.bazel_build_options:
+    if args.bazel_options:
       logging.debug(
-          "Using additional build options: %s", args.bazel_build_options
+          "Using additional build options: %s", args.bazel_options
       )
-      for option in args.bazel_build_options:
+      for option in args.bazel_options:
         requirements_command.append(option)
 
     if args.nightly_update:
@@ -511,11 +511,11 @@ async def main():
 
     # Append additional build options at the end to override any options set in
     # .bazelrc or above.
-    if args.bazel_build_options:
+    if args.bazel_options:
       logging.debug(
-          "Additional Bazel build options: %s", args.bazel_build_options
+          "Additional Bazel build options: %s", args.bazel_options
       )
-      for option in args.bazel_build_options:
+      for option in args.bazel_options:
         wheel_build_command.append(option)
 
     if args.configure_only:

ci/build_artifacts.sh

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-# Build JAX artifacts. 
+# 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"
@@ -46,10 +46,12 @@ if [[ $os =~ "msys_nt" ]] && [[ $arch == "x86_64" ]]; then
 fi
 
 if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
+
   # Build the jax artifact
   if [[ "$artifact" == "jax" ]]; then
     python -m build --outdir $JAXCI_OUTPUT_DIR
   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
@@ -60,27 +62,28 @@ if [[ " ${allowed_artifacts[@]} " =~ " ${artifact} " ]]; then
 
     # 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
+      python build/build.py build --wheels="jaxlib" --bazel_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
+      python build/build.py build --wheels="jax-cuda-plugin" --bazel_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
+      python build/build.py build --wheels="jax-cuda-pjrt" --bazel_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" ]]; then
       ./ci/utilities/run_auditwheel.sh
     fi
-  fi 
+
+  fi
 
 else
-  echo "Error: Invalid artifact '$artifact'. Allowed values are: ${allowed_artifacts[@]}"
+  echo "Error: Invalid artifact: $artifact. Allowed values are: ${allowed_artifacts[@]}"
   exit 1
 fi

ci/envs/default.env

@@ -41,7 +41,6 @@ export JAXCI_OUTPUT_DIR="$(pwd)/dist"
 # #############################################################################
 
 # Docker specifc environment variables. Used by `run_docker_container.sh`
-export JAXCI_RUN_DOCKER_CONTAINER=${JAXCI_RUN_DOCKER_CONTAINER:-1}
 export JAXCI_DOCKER_WORK_DIR="/jax"
 export JAXCI_DOCKER_IMAGE=""
 export JAXCI_DOCKER_ARGS=""

ci/envs/docker.env

@@ -21,9 +21,6 @@ source ci/envs/default.env
 os=$(uname -s | awk '{print tolower($0)}')
 arch=$(uname -m)
 
-export JAXCI_DOCKER_ARGS=""
-export JAXCI_DOCKER_IMAGE=""
-
 # TODO: Set GPU Docker args and GPU Docker images
 # Linux x86 specifc settings
 if [[ $os == "linux" ]] && [[ $arch == "x86_64" ]]; then

ci/utilities/run_auditwheel.sh

@@ -14,7 +14,7 @@
 # limitations under the License.
 # ==============================================================================
 #
-# Runs auditwheel to ensure manylinux compatibility.
+# 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.

ci/utilities/run_docker_container.sh

@@ -49,8 +49,8 @@ if ! docker container inspect jax >/dev/null 2>&1 ; then
   # Start the container. `user_set_jaxci_envs` is read after `jax_ci_envs` to
   # allow the user to override any environment variables set by JAXCI_ENV_FILE.
   docker run $JAXCI_DOCKER_ARGS --name jax \
-      -w /jax -itd --rm \
-      -v "$JAXCI_JAX_GIT_DIR:/jax" \
+      -w "$JAXCI_DOCKER_WORK_DIR" -itd --rm \
+      -v "$JAXCI_JAX_GIT_DIR:$JAXCI_DOCKER_WORK_DIR" \
       "$JAXCI_DOCKER_IMAGE" \
     bash