Skip to content

Commit

Permalink
misc(backend): again
Browse files Browse the repository at this point in the history
  • Loading branch information
mfuntowicz committed Dec 20, 2024
1 parent 8f52d72 commit 742d00d
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 18 deletions.
4 changes: 2 additions & 2 deletions Dockerfile_trtllm
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down
24 changes: 8 additions & 16 deletions backends/trtllm/scripts/setup_sccache.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,34 +24,26 @@ 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__":
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()
if args.is_gha_build:
setup_sccache_for_s3()
else:
setup_sccache_locally()

0 comments on commit 742d00d

Please sign in to comment.