feat: Amazon SageMaker compatible images

[JGalego]

What does this PR do?

Adds support for Amazon SageMaker compatible images. Similar to huggingface/text-generation-inference#147, only for TEI.
Mostly CI stuff and some hacks, since the Amazon SageMaker routes are already implemented.

Implementation:

• Added sagemaker target to Dockerfile-cuda and custom entrypoint
• Added build-and-push-sagemaker-image step to build_* workflows

Who can review?

@OlivierDehaene

[austinmw]

Hi, happen to know if this supports BAAI/bge-reranker-base?

[JGalego]

Hi, happen to know if this supports BAAI/bge-reranker-base?

It does (refs/pr/5 as described in https://github.com/huggingface/text-embeddings-inference#sequence-classification-and-re-ranking).

Running the example request in https://github.com/huggingface/text-embeddings-inference#using-re-rankers-models

[superlaut]

Hi team, is there an expected date for merging this PR? I would like to use these images :)

[andrewrreed]

FYI - here's a working notebook to deploy TEI on SageMaker in the meantime: https://github.com/andrewrreed/hf-notebooks/blob/main/deploy-tei-sagemaker/tei-sagemaker-inference.ipynb

[philschmid]

Hey @JGalego,

Thank you for starting this work. We would need to make some changes to match the environment variables, similar to TGI. https://github.com/huggingface/text-generation-inference/blob/main/sagemaker-entrypoint.sh
Do you want to make those changes? If you are busy we can also contribute directly into your branch.

Also due to complexity of SageMaker, we will build 1 GPU container containing all binaries for the target (https://github.com/huggingface/text-embeddings-inference/blob/main/Dockerfile-cuda-all).
Could you make the changes to the dockerfile-cuda-all?

[JGalego]

Hey @JGalego,

Thank you for starting this work. We would need to make some changes to match the environment variables, similar to TGI. https://github.com/huggingface/text-generation-inference/blob/main/sagemaker-entrypoint.sh Do you want to make those changes? If you are busy we can also contribute directly into your branch.

Also due to complexity of SageMaker, we will build 1 GPU container containing all binaries for the target (https://github.com/huggingface/text-embeddings-inference/blob/main/Dockerfile-cuda-all). Could you make the changes to the dockerfile-cuda-all?

Hey @philschmid,

I can add the env var mappings for HF_MODEL_ID --> MODEL_ID and HF_MODEL_REVISION --> REVISION just as in TGI - anything else?

Regarding the all_image (#186), should I had 1/ a separate target for SM with a custom sagemaker-cuda-all-entrypoint.sh (basically, cuda-all-entrypoint.sh + env var mappings) , 2/ change the existing entrypoint or 3/ move this discussion to a different PR?

[philschmid]

Lets start only with the "all" versions.

[philschmid]

Thank you for working on this! As soon as this is merged we ll work with the SageMaker to make it officially available in side the sagemaker-sdk.

[OlivierDehaene]

Thank you for this and sorry for the delay!

[austinmw]

@JGalego Hi, I have a question. Does this work when using an S3 path for model artifacts instead of HF Hub, as in HuggingFaceModel(model_data="<s3_path>", ...)?

I get an endpoint error if I don't specify HF_MODEL_ID as an env config, however I thought that model_data will be ignored if I specify this?

Appreciate any advice, thanks!

[JGalego]

@JGalego Hi, I have a question. Does this work when using an S3 path for model artifacts instead of HF Hub, as in HuggingFaceModel(model_data="<s3_path>", ...)?

I get an endpoint error if I don't specify HF_MODEL_ID as an env config, however I thought that model_data will be ignored if I specify this?

Appreciate any advice, thanks!

For the SM images, HF_MODEL_ID is mandatory (https://github.com/huggingface/text-embeddings-inference/blob/main/sagemaker-entrypoint-cuda-all.sh#L25). Checking the code (https://github.com/huggingface/text-embeddings-inference/blob/main/router/src/lib.rs#L95) and the docs (https://github.com/huggingface/text-embeddings-inference?tab=readme-ov-file#docker), it only supports model artifacts from the hub.

Options:
      --model-id <MODEL_ID>
          The name of the model to load. Can be a MODEL_ID as listed on <https://hf.co/models> like `thenlper/gte-base`.
          Or it can be a local directory containing the necessary files as saved by `save_pretrained(...)` methods of
          transformers

          [env: MODEL_ID=]
          [default: thenlper/gte-base]

[austinmw]

@JGalego Would it be possible to update this so AWS customers can alternatively use S3 for model storage? I can submit a feature request if needed.

[philschmid]

Hello @austinmw @JGalego,

This works already. Just set the HF_MODEL_ID to /opt/ml/model. SageMaker loads your model from s3 to /opt/ml/model and TEI starts correctly.

[austinmw]

Hi @philschmid,

I set HF_MODEL_ID="/opt/ml/model", but got this deployment error:

Here's my code:

[philschmid]

@austinmw This looks more that your artifact is not correctly created. Can you check https://huggingface.co/docs/sagemaker/inference#create-a-model-artifact-for-deployment

[austinmw]

@philschmid Just tried with new artifacts using the provided steps and got the same exact error:

# delete my fine-tuned model from s3
aws s3 rm s3://sagemaker-us-east-1-197237853439/models/bge-base-en-v1.5/model.tar.gz

# upload fresh artifacts according to steps
git lfs install
git clone https://huggingface.co/BAAI/bge-base-en-v1.5
cd bge-base-en-v1.5/
tar zcvf model.tar.gz *
aws s3 cp model.tar.gz s3://sagemaker-us-east-1-197237853439/models/bge-base-en-v1.5/model.tar.gz

Is it possible this feature is broken for TEI?

[philschmid]

Nope it should work. I am working on a SageMaker end-to-end example with training and deployment and it works. Can you try uploading your model unzipped and then deploy it using the "uncompressed" setting?

s3_path="s3://path/to/ymodel"

model = HuggingFaceModel(
  role=role,
  # path to s3 bucket with model, we are not using a compressed model
  model_data={'S3DataSource':{'S3Uri': s3_path + "/",'S3DataType': 'S3Prefix','CompressionType': 'None'}},
  image_uri=image,
  env=config
)```

[austinmw]

That worked, thanks a ton for your help!