在Sagemaker脚本模式下如何获取Python 3.9+的基础镜像？

核心思路 #

既然你的训练脚本不依赖XGBoost，完全不需要使用XGBoost专属Estimator，改用SageMaker的通用Python Estimator（sagemaker.estimator.Estimator）即可，它支持自由选择对应Python版本的官方镜像。

具体步骤 #

1. 替换Estimator导入
   把原XGBoost Estimator的导入语句替换为通用Estimator：

   from sagemaker.estimator import Estimator
   

2. 选择匹配Python版本的官方镜像
   SageMaker提供了多版本Python的官方训练镜像，根据你的脚本需求选择：

   • 适配Python 3.9：使用镜像763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:0.23-1-cpu-py39（将us-east-1替换为你的AWS区域代码）
   • 适配Python 3.10：使用镜像763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:1.2-1-cpu-py310
     注：这些镜像基于Scikit-learn构建，默认包含完整Python环境，足以运行无特定框架依赖的脚本；若需要更纯净的Python环境，也可选用SageMaker基础Python镜像。

3. 初始化通用Estimator
   修改代码为以下形式：

   est = Estimator(
       image_uri="763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-scikit-learn:0.23-1-cpu-py39", # 替换为对应Python版本的镜像
       entry_point=path_to_script,
       instance_type=instance_type,
       role=exec_role,
       instance_count=1,
       dependencies=[path_to_requirements],
       sagemaker_session=sm_session,
   )
   est.fit()
   

4. 验证镜像Python版本
   若不确定镜像支持的Python版本，可通过两种方式确认：

   • 查看SageMaker官方镜像列表中对应镜像的标签说明
   • 本地拉取镜像后执行docker run --rm <image_uri> python --version查看版本

额外提示 #

• dependencies参数依然有效，会自动安装requirements.txt中的依赖包
• 若脚本需要GPU支持，选择带-gpu后缀的镜像标签即可

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