SageMaker部署HuggingFaceModel时自定义inference.py脚本失效

问题描述 #

基于自有数据微调Llama-2模型后，打包成model.tar.gz（包含模型文件及code/inference.py、code/requirements.txt）上传至S3。使用HuggingFaceModel部署时，端点未执行inference.py中的自定义逻辑（如自定义output_fn的返回格式），仍沿用默认返回格式[{"generated_texts": model_response}]。尝试设置entry_point、source_dir、SAGEMAKER_PROGRAM环境变量均无效。

部署核心代码如下：

from sagemaker.huggingface import get_huggingface_llm_image_uri
from sagemaker.huggingface import HuggingFaceModel

llm_image = get_huggingface_llm_image_uri(
  "huggingface",
  version="0.9.3"
)

huggingface_model = HuggingFaceModel(
    model_data=model_s3_path,
    role=aws_role, 
    image_uri=llm_image,
    env={
      'HF_MODEL_ID': 'meta-llama/Llama-2-7b-hf',
      'SM_NUM_GPUS': '1',
      'MAX_INPUT_LENGTH': '2048',
      'MAX_TOTAL_TOKENS': '4096',
      'MAX_BATCH_TOTAL_TOKENS': '8192', 
      'HUGGING_FACE_HUB_TOKEN': "<my-hf-token>"
    }
)

huggingface_model.deploy(
    initial_instance_count=1,
    instance_type=inference_instance_type,
    endpoint_name=endpoint_name,
    container_startup_health_check_timeout=300
)

解决方案 #

1. 替换LLM专用镜像为通用HuggingFace框架镜像 #

当前使用的get_huggingface_llm_image_uri获取的是TGI（Text Generation Inference）优化镜像，这类镜像默认运行TGI服务，不支持加载自定义inference.py脚本。需切换到通用HuggingFace PyTorch/TensorFlow镜像：

from sagemaker.huggingface import get_huggingface_image_uri

# 获取通用HuggingFace PyTorch镜像（根据实际PyTorch版本调整参数）
image_uri = get_huggingface_image_uri(
    framework="pytorch",
    version="2.0.0",
    py_version="py310",
    instance_type=inference_instance_type
)

2. 调整HuggingFaceModel配置 #

移除TGI专用环境变量（如HF_MODEL_ID、MAX_INPUT_LENGTH等），同时正确指定自定义推理脚本路径：

• 若本地存在code目录（包含inference.py和requirements.txt）：

  huggingface_model = HuggingFaceModel(
      model_data=model_s3_path,
      role=aws_role, 
      image_uri=image_uri,
      entry_point="inference.py",
      source_dir="./code",  # 本地code文件夹路径
      env={
          'SM_NUM_GPUS': '1',
          'HUGGING_FACE_HUB_TOKEN': "<my-hf-token>"
      }
  )
  

• 若model.tar.gz中已包含code目录，可省略本地source_dir：

  huggingface_model = HuggingFaceModel(
      model_data=model_s3_path,
      role=aws_role, 
      image_uri=image_uri,
      entry_point="code/inference.py",
      env={
          'SM_NUM_GPUS': '1',
          'HUGGING_FACE_HUB_TOKEN': "<my-hf-token>"
      }
  )
  

3. 确保inference.py符合SageMaker规范 #

验证自定义脚本中的函数实现正确，示例如下：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def model_fn(model_dir):
    # 加载模型与tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_pretrained(
        model_dir,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    return model, tokenizer

def predict_fn(input_data, model_and_tokenizer):
    # 处理输入并生成结果
    model, tokenizer = model_and_tokenizer
    inputs = tokenizer(input_data["prompt"], return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=200)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def output_fn(prediction, accept):
    # 自定义返回格式
    return {"custom_response": prediction}, accept

4. 检查S3权限与文件结构 #

• 确认IAM角色拥有读取S3中model.tar.gz的权限；
• 验证model.tar.gz解压后结构正确，code目录下的文件无权限问题。

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