如何在Amazon SageMaker Estimator/Processor中列出指定目录文件

无需修改原有脚本的实现方式 #

针对 SageMaker Processor #

如果使用ShellProcessor或其他类型的Processor，直接在command参数中组合目录打印命令与原有任务命令即可。示例：

from sagemaker.processing import ShellProcessor, ProcessingInput, ProcessingOutput

processor = ShellProcessor(
    role=your_iam_role,
    instance_count=1,
    instance_type="ml.t3.medium"
)

processor.run(
    command=["ls", "/opt/ml/input/data/train", "&&", "python", "/opt/ml/processing/input/scripts/transfer_learning.py"],
    inputs=[ProcessingInput(source="s3://your-data-path", destination="/opt/ml/input/data/train")],
    outputs=[ProcessingOutput(source="/opt/ml/processing/output", destination="s3://your-output-path")]
)

该命令会先打印指定目录下的文件列表，再执行原有transfer_learning.py脚本。

针对 SageMaker Estimator #

无论是自定义容器还是官方框架容器（如TensorFlow、PyTorch），都可以通过command参数覆盖默认启动命令，先执行目录打印再运行原有脚本。示例（以PyTorch Estimator为例）：

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="transfer_learning.py",
    source_dir="./your-scripts-dir",
    role=your_iam_role,
    instance_count=1,
    instance_type="ml.m5.xlarge",
    framework_version="2.0",
    command=["bash", "-c", "ls /opt/ml/input/data/train && python transfer_learning.py"]
)

estimator.fit({"train": "s3://your-train-data-path"})

此方式无需修改transfer_learning.py的任何代码，直接通过启动命令实现目录打印。

需引入辅助脚本的实现方式 #

如果上述方法无法满足需求（比如需要更复杂的目录遍历逻辑），可以创建一个辅助启动脚本，在其中完成目录打印后调用原有脚本：

Python 辅助脚本示例 #

创建run_task.py：

import os
import subprocess

# 定义要打印的目标目录
TARGET_DIR = "/opt/ml/input/data/train"

# 打印目录文件列表
print(f"=== Files in {TARGET_DIR} ===")
for root, dirs, files in os.walk(TARGET_DIR):
    level = root.replace(TARGET_DIR, '').count(os.sep)
    indent = ' ' * 2 * level
    print(f"{indent}{os.path.basename(root)}/")
    subindent = ' ' * 2 * (level + 1)
    for file in files:
        print(f"{subindent}{file}")

# 执行原有训练脚本
subprocess.run(["python", "transfer_learning.py"], check=True)

然后将Estimator的entry_point指定为该辅助脚本：

estimator = PyTorch(
    entry_point="run_task.py",
    source_dir="./your-scripts-dir",  # 包含run_task.py和transfer_learning.py
    role=your_iam_role,
    instance_count=1,
    instance_type="ml.m5.xlarge",
    framework_version="2.0"
)

Shell 辅助脚本示例 #

创建start.sh（记得赋予可执行权限chmod +x start.sh）：

#!/bin/bash
set -e

# 打印目标目录的文件与子目录
echo "=== Files in /opt/ml/input/data/train ==="
find /opt/ml/input/data/train -type f -o -type d | sort

# 运行原有脚本
python transfer_learning.py

Estimator配置示例：

estimator = PyTorch(
    entry_point="start.sh",
    source_dir="./your-scripts-dir",
    role=your_iam_role,
    instance_count=1,
    instance_type="ml.m5.xlarge",
    framework_version="2.0"
)

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