Azure ML SDK v2管道组件Docker资源配置及OOM问题解决

你遇到的问题根源是组件定义中的resources字段不支持Docker级别的配置参数，shm_size、docker_args这类设置属于环境（Environment）的范畴，而非组件的集群资源配置，直接放在组件resources里会被SDK过滤，导致调用create_or_update后丢失。

以下是正确的实现方式：

方法1：创建并引用自定义环境（推荐） #

先定义包含Docker配置的环境并注册，再在组件中引用该环境：

1.1 用SDK创建带Docker配置的环境 #

from azure.ai.ml import Environment
from azure.ai.ml.entities import DockerConfiguration

# 配置Docker共享内存及额外参数
docker_config = DockerConfiguration(
    shm_size="16g",  # 根据你的OOM需求调整，比如8g/16g
    docker_args="--ulimit memlock=-1 --ulimit stack=67108864"  # 可选：添加其他Docker运行参数
)

# 定义环境，关联Docker配置
custom_env = Environment(
    name="unet_image_scoring_env",
    description="Custom environment for UNet image scoring with increased shared memory",
    docker_configuration=docker_config,
    conda_file="./conda.yml"  # 你的PyTorch依赖文件路径
)

# 注册环境到Azure ML工作区
ml_client.environments.create_or_update(custom_env)

1.2 在组件YAML中引用该环境 #

name: unet_image_scoring
display_name: UNet Image Scoring
version: 1
type: command
environment: azureml:unet_image_scoring_env:1  # 引用注册好的环境
code: ./src
command: >-
  python score.py --input_data ${{inputs.input_data}} --output_score ${{outputs.output_score}}
inputs:
  input_data:
    type: uri_folder
outputs:
  output_score:
    type: uri_folder
resources:
  instance_count: 1
  instance_type: Standard_NC6  # 根据你的计算需求选择实例类型

方法2：在组件YAML中内嵌环境配置 #

如果不想单独注册环境，也可以直接在组件YAML中定义环境及Docker配置：

name: unet_image_scoring
display_name: UNet Image Scoring
version: 1
type: command
environment:
  conda_file: ./conda.yml
  docker:
    image: mcr.microsoft.com/azureml/openmpi4.1.0-cuda11.3-cudnn8-ubuntu20.04  # 按需选择基础镜像
    docker_configuration:
      shm_size: "16g"
      docker_args: "--ulimit memlock=-1"
code: ./src
command: >-
  python score.py --input_data ${{inputs.input_data}} --output_score ${{outputs.output_score}}
inputs:
  input_data:
    type: uri_folder
outputs:
  output_score:
    type: uri_folder
resources:
  instance_count: 1

关键说明 #

• 组件的resources字段仅用于指定计算集群层面的资源，比如instance_count（实例数量）、instance_type（实例类型），不能存放Docker运行时参数。
• Docker相关配置必须放在环境的docker_configuration字段下，无论是单独注册的环境还是组件内嵌的环境，这样调用create_or_update时配置才会被完整保留。

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