huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

misc(backend): again

Browse Source
This commit is contained in:
Morgan Funtowicz 2024-12-20 14:41:59 +01:00
parent 01289399d6
commit 43ab954b4d
2 changed files with 10 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
Dockerfile_trtllm
View File

@ -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

22
backends/trtllm/scripts/setup_sccache.py
View File

@ -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 @@ if __name__ == "__main__":
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 _:
if args.is_gha_build:
setup_sccache_for_s3()
else:
setup_sccache_locally()