Improve build-te for build-jax step in CI

.github/actions/build-container/action.yml

@@ -66,6 +66,12 @@ inputs:
     description: "Optional named build context to source from a runner registry for the final image handoff"
     required: false
     default: ""
+  mealkit-secret-envs:
+    description: >-
+      Optional newline-delimited BuildKit secret mappings in SECRET_ID=ENV_VAR
+      form. These secrets are exposed only to the mealkit build.
+    required: false
+    default: ""
   require-nsys-jax-version-artifact:
     description: "Whether this container build requires the nsys-jax version.py artifact"
     required: false
@@ -114,8 +120,6 @@ runs:
         keep-state: true
         # TODO: check whether we need the oci.worker specification
 
-
-
     - name: Download nsys-jax version.py
       if: ${{ inputs.require-nsys-jax-version-artifact == 'true' && inputs.CONTAINER_NAME != 'jio' }}
       uses: actions/download-artifact@v8
@@ -160,6 +164,7 @@ runs:
         ssh: default
         secret-files: |
           "SSH_KNOWN_HOSTS=${{ steps.setup-ssh.outputs.known-hosts-file }}"
+        secret-envs: ${{ inputs.mealkit-secret-envs }}
         allow: ${{ inputs.docker-build-allow }}
         network: ${{ inputs.docker-build-network }}
         build-args: |

.github/container/Dockerfile.jax

@@ -27,6 +27,7 @@ ARG BUILD_DATE
 ###############################################################################
 
 FROM ${BASE_IMAGE} AS builder
+ARG TARGETARCH
 ARG URLREF_JAX
 ARG URLREF_TRANSFORMER_ENGINE
 ARG URLREF_XLA
@@ -39,6 +40,8 @@ ARG EXTRA_BAZEL_TARGETS
 ARG EXTRA_BUILD_JAX_ARGS
 ARG GIT_USER_NAME
 ARG GIT_USER_EMAIL
+ARG ENABLE_NVTE_SCCACHE=0
+ARG SCCACHE_REGION=""
 
 RUN --mount=type=ssh \
     --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
@@ -47,7 +50,7 @@ RUN --mount=type=ssh \
     git-clone.sh ${URLREF_XLA} ${SRC_PATH_XLA}
 EOF
 
-ADD build-jax.sh build-te.sh local_cuda_arch /usr/local/bin/
+ADD build-jax.sh build-te.sh build-te-with-sccache.sh install-compiler-cache.sh local_cuda_arch /usr/local/bin/
 # Install bazelisk
 RUN ARCH="$(dpkg --print-architecture)" && \
     wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-${ARCH} && \
@@ -70,13 +73,30 @@ RUN mkdir -p /builder/extra-targets/{bin,python} && \
 RUN --mount=type=ssh \
     --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
     git-clone.sh ${URLREF_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE}
+
+# Install the selected compiler cache. The binary remains in the builder
+# stage and is not copied into the mealkit or final image.
+RUN <<"EOF" bash -ex
+    if [[ "${ENABLE_NVTE_SCCACHE}" == "1" ]]; then
+        : "${SCCACHE_REGION:?ENABLE_NVTE_SCCACHE=1 requires SCCACHE_REGION}"
+        NVTE_CCACHE_BIN=sccache install-compiler-cache.sh sccache
+    elif [[ "${ENABLE_NVTE_SCCACHE}" != "0" ]]; then
+        echo "ENABLE_NVTE_SCCACHE must be 0 or 1" >&2
+        exit 1
+    fi
+EOF
+
 # Populate ${SRC_PATH_TRANSFORMER_ENGINE}/dist with [a] .whl file(s); --no-install
