SageMaker训练任务成功后输出目录缺失model.tar.gz文件

问题根因 #

SageMaker托管PyTorch训练容器有固定的产物上传规则：训练进程退出后，容器只会将内部路径/opt/ml/model下的全部内容打包为model.tar.gz，再上传到estimator配置的output_path对应的S3路径。
从训练日志可以看到，训练产出默认保存在runs/train/exp相对路径下，该路径不属于/opt/ml/model目录，因此容器没有可打包的有效模型文件，最终S3中不会生成预期的model.tar.gz。

修复步骤 #

• 修改训练入口脚本train.py，在训练逻辑执行完成后，将所有需要留存的模型权重、配置文件等训练产物，复制或直接保存到容器内的/opt/ml/model目录下。参考实现代码如下：

import os
import shutil

# 训练逻辑执行完成后添加以下代码
# SageMaker模型固定输出目录，无需修改
SM_MODEL_DIR = "/opt/ml/model"
os.makedirs(SM_MODEL_DIR, exist_ok=True)

# 替换为训练脚本实际的产出保存路径，从日志看默认是runs/train/exp
local_train_output = "runs/train/exp"

# 拷贝根目录下的配置、指标等文件
for f in os.listdir(local_train_output):
    src = os.path.join(local_train_output, f)
    if os.path.isfile(src):
        shutil.copy2(src, os.path.join(SM_MODEL_DIR, f))

# 递归拷贝权重子目录（如YOLO类框架默认生成weights文件夹存储best/last权重）
weights_src = os.path.join(local_train_output, "weights")
if os.path.exists(weights_src):
    shutil.copytree(
        weights_src,
        os.path.join(SM_MODEL_DIR, "weights"),
        dirs_exist_ok=True
    )

• 现有PyTorch estimator的配置无需额外修改，保持原有output_path配置即可。训练容器默认运行用户对/opt/ml/model有完整读写权限，不需要额外配置权限参数。
• 重新启动训练任务后，可通过日志中Uploading generated training model阶段的上传大小判断是否打包成功：如果上传体积和实际生成的模型文件大小匹配，训练完成后即可在estimator.model_data返回的S3路径找到完整的model.tar.gz文件。

注意事项 #

不要将训练产物保存在容器内/opt/ml/model以外的路径，包括代码工作目录、临时目录、输入通道挂载目录。这些路径下的文件不会被容器自动打包上传，容器终止后会被永久清除。

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