build-jax.sh: plumb through building + saving extra targets (#1515)

Example usage:
`--extra-targets=@xla//xla/tools/multihost_hlo_runner:hlo_runner_main`
to build the multi-host HLO runner.
Integrating `--extra-target-dest` into the script makes it useful in
conjunction with `--clean`.

Author: olupton

.github/container/Dockerfile.jax

@@ -1,6 +1,9 @@
 # syntax=docker/dockerfile:1-labs
 ARG BASE_IMAGE=ghcr.io/nvidia/jax:base
 ARG BUILD_PATH_JAXLIB=/opt/jaxlibs
+# Extra targets to build and copy outputs of, can be used for HLO tools. For example
+# @xla//xla/tools/multihost_hlo_runner:hlo_runner_main,@xla//xla/hlo/tools:hlo-opt
+ARG EXTRA_BAZEL_TARGETS=""
 ARG URLREF_JAX=https://github.com/google/jax.git#main
 ARG URLREF_XLA=https://github.com/openxla/xla.git#main
 ARG URLREF_FLAX=https://github.com/google/flax.git#main
@@ -28,6 +31,7 @@ ARG SRC_PATH_TRANSFORMER_ENGINE
 ARG SRC_PATH_XLA
 ARG BAZEL_CACHE
 ARG BUILD_PATH_JAXLIB
+ARG EXTRA_BAZEL_TARGETS
 ARG GIT_USER_NAME
 ARG GIT_USER_EMAIL
 
@@ -45,9 +49,11 @@ RUN ARCH="$(dpkg --print-architecture)" && \
     chmod +x /usr/local/bin/bazel
 # Populate ${BUILD_PATH_JAXLIB} with editable wheels; --no-install because
 # (a) this is the builder stage, and (b) pip-finalize.sh does the install
-RUN build-jax.sh \
+RUN mkdir -p /builder/extra-targets && build-jax.sh \
     --bazel-cache ${BAZEL_CACHE} \
     --build-path-jaxlib ${BUILD_PATH_JAXLIB} \
+    --extra-targets "${EXTRA_BAZEL_TARGETS}" \
+    --extra-target-dest /builder/extra-targets \
     --no-install \
     --src-path-jax ${SRC_PATH_JAX} \
     --src-path-xla ${SRC_PATH_XLA} \
@@ -91,6 +97,7 @@ COPY --from=builder ${BUILD_PATH_JAXLIB} ${BUILD_PATH_JAXLIB}
 COPY --from=builder ${SRC_PATH_JAX} ${SRC_PATH_JAX}
 COPY --from=builder ${SRC_PATH_XLA} ${SRC_PATH_XLA}
 COPY --from=builder /usr/local/bin/bazel /usr/local/bin/bazel
+COPY --from=builder /builder/extra-targets/* /usr/local/bin/
 # Preserve the versions of jax and xla
 COPY --from=builder /opt/manifest.d/git-clone.yaml /opt/manifest.d/git-clone.yaml
 ADD build-jax.sh build-te.sh local_cuda_arch pytest-xdist.sh test-jax.sh /usr/local/bin/

.github/container/build-jax.sh

@@ -51,6 +51,8 @@ usage() {
     echo "    --clean-only                   Do not build, just cleanup"
     echo "    --cpu-arch                     Target CPU architecture, e.g. amd64, arm64, etc."
     echo "    --debug                        Build in debug mode"
+    echo "    --extra-targets T1[,T2[...]    Extra bazel targets that will be built and copied to --extra-target-dest."
+    echo "    --extra-target-dest PATH       Where extra target output files will be copied to."
     echo "    -h, --help                     Print usage."
     echo "    --install                      Install the JAX wheels when build succeeds"
     echo "    --no-install                   Do not install the JAX wheels when build succeeds"
@@ -75,11 +77,13 @@ CLEANONLY=0
 CPU_ARCH="$(dpkg --print-architecture)"
 CUDA_COMPUTE_CAPABILITIES="local"
 DEBUG=0
+EXTRA_TARGETS=()
+EXTRA_TARGET_DEST=""
 INSTALL=1
 SRC_PATH_JAX="/opt/jax"
 SRC_PATH_XLA="/opt/xla"
 
-args=$(getopt -o h --long bazel-cache:,bazel-cache-namespace:,build-param:,build-path-jaxlib:,clean,cpu-arch:,debug,no-clean,clean-only,help,install,no-install,src-path-jax:,src-path-xla:,sm: -- "$@")
+args=$(getopt -o h --long bazel-cache:,bazel-cache-namespace:,build-param:,build-path-jaxlib:,clean,cpu-arch:,debug,extra-targets:,extra-target-dest:,no-clean,clean-only,help,install,no-install,src-path-jax:,src-path-xla:,sm: -- "$@")
 if [[ $? -ne 0 ]]; then
     exit 1
 fi
@@ -123,6 +127,14 @@ while [ : ]; do
             DEBUG=1
             shift 1
             ;;
+        --extra-targets)
+            IFS=',' read -r -a EXTRA_TARGETS <<< "$2"
+            shift 2
+            ;;
+        --extra-target-dest)
+            EXTRA_TARGET_DEST="$2"
+            shift 2
+            ;;
         -h | --help)
             usage 1
             ;;
@@ -202,13 +214,20 @@ print_var CPU_ARCH
 print_var CUDA_COMPUTE_CAPABILITIES
 print_var CUDA_MAJOR_VERSION
 print_var DEBUG
+print_var EXTRA_TARGETS
+print_var EXTRA_TARGET_DEST
 print_var INSTALL
 print_var PYTHON_VERSION
 print_var SRC_PATH_JAX
 print_var SRC_PATH_XLA
 
 echo "=================================================="
 
+if [[ -n "${EXTRA_TARGET_DEST}" && ! -d "${EXTRA_TARGET_DEST}" ]]; then
+    echo "You must pass a directory to --extra-target-dest"
+    exit 1
+fi
+
 set -x
 if [[ ${CLEANONLY} == 1 ]]; then
     clean
@@ -245,7 +264,21 @@ for component in jaxlib "jax-cuda${CUDA_MAJOR_VERSION}-pjrt" "jax-cuda${CUDA_MAJ
     # version, so nvidia-*-cu12 wheels disappear from the lock file
     sed -i "s|^${component}.*$|${component} @ file://${BUILD_PATH_JAXLIB}/${component//-/_}|" build/requirements.in
 done
-bazel run --config=cuda_libraries_from_stubs --verbose_failures=true //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION="${PYTHON_VERSION}"
+# Bazel args to avoid cache invalidation
+BAZEL_ARGS=(
+    --config=cuda_libraries_from_stubs
+    --repo_env=HERMETIC_PYTHON_VERSION="${PYTHON_VERSION}"
+)
+bazel run "${BAZEL_ARGS[@]}" --verbose_failures=true //build:requirements.update
+if (( "${#EXTRA_TARGETS[@]}" > 0 )); then
+    bazel build "${BAZEL_ARGS[@]}" --verbose_failures=true "${EXTRA_TARGETS[@]}"
+    if [[ -n "${EXTRA_TARGET_DEST}" ]]; then
+        for target in "${EXTRA_TARGETS[@]}"; do
+            output_files=$(bazel cquery "${BAZEL_ARGS[@]}" "${target}" --output files)
+            cp -v ${output_files} "${EXTRA_TARGET_DEST}"
+        done
+    fi
+fi
 popd
 
 ## Install the built packages