No bazel cache

NVIDIA · Aug 15, 2024 · 8a9c6e8 · 8a9c6e8
1 parent 04653f2
commit 8a9c6e8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax
@@ -48,7 +48,6 @@ RUN ARCH="$(dpkg --print-architecture)" && \
 # TODO: move this patch into the manifest
 ADD xla-arm64-neon.patch /opt
 RUN build-jax.sh \
-    --bazel-cache ${BAZEL_CACHE} \
     --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
@@ -256,18 +256,25 @@ popd
 
 ## Build jaxlib
 mkdir -p "${BUILD_PATH_JAXLIB}"
-bazel clean --expunge || true
 python /opt/jax/build/build.py 
+
+time python "${SRC_PATH_JAX}/build/build.py" \
     --editable \
     --use_clang \
     --enable_cuda \
     --build_gpu_plugin \
-    --gpu_plugin_cuda_version=12 \
-    --bazel_options=--repo_env=LOCAL_CUDA_PATH="/usr/local/cuda" \
+    --gpu_plugin_cuda_version=$TF_CUDA_MAJOR_VERSION \
+    --cuda_version=$TF_CUDA_VERSION \
+    --cudnn_version=$TF_CUDNN_VERSION \
+    --cuda_compute_capabilities=$TF_CUDA_COMPUTE_CAPABILITIES \
+    --enable_nccl=true \
     --bazel_options=--linkopt=-fuse-ld=lld \
-    --output_path=/opt/jaxlibs \
+    --bazel_options=--repo_env=LOCAL_CUDA_PATH="/usr/local/cuda" \
+    --bazel_options=--override_repository=xla=$SRC_PATH_XLA \
     --bazel_options=--repo_env=LOCAL_CUDNN_PATH="/opt/nvidia/cudnn" \
-    --bazel_options=--repo_env=LOCAL_NCCL_PATH="/usr/local/lib/python3.10/dist-packages/nvidia/nccl"
+    --bazel_options=--repo_env=LOCAL_NCCL_PATH="/usr/local/lib/python3.10/dist-packages/nvidia/nccl" \
+    --output_path=${BUILD_PATH_JAXLIB} \
+    $BUILD_PARAM
 
 # Make sure that JAX depends on the local jaxlib installation
 # https://jax.readthedocs.io/en/latest/developer.html#specifying-dependencies-on-local-wheels