在SageMaker上用LLamaindex部署HuggingFace Embedding报500错误求助

问题背景 #

在Amazon SageMaker平台上使用LLaMA Index实现HuggingFace Embedding功能时，SageMaker端点返回Worker Died（500错误），寻求同类问题经验或解决思路。

当前实现细节 #

• 使用SimpleDirectoryReader读取TXT文档
• 向量存储：QdrantVectorStore
• 文档存储：MongoDocumentStore
• 索引存储：MongoIndexStore

相关代码 #

模型部署代码 #

# Hub Model configuration.
hub = {
    'HF_MODEL_ID':'BAAI/bge-base-en-v1.5', # model_id from hf.co/models
    'HF_TASK':'feature-extraction', # NLP task you want to use for predictions
    'MAX_INPUT_LENGTH': dumps(1024),  # Max length of input text
    'MAX_TOTAL_TOKENS': dumps(2048),  # Max length of the generation (including input text)
    # 'SM_NUM_GPUS': '1',
}

# Configure the HF container
huggingface_model = HuggingFaceModel(
    env=hub, # configuration for loading model from Hub
    role=role, # iam role with permissions to create an Endpoint
    py_version='py310',
    transformers_version="4.37.0", # transformers version used
    pytorch_version="2.1.0", # pytorch version used
)

# Deploy HF embedding
embedding_predictor = huggingface_model.deploy(
    endpoint_name=name_from_base("bge-base-en-v15"),
    initial_instance_count=1,
    instance_type="ml.c6i.xlarge"
)

# Declare embedding_model
embedding_model = SageMakerEmbedding(
    endpoint_name=embedding_predictor.endpoint_name,
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
    region_name=AWS_REGION
)

数据Ingestion Pipeline代码 #

# Ingestion Pipeline Code
pipeline = IngestionPipeline(
        transformations=[
            SentenceSplitter(chunk_size=250, chunk_overlap=50),
            Settings.embed_model
        ],
        vector_store=vector_store,
        docstore=doc_store,
    )
nodes = pipeline.run(documents=documents)

报错信息（翻译后） #

WorkerThread - 9000-52f284e8 Worker disconnected. WORKER_MODEL_LOADED                                                                           
WorkerLifeCycle - Frontend disconnected.                                                                                              
WorkerLifeCycle - Backend worker process died.                                                                                        
WorkerLifeCycle - 回溯（最近的调用最后）：                                                                                  
WorkerLifeCycle -   文件 "/opt/conda/lib/python3.10/site-packages/mms/model_service_worker.py", 第175行, 在 start_worker 中             
WorkerLifeCycle -     self.handle_connection(cl_socket)                                                                               
WorkerLifeCycle -   文件 "/opt/conda/lib/python3.10/site-packages/mms/model_service_worker.py", 第139行, 在 handle_connection 中       
WorkerLifeCycle -     cmd, msg = retrieve_msg(cl_socket)                                                                              
WorkerLifeCycle -   文件 "/opt/conda/lib/python3.10/site-packages/mms/protocol/otf_message_handler.py", 第36行, 在 retrieve_ms 中
WorkerLifeCycle -     cmd = _retrieve_buffer(conn, 1)
WorkerThread - 未知异常                                                                                                                
io.netty.handler.codec.CorruptedFrameException: 消息大小超出限制: 35054253     

解决思路 #

• 调小Embedding请求的批量大小：LLaMA Index调用SageMaker端点时可能一次性发送了过大的文本批量，导致消息超出MMS（模型管理服务）的默认限制。可以在初始化SageMakerEmbedding时指定batch_size参数（比如设置为16或32），减少单请求的数据量。
• 验证文本分块效果：虽然配置了SentenceSplitter(chunk_size=250)，但需确认实际生成的文本块是否存在异常超长的情况（比如文档格式问题导致分块逻辑失效）。可以在分块步骤后打印部分文本块的长度，验证分块是否正常生效。
• 调整MMS的消息大小限制：通过环境变量MMS_MAX_REQUEST_SIZE修改SageMaker端点的MMS配置，在部署HuggingFaceModel时将该变量加入env字典（例如设置为"MMS_MAX_REQUEST_SIZE": "67108864"，即64MB），放宽消息大小限制。
• 检查实例资源状态：当前使用的ml.c6i.xlarge是CPU实例，若批量请求过大可能导致内存不足，进而引发Worker进程崩溃。可以监控实例的内存使用率，或临时切换到带GPU的实例（如ml.g4dn.xlarge）测试是否为资源瓶颈问题。
• 确认Embedding参数匹配：确保SageMakerEmbedding的输入处理逻辑与模型配置一致，比如是否自动处理超长文本的截断，或者是否需要手动设置max_length参数匹配MAX_INPUT_LENGTH=1024的配置。

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