From 5910dabb4e6787457965ee6330d5aac6062e821b Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Thu, 12 Dec 2024 12:45:14 +0100
Subject: [PATCH] misc(ci): provide mecanism to cache inside container

---
 .github/workflows/build_trtllm.yaml |  9 ++++++---
 Dockerfile_trtllm                   | 16 +++++++++++++++-
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/build_trtllm.yaml b/.github/workflows/build_trtllm.yaml
index a350cd1e..4e441b34 100644
--- a/.github/workflows/build_trtllm.yaml
+++ b/.github/workflows/build_trtllm.yaml
@@ -38,6 +38,7 @@ jobs:
           override: true
 
       - name: "Configure AWS Credentials"
+        id: aws-creds
         uses: aws-actions/configure-aws-credentials@v4
         with:
           aws-region: us-east-1
@@ -59,12 +60,14 @@ jobs:
           push: false
           platforms: 'linux/amd64'
           build-args: |
+            AWS_ACCESS_KEY_ID: ${{ steps.aws-creds.outputs.aws-access-key-id }}
+            AWS_SECRET_ACCESS_KEY: ${{ steps.aws-creds.outputs.aws-secret-access-key }}
             BUILD_TYPE=debug
             SCCACHE_BUCKET: ${{ secrets.AWS_S3_BUCKET_GITHUB_TGI_TEST }}
-            SCCACHE_REGION: "us-east-1"
+            SCCACHE_REGION: us-east-1
             SCCACHE_S3_USE_SSL: false
-            SCCACHE_S3_KEY_PREFIX: "tgi+trtllm"
-            RUSTC_WRAPPER: sccache
+            SCCACHE_S3_KEY_PREFIX: trtllm
+
 
 
 
diff --git a/Dockerfile_trtllm b/Dockerfile_trtllm
index 769aa927..4f315b59 100644
--- a/Dockerfile_trtllm
+++ b/Dockerfile_trtllm
@@ -2,6 +2,13 @@ ARG CUDA_ARCH_LIST="75-real;80-real;86-real;89-real;90-real"
 ARG OMPI_VERSION="4.1.7rc1"
 ARG BUILD_TYPE=release
 
+# Additional arguments if we enable sccache in the CI/CD
+ARG AWS_ACCESS_KEY_ID
+ARG AWS_SECRET_ACCESS_KEY
+ARG SCCACHE_BUCKET=""
+ARG SCCACHE_REGION=""
+ARG SCCACHE_S3_KEY_PREFIX=""
+
 # CUDA dependent dependencies resolver stage
 FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu24.04 AS cuda-builder
 
@@ -68,11 +75,18 @@ ENV CMAKE_PREFIX_PATH="/usr/local/mpi:/usr/local/tensorrt:$CMAKE_PREFIX_PATH"
 ENV LD_LIBRARY_PATH="/usr/local/mpi/lib:$LD_LIBRARY_PATH"
 ENV PKG_CONFIG_PATH="/usr/local/mpi/lib/pkgconfig:$PKG_CONFIG_PATH"
 
+# Retrieve potential build arg for S3 caching
+ENV AWS_ACCESS_KEY_ID ${AWS_ACCESS_KEY_ID}
+ENV AWS_SECRET_ACCESS_KEY ${AWS_SECRET_ACCESS_KEY}
+ENV SCCACHE_BUCKET ${SCCACHE_BUCKET}
+ENV SCCACHE_REGION ${SCCACHE_REGION}
+ENV SCCACHE_S3_KEY_PREFIX ${SCCACHE_S3_KEY_PREFIX}
+
 COPY . .
 COPY --from=trt-builder /usr/local/tensorrt /usr/local/tensorrt
 COPY --from=mpi-builder /usr/local/mpi /usr/local/mpi
 RUN mkdir $TGI_INSTALL_PREFIX && mkdir "$TGI_INSTALL_PREFIX/include" && mkdir "$TGI_INSTALL_PREFIX/lib" && \
-    CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX sccache build --${BUILD_TYPE} --package text-generation-backends-trtllm --bin text-generation-backends-trtllm
+    CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX cargo build --${BUILD_TYPE} --package text-generation-backends-trtllm --bin text-generation-backends-trtllm
 
 FROM nvidia/cuda:12.6.3-cudnn-runtime-ubuntu24.04 AS runtime
 RUN apt update && apt install -y libucx0 pipx python3-minimal python3-dev python3-pip python3-venv && \