AWS SageMaker部署Instructor模型调用predict()报错的排查与脚本优化问询

核心问题分析 #

从日志中的MemoryError和OSError: [Errno 14] Bad address可以明确：

1. 内存不足：ml.r5.xlarge是CPU实例（16GB内存），无法加载10亿参数的hkunlp/instructor-xl大模型，导致加载失败。
2. 输入参数拼写错误：调用predict时传入的是'embeding_instruction'，但脚本中读取的是'embedding_instruction'（少了一个字母d），即使模型加载成功也会触发参数缺失错误。
3. 输入解析逻辑缺失：默认情况下SageMaker会把JSON输入传递为字符串，input_fn未做反序列化处理，导致后续无法读取字典参数。
4. 模型加载效率低：直接在model_fn中从Hugging Face Hub下载模型，会导致每次启动都重复下载，且未利用SageMaker本地模型目录缓存。

分步修复方案 #

1. 更换实例类型（解决内存核心问题） #

将实例替换为带GPU的ml.g4dn.xlarge（16GB GPU内存）或更高规格，GPU加载大模型的内存效率远高于CPU，可顺利加载instructor-xl。

2. 修正推理脚本inference.py #

from InstructorEmbedding import INSTRUCTOR
import torch
import json

def model_fn(model_dir):
    # 优先从本地模型目录加载，避免重复下载
    try:
        model = INSTRUCTOR(model_dir)
    except:
        # 本地无模型时从Hub下载（首次部署用）
        model = INSTRUCTOR('hkunlp/instructor-xl')
    model.max_seq_length = 768
    # 自动迁移到可用设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model

def input_fn(input_data, content_type):
    # 解析JSON输入为字典格式
    if content_type == 'application/json':
        return json.loads(input_data)
    raise ValueError(f"不支持的内容类型: {content_type}")

def output_fn(prediction, accept):
    # 确保返回JSON格式结果
    if accept == 'application/json':
        return json.dumps(prediction), accept
    raise ValueError(f"不支持的返回格式: {accept}")

def predict_fn(processed_data, model):
    try:
        data = processed_data['data']
        embedding_instruction = processed_data['embedding_instruction']
        
        documents = data['documents']
        metadatas = data.get('metadatas', [])
        ids = data.get('ids', [])
        
        # 构造模型要求的输入结构：[[指令, 文本], ...]
        inputs = [[embedding_instruction, doc] for doc in documents]
        # 使用模型已加载的设备，无需重复指定
        embeddings = model.encode(inputs)
        
        return {
            "embeddings": embeddings.tolist(), 
            "metadatas": metadatas, 
            "ids": ids
        }
    except KeyError as e:
        raise ValueError(f"缺失必填参数: {str(e)}")

3. 修正requirements.txt版本兼容性 #

将transformers版本与SageMaker指定的4.26对齐，避免版本冲突：

InstructorEmbedding>=1.0.1
transformers==4.26.0
datasets>=2.2.0
jsonlines
numpy
requests>=2.26.0
scikit_learn>=1.0.2
scipy
sentence_transformers>=2.2.0
torch==1.13.1
tqdm
rich

4. 修正调用时的输入参数拼写 #

data = {
    'data': {
        'documents': ['Which paper discuss anemia?'],
        'metadatas': ['NONE'],
        'ids': ['NONE']
    },
    # 修正拼写为embedding_instruction
    'embedding_instruction': "Represent the Medical title for retrieving relevant documents:"
}
hf_predictor.predict(data)

5. 优化模型打包流程 #

提前下载模型到本地再打包，避免部署时重复下载：

# 预先下载模型到本地model目录
python -c "from InstructorEmbedding import INSTRUCTOR; model = INSTRUCTOR('hkunlp/instructor-xl'); model.save_pretrained('./model')"

# 将code目录和本地模型文件一起打包
tar zcvf model.tar.gz code/ model/

# 上传到S3存储桶
aws s3 cp model.tar.gz $s3_location

验证步骤 #

1. 使用新打包的model.tar.gz重新创建HuggingFaceModel并部署到GPU实例。
2. 调用predict时传入修正后的参数。
3. 查看CloudWatch日志，确认模型加载成功且无内存错误。

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