Failing deployment on AWS Sagemaker endpoints

[CoolFish88]

System Info

Hello,

When attempting to deploy TEI 1.6.1 images on AWS Sagemaker GPU endpoints (e.g ml.g5.2xlarge), various errors led to a failed deployment , as summarized by the following CloudWatch logs:

• ghcr.io/huggingface/text-embeddings-inference:cuda-1.6.1 and ghcr.io/huggingface/text-embeddings-inference:cuda-sha-7d4d9ec

./entrypoint.sh: line 10: [: -eq: unary operator expected
./entrypoint.sh: line 13: [: too many arguments
./entrypoint.sh: line 16: [: -eq: unary operator expected
cuda compute cap is not supported

• ghcr.io/huggingface/text-embeddings-inference:86-1.6.1

error:unexpected argument 'serve' found
Usage: text-embeddings-router [OPTIONS]
For more information, try '--help'.

Each deployment referenced model artifacts from jinaai/jina-embeddings-v2-small-en supplied as an S3 archive to different deployment strategies leveraging: (1) HuggingFaceModeland (2) Sagemaker Model and associated Endpoint Config created with a boto3 Sagemaker client. Similarly, the deployment failed when instructing the endpoint to fetch model artifacts from the hub (approach using HuggingFaceModel)

Remark: When deploying TEI 1.4.0 using the HuggingFaceModel approach, asa retrieved by the following code:

from sagemaker.huggingface import get_huggingface_llm_image_uri
tei_image_uri = get_huggingface_llm_image_uri("huggingface-tei", version="1.4.0")

the process completes without errors as long as model artifacts are fetched from the hub. When supplying model artifacts in an S3 archive, deployment fails due to the incorrect backend being initialized, as discussed here and addressed in #559

Information

• Docker
• The CLI directly

Tasks

• An officially supported command
• My own modifications

Reproduction

Deployment instructions using HuggingFaceModel:

• With model artifacts fetched from the hub

tei_image_uri = <image_uri>
emb_model = HuggingFaceModel(
    name="my-tei-model",
    role=role,
    #model_data=<s3_path_to_optional_model_artifacts),
    sagemaker_session= <sm_session>,
    image_uri=tei_image_uri,
    env={"HF_TASK": "feature-extraction",
         "HF_MODEL_ID": "jinaai/jina-embeddings-v2-small-en",
    },
)

emb_predictor = emb_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge",
    endpoint_name="jina-embeddings-tei"
)

• With model artifacts stored in S3:

Modify the above code such that:
- model_data points to an S3 tar.gz. archive storing model artifacts
- HF_MODEL_ID points to /opt/ml/model

Expected behavior

Deployment completes without errors

[CoolFish88]

Rebuilt the container locally following the solution discussed here: #566 (comment) using: docker build . -f Dockerfile-cuda-all --build-arg CUDA_COMPUTE_CAP=86.

Deployed the container to a A10G GPU Sagemaker instance and the process fails with the above errors.

[CoolFish88]

@alvarobartt Failed deployments mentioning cuda compute cap is not supported may be caused by the mismatch between the CUDA version in the container and Sagemaker instance.

Experimented with the production variants listed here: https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_ProductionVariant.html, ensuring newer diver versions are mounted on the instance.

This approach addressed the CUDA error, but deployments failed with error:

error:unexpected argument 'serve' found
Usage: text-embeddings-router [OPTIONS]
For more information, try '--help'.

[CoolFish88]

Built the image locally for the target GPU using docker build . -f Dockerfile-cuda-all --build-arg CUDA_COMPUTE_CAP=86.
Deployed on a Sagemaker instance with A10G processor. Deployment fails with:

error: unexpected argument 'serve' found
Usage: text-embeddings-router-80 [OPTIONS]
For more information, try '--help'.

[CoolFish88]

The error suggests that an argument 'serve' is being passed to the text-embeddings-router command, but this argument isn't expected by the router binary.

In the SageMaker context, this likely occurs because SageMaker automatically appends 'serve' as an argument when starting containers on endpoints. This is part of SageMaker's standard container interface.

[CoolFish88]

Managed to bypass the "serve" error by setting --target sagemaker when building the image. This tells Docker to build the image up to the stage in Dockerfile-cuda-all named sagemaker

This is important because SageMaker has specific container interface requirements (i.e. automatically adds 'serve' as an argument when starting containers). Without --target sagemaker, Docker would build the default image, which will not work correctly in the SageMaker environment.

[alvarobartt]

Yes that's true when building the SageMaker from the repository, but shouldn't be required to specify when running the image on AWS SageMaker with the official Hugging Face image, unless it's the case and something's misconfigured.

Just to understand it a bit better, the issue is that the latest image released on AWS SageMaker is outdated and you need a newer one, right? In any case we're already working on shipping the recently released TEI 1.7.0 container in SageMaker 🤗

[CoolFish88]

That's right. The latest version I could fetch that is compatible with AWS Sagemaker was 1.4.0, using:

from sagemaker.huggingface import get_huggingface_llm_image_uri
tei_image_uri = get_huggingface_llm_image_uri("huggingface-tei",version="1.4.0")

Unfortunately, this version doesn't incorporate the features and bug fixes added in newer versions.
This is the reason I had to build a Sagemaker compatible version from the repository.

A few questions regarding the new container you're making available for Sagemaker:

• Can it be retrieved using get_huggingface_llm_image_uri?
• Will it have support for jina rerankers (e.g. https://huggingface.co/jinaai/jina-reranker-v1-turbo-en) ?

Thank you!

[alvarobartt]

Hey @CoolFish88 to answer your questions:

• Yes, once released I'll share with you all the information and details on how to deploy it on AWS SageMaker!
• Yes, but it's an issue with the config.json rather than TEI, as it's indeed compatible, feel free to use the following revision in the meantime https://huggingface.co/jinaai/jina-reranker-v1-turbo-en/discussions/13 🤗

Thanks!

[CoolFish88]

Thank you once again @alvarobartt for the continued support on these issues

[CoolFish88]

@alvarobartt
I tried deploying the jina-reranker model but it fails with the error reported here: #556. Does this fix #559 need to be extended for this model as well in order to work?