如何在SageMaker Pipeline预处理步骤中安装依赖包

你之前用setup.py没生效是因为SKLearn官方镜像不会自动执行setup.py的安装逻辑，换这几种方式就能解决：

方法一：用requirements.txt + source_dir（推荐） #

把预处理代码和依赖清单放在同一个目录下，比如创建processing_code文件夹，里面放preprocess.py和requirements.txt：

• requirements.txt内容：

sagemaker==2.93.0
matplotlib

然后修改sklearn_processor.run的参数，加上source_dir指定这个目录：

step_args = sklearn_processor.run(
    outputs=[
        ProcessingOutput(output_name="train", source="/opt/ml/processing/train"),
        ProcessingOutput(output_name="validation", source="/opt/ml/processing/validation"),
        ProcessingOutput(output_name="test", source="/opt/ml/processing/test"),
    ],
    source_dir=os.path.join(BASE_DIR, "processing_code"),  # 新增这行，指向代码目录
    code="preprocess.py",  # 这里只写文件名，不用全路径
    arguments=["--input-data", input_data],
)

SKLearnProcessor会自动识别source_dir里的requirements.txt，启动任务时先安装所有依赖。

方法二：自定义镜像 #

如果需要复杂环境配置（比如系统依赖、特殊版本库），可以基于官方SKLearn镜像构建自己的镜像：

1. 创建Dockerfile：

FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:0.23-1-cpu-py3
RUN pip install sagemaker==2.93.0 matplotlib

2. 把镜像推送到ECR，初始化SKLearnProcessor时指定image_uri：

sklearn_processor = SKLearnProcessor(
    framework_version="0.23-1",
    instance_type=processing_instance_type,
    instance_count=processing_instance_count,
    base_job_name=f"{base_job_prefix}/job-name",
    sagemaker_session=pipeline_session,
    role=role,
    image_uri="你的ECR镜像地址",  # 指定自定义镜像
)

方法三：预处理代码中直接安装（临时应急） #

在preprocess.py最开头添加pip安装命令，虽不优雅但能快速生效：

import subprocess
import sys

def install(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

install("sagemaker==2.93.0")
install("matplotlib")

# 你的预处理代码...

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