在AWS Lambda部署SentenceTransformer遇CUFFT库映射失败问题求助

问题原因 #

AWS Lambda运行环境无GPU支持，但你当前安装的依赖（如PyTorch）默认包含CUDA相关库，Lambda加载这些需要GPU的库时，就会抛出/var/lang/lib/python3.11/site-packages/nvidia/cufft/lib/libcufft.so.10: failed to map segment from shared object错误。

解决步骤 #

1. 修正requirements.txt，安装CPU版本依赖 #

替换原依赖为CPU专属版本，避免引入GPU相关库：

sentence-transformers==2.2.2
torch==2.0.1+cpu
torchvision==0.15.2+cpu
transformers==4.30.2

注：安装时需确保拉取CPU版本包，可通过指定PyTorch CPU源实现。

2. 调整Dockerfile，适配Lambda环境 #

基于Lambda官方Python镜像构建，确保依赖安装过程正确获取CPU版本：

FROM public.ecr.aws/lambda/python:3.11

COPY requirements.txt .
RUN pip install -r requirements.txt --index-url https://download.pytorch.org/whl/cpu

COPY main.py .

CMD ["main.handler"]

3. 修改handler代码，强制使用CPU加载模型 #

初始化SentenceTransformer时明确指定device为CPU，避免自动检测GPU：

from sentence_transformers import SentenceTransformer

# 强制使用CPU加载模型
model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')

def handler(event, context):
    texts = event.get('texts', [])
    embeddings = model.encode(texts)
    return {
        'embeddings': embeddings.tolist()
    }

额外优化建议 #

• 提前下载模型文件打包到镜像中，避免Lambda冷启动时重复下载模型，降低延迟。
• 清理镜像冗余文件（如pip缓存、无用依赖），缩小镜像体积，加快冷启动速度。

内容的提问来源于stack exchange，提问作者Pradeep Sanjaya