如何将自定义模型集成到Amazon SageMaker Pipelines中？

针对JAX、PyMC3这类非开箱即用的框架，把它们集成到SageMaker流水线的预处理、训练等步骤中，主要有以下几种通用实践：

1. 脚本模式结合官方基础镜像快速适配 #

SageMaker的脚本模式允许你基于官方提供的Python基础镜像，通过自定义脚本和依赖配置引入第三方框架，无需构建完整镜像。

核心步骤： #

• 编写训练/预处理脚本（比如train_jax.py），在脚本中直接导入目标框架并实现业务逻辑；
• 准备requirements.txt文件，列出所需依赖（比如jax==0.4.10、pymc3==3.11.5）；
• 在流水线的训练/处理步骤中，使用Estimator或Processor类，指定官方Python镜像，通过dependencies参数传入依赖文件，也可在脚本开头通过pip install安装依赖。

代码示例（训练步骤）： #

from sagemaker.estimator import Estimator
from sagemaker.session import Session

# 使用官方Python 3.9基础镜像（根据实际区域调整账号ID）
image_uri = "763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-python-sdk:2.100.0"

estimator = Estimator(
    image_uri=image_uri,
    role="SageMakerExecutionRole",
    instance_count=1,
    instance_type="ml.m5.xlarge",
    entry_point="train_jax.py",
    dependencies=["requirements.txt"],
    sagemaker_session=Session()
)

# 启动训练（可集成到流水线的TrainingStep中）
estimator.fit({"train": "s3://your-bucket/train-data/"})

2. 构建自定义Docker镜像（适合复杂依赖场景） #

如果框架依赖特殊系统库、特定版本的底层软件，或需要固定运行环境，推荐构建自定义Docker镜像并推送到Amazon ECR，再在流水线中引用。

核心步骤： #

• 编写Dockerfile，基于基础镜像（比如python:3.9-slim）安装系统依赖、Python包；
• 配置SageMaker要求的容器入口点（比如指定train脚本，或使用sagemaker-containers工具包）；
• 将镜像推送到你的ECR仓库；
• 在流水线步骤中直接指定该ECR镜像的URI。

Dockerfile示例： #

# 基础镜像
FROM python:3.9-slim

# 安装系统依赖（比如JAX需要的编译依赖）
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc gfortran libopenblas-dev && \
    rm -rf /var/lib/apt/lists/*

# 安装Python依赖
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 设置SageMaker容器入口
ENV SAGEMAKER_PROGRAM train_jax.py
COPY train_jax.py /opt/ml/code/

WORKDIR /opt/ml/code

推送镜像到ECR的关键命令： #

# 登录ECR
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin <你的账号ID>.dkr.ecr.us-east-1.amazonaws.com

# 构建镜像
docker build -t custom-jax-image .

# 打标签
docker tag custom-jax-image:latest <你的账号ID>.dkr.ecr.us-east-1.amazonaws.com/custom-jax-image:latest

# 推送镜像
docker push <你的账号ID>.dkr.ecr.us-east-1.amazonaws.com/custom-jax-image:latest

3. 继承SageMaker框架基类封装可复用组件 #

如果需要将自定义框架封装成可复用的SageMaker组件，可以继承SageMaker SDK中的Framework或Estimator基类，实现必要的初始化和训练逻辑。

代码示例（自定义JAX Estimator）： #

from sagemaker.framework import Framework

class JAXEstimator(Framework):
    def __init__(
        self,
        entry_point,
        role,
        instance_count=1,
        instance_type="ml.m5.xlarge",
        framework_version="0.4.10",
        py_version="py39",
        **kwargs
    ):
        # 指定基础镜像
        image_uri = "763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-python-sdk:2.100.0"
        super().__init__(
            entry_point,
            role,
            instance_count,
            instance_type,
            image_uri=image_uri,
            py_version=py_version,
            **kwargs
        )
        self.framework_version = framework_version

# 在流水线中使用自定义Estimator
jax_estimator = JAXEstimator(
    entry_point="train_jax.py",
    role="SageMakerExecutionRole",
    instance_type="ml.m5.xlarge",
    dependencies=["requirements.txt"]
)

4. 通过SageMaker Processing集成自定义预处理/后处理 #

对于流水线中的预处理、后处理步骤，可以使用ProcessingJob结合自定义脚本或镜像，灵活集成目标框架的数据处理逻辑。

代码示例（预处理步骤）： #

from sagemaker.processing import Processor, ProcessingInput, ProcessingOutput

# 使用自定义镜像或官方Python镜像
processor = Processor(
    image_uri="<你的ECR镜像URI>",
    role="SageMakerExecutionRole",
    instance_count=1,
    instance_type="ml.m5.xlarge"
)

# 定义处理输入输出
processor.run(
    inputs=[ProcessingInput(source="s3://your-bucket/raw-data/", destination="/opt/ml/processing/input")],
    outputs=[ProcessingOutput(source="/opt/ml/processing/output", destination="s3://your-bucket/preprocessed-data/")],
    code="preprocess_pymc3.py"
)

# 将此步骤集成到流水线的ProcessingStep中

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