From 64979643427db171ea79822fbdeb076321be6ef9 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 18 Dec 2024 15:18:27 +0100 Subject: [PATCH] misc(backend): lets try... --- .github/workflows/build_trtllm.yaml | 9 ++-- Dockerfile_trtllm | 28 ++++++------ backends/trtllm/scripts/setup_sccache.py | 57 ++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 18 deletions(-) create mode 100644 backends/trtllm/scripts/setup_sccache.py diff --git a/.github/workflows/build_trtllm.yaml b/.github/workflows/build_trtllm.yaml index fc132404..88bd2834 100644 --- a/.github/workflows/build_trtllm.yaml +++ b/.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 diff --git a/Dockerfile_trtllm b/Dockerfile_trtllm index 6c1445d2..d7049cb7 100644 --- a/Dockerfile_trtllm +++ b/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 && \ diff --git a/backends/trtllm/scripts/setup_sccache.py b/backends/trtllm/scripts/setup_sccache.py new file mode 100644 index 00000000..6e3c2ab2 --- /dev/null +++ b/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()