如何在含setup.py的Python包中正确定义SageMaker Estimator

问题场景 #

目录结构如下：

|-my_directory
|----- README.md
|----- setup.py
|----- src
|---------- my_train_script.py
|---------- __init__.py
|----- requirements.txt

使用SageMaker TensorFlow Estimator时，若指定source_dir="my_directory"、entry_point="src/my_train_script.py"，会触发错误：No module named src/my_train_script。将my_train_script.py移至根目录并设置entry_point="my_train_script.py"，或删除my_directory下的setup.py，代码均可正常运行，但需保留setup.py用于其他用途，需找到正确的Estimator定义方式。

错误配置示例：

TensorFlow(
    entry_point="src/my_train_script.py",         
    source_dir="my_directory", 

    role=get_execution_role(),
    instance_count=1, 
    instance_type="ml.m5.2xlarge",
    framework_version="2.10.1",
    py_version="py39",
    debugger_hook_config=None,
    disable_profiler=True,
    base_job_name="base_job_name",  
)

解决方法 #

当source_dir目录下存在setup.py时，SageMaker会将该目录作为Python包执行安装（运行pip install .），此时entry_point需以**Python模块路径（点分隔）**形式指定，而非文件路径。

修改后的正确配置：

TensorFlow(
    entry_point="src.my_train_script",  # 将文件路径替换为模块路径，斜杠改点
    source_dir="my_directory", 

    role=get_execution_role(),
    instance_count=1, 
    instance_type="ml.m5.2xlarge",
    framework_version="2.10.1",
    py_version="py39",
    debugger_hook_config=None,
    disable_profiler=True,
    base_job_name="base_job_name",  
)

原理说明 #

• 存在setup.py时，SageMaker会把my_directory内容打包安装到训练环境的Python路径中，此时src成为可导入的Python包，my_train_script.py是该包下的模块。
• 用模块路径src.my_train_script指定entry_point，SageMaker能正确定位训练脚本；原文件路径写法会因包安装后的目录结构变化，导致文件找不到。

内容的提问来源于stack exchange，提问作者Mohamed T. Bouzid