Azure ML中使用含训练脚本的自定义Docker容器训练可行性咨询

核心结论 #

完全可以将包含训练脚本的自定义Docker容器上传至Azure Artifact Registry，并在Azure ML中触发训练，实现逻辑和Vertex AI的方式类似——核心是构建符合Azure ML要求的自定义容器，再通过AML作业定义指定该容器镜像即可。

具体实现步骤 #

1. 构建符合要求的自定义容器

   • 基础镜像可选择Azure ML官方提供的镜像（如mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04），也可使用自定义基础镜像。
   • 将训练脚本直接拷贝到容器内的指定路径，同时安装训练所需的依赖包。示例Dockerfile：

     # 基于Azure ML基础镜像
     FROM mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04
     
     # 安装训练依赖
     RUN pip install torch pandas scikit-learn
     
     # 拷贝本地训练脚本到容器内
     COPY train.py /app/train.py
     
     # 设置容器工作目录
     WORKDIR /app
     
     # 定义默认启动命令（AML作业可按需覆盖）
     CMD ["python", "train.py"]
     

2. 推送容器镜像到Azure Artifact Registry

   • 登录到你的Artifact Registry实例：

     az login
     az acr login --name <你的Artifact Registry名称>
     

   • 构建并标记镜像：

     docker build -t <你的Artifact Registry名称>.azurecr.io/training-container:v1 .
     

   • 推送镜像至Registry：

     docker push <你的Artifact Registry名称>.azurecr.io/training-container:v1
     

3. 在Azure ML中触发训练作业

   • 使用Azure ML Python SDK v2定义作业，指定自定义容器镜像路径与计算目标，示例代码：

     from azure.ai.ml import MLClient
     from azure.ai.ml.entities import CommandJob
     from azure.identity import DefaultAzureCredential
     
     # 初始化MLClient连接工作区
     ml_client = MLClient(
         credential=DefaultAzureCredential(),
         subscription_id="<你的订阅ID>",
         resource_group_name="<你的资源组名称>",
         workspace_name="<你的AML工作区名称>",
     )
     
     # 定义命令式训练作业
     job = CommandJob(
         command="python train.py",  # 可覆盖容器默认CMD，比如添加参数：python train.py --epochs 20
         environment="<你的Artifact Registry名称>.azurecr.io/training-container:v1",
         compute="<你的AML计算集群名称>",
         display_name="custom-container-training-job",
     )
     
     # 提交并监控作业
     returned_job = ml_client.jobs.create_or_update(job)
     ml_client.jobs.stream(returned_job.name)
     

关键注意事项 #

• 需为Azure ML工作区的系统分配身份添加Artifact Registry的AcrPull角色权限，确保工作区能拉取镜像。
• 若训练需要访问AML数据集、模型等资产，可在作业中配置数据挂载，或在容器内通过Azure ML SDK直接访问。
• 作业中的command参数可灵活调整，无需重新构建容器即可修改训练参数或执行逻辑。

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