Airflow优化：避免PythonOperator中get_cluster_id函数提前执行

我有多个包含EMR步骤的Airflow DAG，编写了get_cluster_id函数用于挑选活跃步骤（PENDING/RUNNING状态）最少的EMR集群ID，然后将该ID传入PythonOperator的op_kwargs中。当前遇到的问题是：这个函数会在DAG未运行的解析阶段就执行，每秒约触发5次API调用，最终触发ThrottlingException。需要修改代码，让get_cluster_id仅在DAG运行时才被调用。

原始代码示例：

def get_cluster_id(**kwargs):
        emr_client = boto3.client("emr", 'us-east-1')
        cluster_details = emr_client.list_clusters(ClusterStates=['STARTING','BOOTSTRAPPING','RUNNING','WAITING'])['Clusters']
        cluster_list = []
        for i in range(len(cluster_details)):
            if cluster_details[i]['Name'] in ['cluster1', 'cluster2']: 
                cluster_id = cluster_details[i]['Id']
                nr_of_active_steps = len(emr_client.list_steps(ClusterId=cluster_id, StepStates=['PENDING','RUNNING'])['Steps'])
                cluster_list.append({'cluster_id':cluster_id, 'nr_of_active_steps':nr_of_active_steps})
        return min(cluster_list, key=lambda x:x['nr_of_active_steps'])['cluster_id']   

# Example operator
task1 = PythonOperator(
    task_id='task1',
    provide_context=True,
    python_callable=emr_step_adder,
    retries = 1,
    retry_delay=timedelta(minutes=2),
    op_kwargs={'task_name': 'task1'
        ,'cluster_id': get_cluster_id()
        ,'arg_list':['spark-submit', 's3://test/test.py']
        ,'timeout_in_seconds': 1800
    },
    dag=dag
)

问题根源 #

你当前的代码在定义PythonOperator时，直接在op_kwargs里调用了get_cluster_id()，而Airflow会定期解析DAG文件（频率由配置决定），每次解析都会执行这个函数，因此才会频繁触发API调用，导致AWS限流。

解决方案 #

方法1：将集群ID获取逻辑整合到目标函数中 #

直接把get_cluster_id的逻辑放到emr_step_adder函数内部，这样只有当任务实际运行时才会执行API调用：

def emr_step_adder(**kwargs):
    # 先获取目标集群ID
    emr_client = boto3.client("emr", 'us-east-1')
    cluster_details = emr_client.list_clusters(ClusterStates=['STARTING','BOOTSTRAPPING','RUNNING','WAITING'])['Clusters']
    cluster_list = []
    # 优化循环写法，更简洁
    for cluster in cluster_details:
        if cluster['Name'] in ['cluster1', 'cluster2']: 
            cluster_id = cluster['Id']
            active_steps = emr_client.list_steps(ClusterId=cluster_id, StepStates=['PENDING','RUNNING'])['Steps']
            cluster_list.append({'cluster_id': cluster_id, 'nr_of_active_steps': len(active_steps)})
    target_cluster_id = min(cluster_list, key=lambda x: x['nr_of_active_steps'])['cluster_id']

    # 执行原本的EMR步骤添加逻辑
    # ... 这里写你原来emr_step_adder的代码，使用target_cluster_id即可

task1 = PythonOperator(
    task_id='task1',
    provide_context=True,
    python_callable=emr_step_adder,
    retries=1,
    retry_delay=timedelta(minutes=2),
    op_kwargs={
        'task_name': 'task1',
        'arg_list': ['spark-submit', 's3://test/test.py'],
        'timeout_in_seconds': 1800
    },
    dag=dag
)

方法2：拆分任务，用XCom传递集群ID #

单独创建一个获取集群ID的任务，将结果通过Airflow的XCom机制传递给后续的EMR任务，这样只有当这个获取任务运行时才会调用API：

def get_cluster_id(**kwargs):
    emr_client = boto3.client("emr", 'us-east-1')
    cluster_details = emr_client.list_clusters(ClusterStates=['STARTING','BOOTSTRAPPING','RUNNING','WAITING'])['Clusters']
    cluster_list = []
    for cluster in cluster_details:
        if cluster['Name'] in ['cluster1', 'cluster2']: 
            cluster_id = cluster['Id']
            active_steps = emr_client.list_steps(ClusterId=cluster_id, StepStates=['PENDING','RUNNING'])['Steps']
            cluster_list.append({'cluster_id': cluster_id, 'nr_of_active_steps': len(active_steps)})
    target_cluster_id = min(cluster_list, key=lambda x: x['nr_of_active_steps'])['cluster_id']
    # 将结果推送到XCom
    kwargs['ti'].xcom_push(key='target_cluster_id', value=target_cluster_id)

def emr_step_adder(**kwargs):
    # 从XCom中拉取集群ID
    target_cluster_id = kwargs['ti'].xcom_pull(task_ids='fetch_cluster_id', key='target_cluster_id')
    # 执行原本的EMR步骤添加逻辑
    # ... 这里写你原来emr_step_adder的代码

# 创建获取集群ID的任务
fetch_cluster_task = PythonOperator(
    task_id='fetch_cluster_id',
    provide_context=True,
    python_callable=get_cluster_id,
    dag=dag
)

# 创建EMR任务
task1 = PythonOperator(
    task_id='task1',
    provide_context=True,
    python_callable=emr_step_adder,
    retries=1,
    retry_delay=timedelta(minutes=2),
    op_kwargs={
        'task_name': 'task1',
        'arg_list': ['spark-submit', 's3://test/test.py'],
        'timeout_in_seconds': 1800
    },
    dag=dag
)

# 设置任务依赖：先获取集群ID，再执行EMR任务
fetch_cluster_task >> task1

方法3：使用可调用对象传入op_kwargs（Airflow 2.x+适用） #

如果你使用的是Airflow 2.x及以上版本，可以在op_kwargs中传入函数对象而非函数调用结果，Airflow会在任务运行时自动执行该函数：

def get_cluster_id():
    emr_client = boto3.client("emr", 'us-east-1')
    cluster_details = emr_client.list_clusters(ClusterStates=['STARTING','BOOTSTRAPPING','RUNNING','WAITING'])['Clusters']
    cluster_list = []
    for cluster in cluster_details:
        if cluster['Name'] in ['cluster1', 'cluster2']: 
            cluster_id = cluster['Id']
            active_steps = emr_client.list_steps(ClusterId=cluster_id, StepStates=['PENDING','RUNNING'])['Steps']
            cluster_list.append({'cluster_id': cluster_id, 'nr_of_active_steps': len(active_steps)})
    return min(cluster_list, key=lambda x: x['nr_of_active_steps'])['cluster_id']

def emr_step_adder(**kwargs):
    # 调用传入的函数获取集群ID
    target_cluster_id = kwargs['cluster_id']()
    # 执行原本的EMR步骤添加逻辑
    # ... 这里写你原来emr_step_adder的代码

task1 = PythonOperator(
    task_id='task1',
    provide_context=True,
    python_callable=emr_step_adder,
    retries=1,
    retry_delay=timedelta(minutes=2),
    op_kwargs={
        'task_name': 'task1',
        'cluster_id': get_cluster_id,  # 这里传入函数对象，不要加括号调用
        'arg_list': ['spark-submit', 's3://test/test.py'],
        'timeout_in_seconds': 1800
    },
    dag=dag
)

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