Default cache namespace; only pass it when using remote cache

NVIDIA · Nov 5, 2024 · 83dbdca · 83dbdca
1 parent ddb39a1
commit 83dbdca
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax
@@ -49,7 +49,6 @@ RUN ARCH="$(dpkg --print-architecture)" && \
 ADD xla-arm64-neon.patch /opt
 RUN build-jax.sh \
     --bazel-cache ${BAZEL_CACHE} \
-    --bazel-cache-namespace "${CUDA_BASE_IMAGE:-jax}" \
     --build-path-jaxlib ${BUILD_PATH_JAXLIB} \
     --src-path-jax ${SRC_PATH_JAX} \
     --src-path-xla ${SRC_PATH_XLA} \

diff --git a/.github/container/build-jax.sh b/.github/container/build-jax.sh
@@ -52,7 +52,7 @@ usage() {
 
 # Set defaults
 BAZEL_CACHE=""
-BAZEL_CACHE_NAMESPACE=""
+BAZEL_CACHE_NAMESPACE="jax${CUDA_BASE_IMAGE:+:}${CUDA_BASE_IMAGE}"
 BUILD_PATH_JAXLIB="/opt/jaxlibs"
 BUILD_PARAM=""
 CLEAN=0
@@ -202,14 +202,13 @@ 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
 
-if [[ -n "${BAZEL_CACHE_NAMESPACE}" ]]; then
-    BUILD_PARAM="${BUILD_PARAM} --bazel_options=--remote_instance_name=${BAZEL_CACHE_NAMESPACE}"
-fi
-
 if [[ "$DEBUG" == "1" ]]; then
     BUILD_PARAM="${BUILD_PARAM} --bazel_options=-c --bazel_options=dbg --bazel_options=--strip=never --bazel_options=--cxxopt=-g --bazel_options=--cxxopt=-O0"
 fi