build-jax.sh: --bazel-cache-namespace option (#1141)

When building the JAX container, set this to the name of the base
container to provide some isolation within the shared cache between
builds using different base containers. Bazel does not track system
dependencies.

Author: olupton

.github/container/Dockerfile.base

@@ -8,27 +8,29 @@ ARG CLANG_VERSION=18
 ## Obtain GCP's NCCL TCPx plugin
 ###############################################################################
 
-FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx:v3.1.10 as tcpx-installer-amd64
+FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx:v3.1.10 AS tcpx-installer-amd64
 
 # make a stub arm64 container because GCP does not provide an arm64 version of the plugin
-FROM ubuntu as tcpx-installer-arm64
+FROM ubuntu AS tcpx-installer-arm64
 RUN <<"OUTEREOF" bash -ex
 mkdir -p /scripts /var/lib/tcpx/lib64
 echo '#!/bin/bash' > /scripts/container_entry.sh
 chmod +x /scripts/container_entry.sh
 OUTEREOF
 
-FROM tcpx-installer-${TARGETARCH} as tcpx-installer
+FROM tcpx-installer-${TARGETARCH} AS tcpx-installer
 RUN /scripts/container_entry.sh install
 
 ###############################################################################
 ## Build base image
 ###############################################################################
 
 FROM ${BASE_IMAGE}
+ARG BASE_IMAGE
 ARG GIT_USER_EMAIL
 ARG GIT_USER_NAME
 ARG CLANG_VERSION
+ENV CUDA_BASE_IMAGE=${BASE_IMAGE}
 
 ###############################################################################
 ## Install Python and essential tools

.github/container/build-jax.sh

@@ -29,6 +29,7 @@ usage() {
     echo ""
     echo "    OPTIONS                        DESCRIPTION"
     echo "    --bazel-cache URI              Path for local bazel cache or URL of remote bazel cache"
+    echo "    --bazel-cache-namespace NAME   Namespace for bazel cache content"
     echo "    --build-param PARAM            Param passed to the jaxlib build command. Can be passed many times."
     echo "    --build-path-jaxlib PATH       Editable install prefix for jaxlib and plugins"
     echo "    --clean                        Delete local configuration and bazel cache"
@@ -51,6 +52,7 @@ usage() {
 
 # Set defaults
 BAZEL_CACHE=""
+BAZEL_CACHE_NAMESPACE="jax${CUDA_BASE_IMAGE:+:}${CUDA_BASE_IMAGE}"
 BUILD_PATH_JAXLIB="/opt/jaxlibs"
 BUILD_PARAM=""
 CLEAN=0
@@ -64,7 +66,7 @@ SRC_PATH_JAX="/opt/jax"
 SRC_PATH_XLA="/opt/xla"
 XLA_ARM64_PATCH_LIST=""
 
-args=$(getopt -o h --long bazel-cache:,build-param:,build-path-jaxlib:,clean,cpu-arch:,debug,jaxlib_only,no-clean,clean-only,dry,help,src-path-jax:,src-path-xla:,sm:,xla-arm64-patch: -- "$@")
+args=$(getopt -o h --long bazel-cache:,bazel-cache-namespace:,build-param:,build-path-jaxlib:,clean,cpu-arch:,debug,jaxlib_only,no-clean,clean-only,dry,help,src-path-jax:,src-path-xla:,sm:,xla-arm64-patch: -- "$@")
 if [[ $? -ne 0 ]]; then
     exit 1
 fi
@@ -76,6 +78,10 @@ while [ : ]; do
             BAZEL_CACHE=$2
             shift 2
             ;;
+        --bazel-cache-namespace)
+            BAZEL_CACHE_NAMESPACE=$2
+            shift 2
+            ;;
         --build-param)
             BUILD_PARAM="$BUILD_PARAM $2"
             shift 2
@@ -196,6 +202,9 @@ fi
 if [[ "${BAZEL_CACHE}" == http://* ]] || \
    [[ "${BAZEL_CACHE}" == grpc://* ]]; then
     BUILD_PARAM="${BUILD_PARAM} --bazel_options=--remote_cache=${BAZEL_CACHE}"
+    if [[ -n "${BAZEL_CACHE_NAMESPACE}" ]]; then
+        BUILD_PARAM="${BUILD_PARAM} --bazel_options=--remote_instance_name=${BAZEL_CACHE_NAMESPACE}"
+    fi
 elif [[ ! -z "${BAZEL_CACHE}" ]] ; then
     BUILD_PARAM="${BUILD_PARAM} --bazel_options=--disk_cache=${BAZEL_CACHE}"
 fi