如何在SageMaker Pipeline中配置私有镜像训练的TrainingImageConfig及AccessMode？

问题场景 #

在搭建包含处理和训练步骤的SageMaker Pipeline时，训练步骤从私有容器仓库拉取镜像，且已在Estimator中配置子网、VPC等网络参数，但运行时触发以下错误：

clienterror : Failed to invoke sagemaker: CreateTrainingJob.
Error Details: TrainingImageConfig with TrainingRepositoryAccessMode
set to VPC must be provided when using a training image from a private docker
registry...

现有核心代码片段：

from sagemaker.processing import ScriptProcessor, ProcessingInput, ProcessingOutput
from sagemaker.estimator import Estimator
...

security_group = ["group_id"]
my_subnets = ["subnet-public", "subnet-private"]

network_config = NetworkConfig(enable_network_isolation=False, 
                               security_group_ids=security_group, 
                               subnets=my_subnets)

# 训练步骤原有配置
estimator = Estimator(
    entry_point="src/train.py", 
    network_config=network_config, 
    image_uri='private_docker_registry/some_image', 
    ...  # 其他参数
)

解决方案 #

需要为Estimator添加training_image_config参数，明确指定通过VPC访问私有镜像仓库：

1. 导入TrainingImageConfig类（若未导入）：

from sagemaker.estimator import TrainingImageConfig

2. 创建镜像访问配置实例：

training_image_config = TrainingImageConfig(
    training_repository_access_mode='VPC'
)

3. 将配置传入Estimator：

# 修改后的训练步骤配置
estimator = Estimator(
    entry_point="src/train.py", 
    network_config=network_config, 
    image_uri='private_docker_registry/some_image',
    training_image_config=training_image_config,  # 添加该配置
    ...  # 保留其他原有参数
)

补充说明 #

• training_repository_access_mode='VPC'会让SageMaker训练任务通过你配置的子网和安全组访问私有镜像仓库，与已有的network_config配合生效。
• 若私有镜像仓库需要额外身份验证（如非AWS ECR的第三方仓库），可在TrainingImageConfig中添加training_repository_auth_config参数，传入包含凭证的字典配置。

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