通过CloudFormation将预训练模型部署至SageMaker端点

核心问题解答 #

1. 能否实现部署？
   完全可以。SageMaker原生支持HuggingFace预训练模型的部署，结合CloudFormation能实现全流程自动化，无需手动操作控制台。

2. 模型名称怎么填？
   你目标模型all-MiniLM-L6-v2就是HuggingFace Hub上的官方模型ID，直接用这个字符串即可。SageMaker的HuggingFace推理容器会自动从Hub拉取对应模型。

3. 是否需要容器？
   需要，但不用从零构建——AWS提供了预构建的HuggingFace推理容器，支持PyTorch、TensorFlow等框架，直接使用这些镜像就行。比如（以us-east-1区域的CPU版PyTorch镜像为例）：
   763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04
   不同区域、框架版本的镜像前缀为对应区域的AWS账号ID，后面跟着具体的镜像标签，按需选择即可。

4. 容器怎么配置处理请求？
   需要准备一个inference.py脚本，上传到S3的指定目录（或和模型打包），脚本里实现两个核心函数：

• model_fn: 加载预训练模型和分词器
• predict_fn: 接收输入文本，生成嵌入并返回

示例inference.py代码：

from transformers import AutoModel, AutoTokenizer
import torch

def model_fn(model_dir):
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModel.from_pretrained(model_dir)
    return model, tokenizer

def predict_fn(input_data, model_and_tokenizer):
    model, tokenizer = model_and_tokenizer
    inputs = tokenizer(input_data["text"], return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    # 取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>token的嵌入作为文本表示
    embeddings = outputs.last_hidden_state[:, 0, :].numpy()
    return {"embeddings": embeddings.tolist()}

同时在CloudFormation的Model资源中设置环境变量HF_TASK=feature-extraction，容器会自动适配任务逻辑；如果用官方默认处理流程，甚至可以不用自定义inference.py，仅靠环境变量就能完成配置。

CloudFormation模板示例 #

以下是极简版的CloudFormation YAML模板，包含部署所需的核心资源：

AWSTemplateFormatVersion: '2010-09-09'
Resources:
  SageMakerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess

  HuggingFaceModel:
    Type: AWS::SageMaker::Model
    Properties:
      ExecutionRoleArn: !GetAtt SageMakerExecutionRole.Arn
      PrimaryContainer:
        Image: 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-inference:2.0.0-transformers4.28.1-cpu-py310-ubuntu20.04
        Environment:
          HF_MODEL_ID: all-MiniLM-L6-v2
          HF_TASK: feature-extraction
          # 若使用自定义inference.py，需设置S3路径：
          # SAGEMAKER_SUBMIT_DIRECTORY: s3://your-bucket/path/to/model-dir/

  SageMakerEndpointConfig:
    Type: AWS::SageMaker::EndpointConfig
    Properties:
      ProductionVariants:
        - ModelName: !Ref HuggingFaceModel
          VariantName: AllTraffic
          InitialInstanceCount: 1
          InstanceType: ml.t2.medium

  SageMakerEndpoint:
    Type: AWS::SageMaker::Endpoint
    Properties:
      EndpointConfigName: !Ref SageMakerEndpointConfig

说明：

• 替换Image字段为你所在区域的对应镜像
• 若使用自定义inference.py，将脚本打包成zip上传到S3，再设置SAGEMAKER_SUBMIT_DIRECTORY环境变量指向该路径
• 根据需求调整InstanceType（比如用GPU实例加速推理）

关于Lambda的替代说明 #

你提到的Lambda+层体积超限制问题确实存在，Lambda的250MB限制很难容纳HuggingFace模型和依赖包。直接用CloudFormation部署SageMaker资源是更合理的方案——全程无需Lambda参与，CloudFormation直接调用SageMaker API完成模型、端点配置和端点的创建。

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