From 742d00db94b2b3f8c3abe6da4e90df9b37d6bf64 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Fri, 20 Dec 2024 14:41:59 +0100
Subject: [PATCH] misc(backend): again

---
 Dockerfile_trtllm                        |  4 ++--
 backends/trtllm/scripts/setup_sccache.py | 24 ++++++++----------------
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/Dockerfile_trtllm b/Dockerfile_trtllm
index 88bd990a527..b940e808cd9 100644
--- a/Dockerfile_trtllm
+++ b/Dockerfile_trtllm
@@ -79,7 +79,7 @@ ARG sccache_s3_key_prefix
 ARG sscache_region
 
 ENV AWS_ACCESS_KEY_ID=$aws_access_key_id
-ENV AWS_SECRET_KEY_ID=$aws_secret_key_id
+ENV AWS_SECRET_ACCESS_KEY=$aws_secret_access_key
 ENV AWS_SESSION_TOKEN=$aws_session_token
 ENV SCCACHE_BUCKET=$sccache_bucket
 ENV SCCACHE_S3_KEY_PREFIX=$SCCACHE_S3_KEY_PREFIX
@@ -96,7 +96,7 @@ 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" && \
-    python3 backends/trtllm/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 && \
+    python3 backends/trtllm/scripts/setup_sccache.py --is-gha-build ${IS_GHA_BUILD} && \
     RUSTC_WRAPPER=sccache CMAKE_INSTALL_PREFIX=$TGI_INSTALL_PREFIX cargo build --profile ${build_type} --package text-generation-backends-trtllm --bin text-generation-backends-trtllm && \
     sccache --show-stats
 
diff --git a/backends/trtllm/scripts/setup_sccache.py b/backends/trtllm/scripts/setup_sccache.py
index 6e3c2ab2b88..982f8c77d23 100644
--- a/backends/trtllm/scripts/setup_sccache.py
+++ b/backends/trtllm/scripts/setup_sccache.py
@@ -24,12 +24,13 @@ def setup_sccache_locally():
                 del environ[envvar]
 
 
-def setup_sccache_for_s3(s3_args):
+def setup_sccache_for_s3():
     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)
+    for envvar in AWS_S3_CACHING_VARIABLES.keys():
+        if not envvar in environ or not environ[envvar] or len(environ[envvar]) == 0:
+            print(f"Missing definition for environment variable {envvar}")
 
 
 if __name__ == "__main__":
@@ -37,21 +38,12 @@ def setup_sccache_for_s3(s3_args):
 
     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()
+    if args.is_gha_build:
+        setup_sccache_for_s3()
+    else:
+        setup_sccache_locally()