如何为Airflow Worker动态申请/释放EC2节点适配ML训练任务？

需求可行性确认 #

完全可行。结合EKS Cluster Autoscaler与Airflow的KubernetesExecutor，可实现：

• 触发大资源任务时自动扩容出匹配规格的EC2节点（如m5.8xlarge）
• 任务结束后，Pod销毁，Cluster Autoscaler自动缩容释放闲置节点
• 支持为不同DAG/任务配置独立的资源与调度规则

核心配置步骤 #

1. 确保EKS集群已部署Cluster Autoscaler #

这是节点自动扩缩容的基础，需提前在EKS集群中部署Cluster Autoscaler，并配置对应节点组的扩缩容范围（如最小0，最大N）。可调整缩容延迟参数（如scale-down-unneeded-time: 10m），确保任务结束后节点快速释放。

2. 修改Airflow Helm Chart的values.yaml配置 #

针对全局默认Pod模板，或为特定任务配置基础模板，需调整以下项：

启用KubernetesExecutor（已确认启用可跳过）

executor: KubernetesExecutor

配置默认Pod模板的资源与调度规则

在airflow.kubernetesPodTemplate下定义大资源任务的默认配置：

airflow:
  kubernetesPodTemplate:
    # 资源请求与限制，匹配m5.8xlarge的规格（24vCPU/64GiB）
    resources:
      requests:
        cpu: "24"
        memory: "64Gi"
      limits:
        cpu: "24"
        memory: "64Gi"
    
    # 节点选择器：指定调度到m5.8xlarge节点
    nodeSelector:
      kubernetes.io/instance-type: m5.8xlarge
    
    # 若大规格节点配置了污点，添加容忍规则
    tolerations:
      - key: "ml-large-node"
        operator: "Exists"
        effect: "NoSchedule"
    
    # 任务结束后自动删除Pod（触发缩容的前提）
    isDeleteOperatorPod: true

3. DAG/Operator层面的自定义调度（针对不同任务） #

若需要为不同DAG配置独立的节点规格或调度规则，直接在KubernetesPodOperator中覆盖全局配置即可：

from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

# 示例：ML训练任务的自定义Operator
ml_training_task = KubernetesPodOperator(
    task_id="weekly_ml_train",
    name="ml-train-pod",
    image="your-ml-training-image:v1",
    # 单独配置资源
    resources={
        "requests": {"cpu": "24", "memory": "64Gi"},
        "limits": {"cpu": "24", "memory": "64Gi"}
    },
    # 指定调度到m5.8xlarge节点
    node_selector={"kubernetes.io/instance-type": "m5.8xlarge"},
    # 容忍节点污点（若有）
    tolerations=[{"key": "ml-large-node", "operator": "Exists", "effect": "NoSchedule"}],
    # 任务结束立即删除Pod
    is_delete_operator_pod=True,
    namespace="airflow",
)

也可以通过pod_template_file参数引用外部Pod模板文件，实现更复杂的配置复用：

ml_training_task = KubernetesPodOperator(
    task_id="weekly_ml_train",
    pod_template_file="/path/to/ml-large-pod-template.yaml",
    is_delete_operator_pod=True,
    namespace="airflow",
)

关键注意事项 #

• 确保节点组的实例类型包含m5.8xlarge，且Cluster Autoscaler有权限调整节点组规模
• 资源请求需与节点规格匹配，避免因资源不足导致Pod pending或调度到错误节点
• is_delete_operator_pod必须设为true，否则Pod会残留，节点无法被缩容释放

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