如何使用SageMaker notebook instance本地测试SageMaker endpoint部署

SageMaker Notebook Instance支持本地部署模式，无需调用SageMaker托管端点调度资源，直接在实例本地启动推理容器，通常启动耗时在10秒以内，可以彻底解决Studio冷启动容器慢的问题。操作流程如下：

前置准备 #

• 确认Notebook Instance环境安装了最新版SageMaker SDK，执行命令更新：
  pip install --upgrade sagemaker
• 提前准备好你的模型推理脚本、模型权重S3路径，和你部署到托管端点的配置保持一致即可，不需要额外修改逻辑。

具体操作步骤 #

1. 初始化本地会话

import sagemaker
from sagemaker.local import LocalSession
# 替换为你实际使用的模型类，比如PyTorchModel、TensorFlowModel、SKLearnModel都适用
from sagemaker.huggingface import HuggingFaceModel

sagemaker_session = LocalSession()
sagemaker_session.config = {'local': {'local_code': True}}
role = sagemaker.get_execution_role()

2. 定义模型配置

model = HuggingFaceModel(
    model_data="s3://<你的模型存储桶路径>/model.tar.gz",
    role=role,
    entry_point="inference.py", # 你的推理入口脚本名称
    source_dir="./src", # 推理脚本所在的本地目录
    # 下面的框架版本根据你实际使用的配置替换即可
    transformers_version="4.28.1",
    pytorch_version="2.0.0",
    py_version="py310"
)

3. 本地部署端点

predictor = model.deploy(
    initial_instance_count=1,
    # 核心参数：指定instance_type为local，就会在Notebook Instance本地启动容器，不需要调度远程资源
    instance_type="local"
)

4. 测试推理

# 调用方式和托管端点完全一致，不需要修改调用逻辑
test_payload = {"inputs": "测试样例输入"}
response = predictor.predict(test_payload)
print(response)

5. 测试完成后清理资源

predictor.delete_endpoint()

常见异常处理 #

如果启动时提示docker权限错误，执行sudo usermod -aG docker ec2-user，重启Notebook内核后重试即可。
模型权重会自动从S3拉取到本地缓存，第二次部署相同模型时不需要重复拉取，启动速度会更快。

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