Add ccache support to TE and use it in GitHub actions (#2444)

* Add ccache support to TE and use it in GitHub actions
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Move to allowed action with sccache
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Properly handle sccache
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Fix typo
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Removing ccache from the custom docker workflows where we can't run the
action in the container
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* JAX already uses same cmake options to build the extension so there is
no need to set CXX too
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

* Removed the unnecessary env variables
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

.github/workflows/build.yml

@@ -24,11 +24,14 @@ jobs:
         uses: actions/checkout@v3
         with:
           submodules: recursive
+      - name: ccache
+        uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad
       - name: 'Build'
-        run: pip install --no-build-isolation . -v
+        run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v
         env:
           NVTE_FRAMEWORK: none
           MAX_JOBS: 1
+          SCCACHE_GHA_ENABLED: "true"
       - name: 'Sanity check'
         run: python3 -c "import transformer_engine"
         working-directory: /
@@ -94,11 +97,15 @@ jobs:
         uses: actions/checkout@v3
         with:
           submodules: recursive
+      - name: ccache
+        uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad
       - name: 'Build'
-        run: pip install --no-build-isolation . -v
+        run: |
+          NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v
         env:
           NVTE_FRAMEWORK: jax
           MAX_JOBS: 1
+          SCCACHE_GHA_ENABLED: "true"
       - name: 'Sanity check'
         run: python3 tests/jax/test_sanity_import.py
   all:
@@ -140,7 +147,6 @@ jobs:
             pip install pybind11[global] einops onnxscript && \
             pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130
           '
-
       - name: 'Build'
         run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps'
         env:

build_tools/build_ext.py

@@ -61,6 +61,12 @@ class CMakeExtension(setuptools.Extension):
             f"-DCMAKE_BUILD_TYPE={build_type}",
             f"-DCMAKE_INSTALL_PREFIX={install_dir}",
         ]
+        if bool(int(os.getenv("NVTE_USE_CCACHE", "0"))):
+            ccache_bin = os.getenv("NVTE_CCACHE_BIN", "ccache")
+            configure_command += [
+                f"-DCMAKE_CXX_COMPILER_LAUNCHER={ccache_bin}",
+                f"-DCMAKE_CUDA_COMPILER_LAUNCHER={ccache_bin}",
+            ]
         configure_command += self.cmake_flags
 
         import pybind11