-# because (a) this is the builder stage, and (b) pip-finalize.sh does the install
-RUN build-te.sh \
+# because (a) this is the builder stage, and (b) pip-finalize.sh does the install.
+RUN --mount=type=secret,id=AWS_ACCESS_KEY_ID,required=false \
+    --mount=type=secret,id=AWS_SECRET_ACCESS_KEY,required=false \
+    --mount=type=secret,id=AWS_SESSION_TOKEN,required=false \
+    --mount=type=secret,id=SCCACHE_BUCKET,required=false \
+    build-te-with-sccache.sh \
     --clean \
     --no-install \
-    --src-path-te ${SRC_PATH_TRANSFORMER_ENGINE} \
-    --src-path-xla ${SRC_PATH_XLA}
+    --src-path-te "${SRC_PATH_TRANSFORMER_ENGINE}" \
+    --src-path-xla "${SRC_PATH_XLA}"
 
 ###############################################################################
 ## Pack jaxlib wheel and various source dirs into a pre-installation image
@@ -107,7 +127,7 @@ COPY --from=builder /builder/extra-targets/python/*.so /usr/lib/python3/dist-pac
 
 # 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/
+ADD build-jax.sh build-te.sh install-compiler-cache.sh local_cuda_arch pytest-xdist.sh test-jax.sh /usr/local/bin/
 
 RUN mkdir -p /opt/pip-tools.d
 

.github/container/build-te-with-sccache.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+# Never trace this script: it handles credentials mounted by BuildKit.
+set +x
+
+compiler_cache_args=()
+case "${ENABLE_NVTE_SCCACHE:-0}" in
+    1)
+        required_sccache_secrets=(
+            AWS_ACCESS_KEY_ID
+            AWS_SECRET_ACCESS_KEY
+            AWS_SESSION_TOKEN
+            SCCACHE_BUCKET
+        )
+        for secret_name in "${required_sccache_secrets[@]}"; do
+            if [[ ! -s "/run/secrets/${secret_name}" ]]; then
+                echo "ENABLE_NVTE_SCCACHE=1 requires the ${secret_name} BuildKit secret" >&2
+                exit 1
+            fi
+        done
+
+        for secret_name in "${required_sccache_secrets[@]}"; do
+            printf -v "${secret_name}" '%s' "$(<"/run/secrets/${secret_name}")"
+            export "${secret_name}"
+        done
+
+        : "${SCCACHE_REGION:?ENABLE_NVTE_SCCACHE=1 requires SCCACHE_REGION}"
+        : "${TARGETARCH:?BuildKit did not provide TARGETARCH}"
+        export NVTE_CCACHE_BIN=sccache
+        export SCCACHE_REGION
+        export SCCACHE_S3_KEY_PREFIX="sccache/${TARGETARCH}"
+        export SCCACHE_S3_USE_SSL=true
+        compiler_cache_args+=(--ccache)
+        ;;
+    0)
+        ;;
+    *)
+        echo "ENABLE_NVTE_SCCACHE must be 0 or 1" >&2
+        exit 1
+        ;;
+esac
+
+exec build-te.sh "${compiler_cache_args[@]}" "$@"

.github/container/build-te.sh

@@ -11,7 +11,7 @@ usage() {
     echo "    OPTIONS                        DESCRIPTION"
     echo "    --clean                        Clear build caches under --src-path-te."
     echo "    -h, --help                     Print usage."
-    echo "    --ccache                       Use ccache to build TransformerEngine. It will be installed, but the caller should configure it."
+    echo "    --ccache                       Use a compiler cache to build TransformerEngine."
     echo "    --no-install                   Only build a wheel; do not install."
     echo "    --src-path-te                  Path to TransformerEngine source code."
     echo "    --src-path-xla                 Path to XLA source code."
@@ -148,25 +148,78 @@ subprocess.run(
     + [f"nvidia-cudnn-frontend=={os.environ['CUDNN_FRONTEND_VERSION']}"]
 )
 EOF
+# Transformer Engine adds NVTE_CCACHE_BIN as both its C++ and CUDA CMake
+# compiler launcher when NVTE_USE_CCACHE is set. Do not also wrap CXX: that
+# would produce a recursive "ccache ccache g++" invocation.
 if [[ "${CCACHE}" == "1" ]]; then
-    # Install ccache if not present (needs >= 4.1 for Redis remote storage support)
-    if ! command -v ccache &> /dev/null; then
-        apt-get update && apt-get install -y --no-install-recommends ccache
+    NVTE_CCACHE_BIN="${NVTE_CCACHE_BIN:-ccache}"
+    CACHE_PROGRAM="$(basename -- "${NVTE_CCACHE_BIN}")"
+
+    if ! command -v "${NVTE_CCACHE_BIN}" &> /dev/null; then
+        # Keep direct callers backwards compatible while the Docker build
+        # installs its selected cache explicitly.
+        CACHE_INSTALLER="${SCRIPT_DIR}/install-compiler-cache.sh"
+        if [[ ! -x "${CACHE_INSTALLER}" ]]; then
+            CACHE_INSTALLER="$(command -v install-compiler-cache.sh || true)"
+        fi
+        if [[ -z "${CACHE_INSTALLER}" || ! -x "${CACHE_INSTALLER}" ]]; then
+            echo "${NVTE_CCACHE_BIN} is not installed and install-compiler-cache.sh is unavailable"
+            exit 1
+        fi
+        "${CACHE_INSTALLER}" "${NVTE_CCACHE_BIN}"
+    fi
+    if ! command -v "${NVTE_CCACHE_BIN}" &> /dev/null; then
+        echo "Compiler cache installation did not provide ${NVTE_CCACHE_BIN}"
+        exit 1
     fi
-    export CXX="ccache g++"
+
     export NVTE_USE_CCACHE=1
-    ccache --zero-stats
+    export NVTE_CCACHE_BIN
+    case "${CACHE_PROGRAM}" in
+        sccache)
+            # Keep the daemon alive for long CUDA builds and retain an error log
+            # so backend failures are visible alongside the cache statistics.
+            export SCCACHE_IDLE_TIMEOUT="${SCCACHE_IDLE_TIMEOUT:-0}"
+            export SCCACHE_ERROR_LOG="${SCCACHE_ERROR_LOG:-/tmp/sccache-server.log}"
+            rm -f -- "${SCCACHE_ERROR_LOG}"
+            "${NVTE_CCACHE_BIN}" --start-server
+            "${NVTE_CCACHE_BIN}" --zero-stats
+            ;;
+        ccache)
+            export CCACHE_DIR="${CCACHE_DIR:-/root/.cache/ccache}"
+            "${NVTE_CCACHE_BIN}" --zero-stats
+            ;;
+        *)
+            echo "Unsupported compiler cache: ${NVTE_CCACHE_BIN}"
+            exit 1
+            ;;
+    esac
+    echo "Transformer Engine compiler cache: ${CACHE_PROGRAM}"
 fi
 
 # The wheel filename includes the TE commit; if this has changed since the last
 # incremental build then we would end up with multiple wheels.
 rm -fv dist/*.whl
 python setup.py bdist_wheel
-ls dist/
 popd
 
+CACHE_STATUS=0
 if [[ "${CCACHE}" == "1" ]]; then
-    ccache --show-stats --verbose
+    case "${CACHE_PROGRAM}" in
+        sccache)
+            "${NVTE_CCACHE_BIN}" --show-stats || CACHE_STATUS=$?
+            if [[ -s "${SCCACHE_ERROR_LOG}" ]]; then
+                echo "WARNING: sccache reported backend errors:"
+                tail -n 200 "${SCCACHE_ERROR_LOG}"
+            fi
+            # Stopping the daemon drains any in-flight cache uploads.
+            "${NVTE_CCACHE_BIN}" --stop-server || true
+            rm -f -- "${SCCACHE_ERROR_LOG}"
+            ;;
+        ccache)
+            "${NVTE_CCACHE_BIN}" --show-stats --verbose || CACHE_STATUS=$?
+            ;;
+    esac
 fi
 
 ## Install the built packages