misc(backend): lets try...

.github/workflows/build_trtllm.yaml

@@ -63,15 +63,14 @@ jobs:
           load: true
           platforms: 'linux/amd64'
           build-args: |
-            BUILD_TYPE=dev
-            IS_GHA_BUILD=true
-          secrets: |
+            build_type=dev
+            is_gha_build=true
             aws_access_key_id=${{ steps.aws-creds.outputs.aws-access-key-id }}
             aws_secret_access_key=${{ steps.aws-creds.outputs.aws-secret-access-key }}
             aws_session_token=${{ steps.aws-creds.outputs.aws-session-token }}
-            sccache_s3_bucket=${{ secrets.AWS_S3_BUCKET_GITHUB_TGI_TEST }}
+            sccache_bucket=${{ secrets.AWS_S3_BUCKET_GITHUB_TGI_TEST }}
             sccache_s3_key_prefix=trtllm-${{ env.TENSORRT_LLM_VERSION }}
-            sccache_s3_region=us-east-1
+            sccache_region=us-east-1
           cache-from: type=s3,region=us-east-1,bucket=ci-docker-buildx-cache,name=text-generation-inference-cache-${{ env.LABEL }}-trtllm,mode=min,access_key_id=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_ACCESS_KEY_ID }},secret_access_key=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_SECRET_ACCESS_KEY }},mode=min
           cache-to: type=s3,region=us-east-1,bucket=ci-docker-buildx-cache,name=text-generation-inference-cache-${{ env.LABEL }}-trtllm,mode=min,access_key_id=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_ACCESS_KEY_ID }},secret_access_key=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_SECRET_ACCESS_KEY }},mode=min
 

Dockerfile_trtllm

@@ -63,27 +63,29 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | bash -s -- -y &&
     cargo install sccache --locked
 
 # Build actual TGI
-ARG BUILD_TYPE
-ARG CUDA_ARCH_LIST
-ARG IS_GHA_BUILD
-ARG SCCACHE_S3_KEY_PREFIX
+ARG build_type
+ARG cuda_arch_list
+ARG is_gha_build
+
+# SCCACHE Specifics args - before finding a better, more generic, way...
+ARG aws_access_key_id
+ARG aws_secret_key_id
+ARG aws_session_token
+ARG sccache_bucket
+ARG sccache_s3_key_prefix
+ARG sscache_region
 
-ENV IS_GHA_BUILD=${IS_GHA_BUILD}
 ENV CMAKE_PREFIX_PATH="/usr/local/mpi:/usr/local/tensorrt:$CMAKE_PREFIX_PATH"
+ENV CUDA_ARCH_LIST=${cuda_arch_list}
 ENV LD_LIBRARY_PATH="/usr/local/mpi/lib:$LD_LIBRARY_PATH"
 ENV PKG_CONFIG_PATH="/usr/local/mpi/lib/pkgconfig:$PKG_CONFIG_PATH"
 
 COPY . .
 COPY --from=trt-builder /usr/local/tensorrt /usr/local/tensorrt
 COPY --from=mpi-builder /usr/local/mpi /usr/local/mpi
-RUN --mount=type=secret,id=aws_access_key_id,env=AWS_ACCESS_KEY_ID  \
-    --mount=type=secret,id=aws_secret_access_key,env=AWS_SECRET_ACCESS_KEY \
-    --mount=type=secret,id=aws_session_token,env=AWS_SESSION_TOKEN \
-    --mount=type=secret,id=sccache_s3_bucket,env=SCCACHE_BUCKET \
-    --mount=type=secret,id=sccache_s3_key_prefix,env=SCCACHE_S3_KEY_PREFIX \
-    --mount=type=secret,id=sccache_s3_region,env=SCCACHE_REGION \
-    mkdir $TGI_INSTALL_PREFIX && mkdir "$TGI_INSTALL_PREFIX/include" && mkdir "$TGI_INSTALL_PREFIX/lib" && \
-    RUSTC_WRAPPER=sccache CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX cargo build --profile ${BUILD_TYPE} --package text-generation-backends-trtllm --bin text-generation-backends-trtllm
+RUN mkdir $TGI_INSTALL_PREFIX && mkdir "$TGI_INSTALL_PREFIX/include" && mkdir "$TGI_INSTALL_PREFIX/lib" && \
+    python3 scripts/setup_sccache.py --is-gha-build ${is_gha_build} -k ${aws_access_key_id} -s ${aws_secret_key_id} -t ${aws_session_token} -b ${sccache_bucket} -r ${sscache_region} -p ${sccache_s3_key_prefix }&& \
+    RUSTC_WRAPPER=sccache CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX cargo build --profile ${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 && \

backends/trtllm/scripts/setup_sccache.py

@@ -0,0 +1,57 @@
+from argparse import ArgumentParser
+
+AWS_S3_CACHING_VARIABLES = {
+    "AWS_ACCESS_KEY_ID": "aws_access_key_id",
+    "AWS_SECRET_ACCESS_KEY": "aws_secret_access_key",
+    "AWS_SESSION_TOKEN": "aws_session_token",
+    "SCCACHE_REGION": "s3_region",
+    "SCCACHE_BUCKET": "s3_bucket_name",
+}
+
+ALL_CACHING_STORAGE_VARIABLES = {
+    "AWS_S3_CACHING_VARIABLES"
+}
+
+
+def setup_sccache_locally():
+    from os import environ
+
+    print("Setting up Local Caching Layer")
+    for target in ALL_CACHING_STORAGE_VARIABLES:
+        for envvar in globals()[target].keys():
+            if envvar in environ:
+                print(f"Deleted {envvar} from environment variables")
+                del environ[envvar]
+
+
+def setup_sccache_for_s3(s3_args):
+    from os import environ
+
+    print("Setting up AWS S3 Caching Layer")
+    for envvar, field in AWS_S3_CACHING_VARIABLES.items():
+        environ[envvar] = getattr(s3_args, field)
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser("TensorRT-LLM Build Caching Setup")
+
+    parser.add_argument("--is-gha-build", type=str, default="FALSE",
+                        help="Indicate if the build is from Github Actions")
+    parser.add_argument("--aws-access-key-id", "-k", type=str, required=True, help="AWS Access Key ID to use")
+    parser.add_argument("--aws-secret-access-key", "-s", type=str, required=True,
+                        help="AWS Secret Access Key to use")
+    parser.add_argument("--aws-session-token", "-t", type=str, required=True, help="AWS Session Token to use")
+    parser.add_argument("--s3-bucket-name", "-b", type=str, required=True, help="AWS target S3 Bucket")
+    parser.add_argument("--s3-bucket-prefix", "-p", type=str, required=True, help="AWS target S3 Bucket folder prefix")
+    parser.add_argument("--s3-region", "-r", type=str, required=True, help="AWS target S3 region")
+
+    # Parse args
+    args = parser.parse_args()
+    args.is_gha_build = args.is_gha_build.lower() in {"on", "true", "1"}
+    print(args)
+
+    match args.store:
+        case "s3":
+            setup_sccache_for_s3(args)
+        case _:
+            setup_sccache_locally()