如何通过Airflow实现Dataproc集群首次创建失败时使用不同配置重试

方案1：修正自定义DataprocCreateClusterOperator实现 #

你之前自定义Operator无法获取task_instance是对参数逻辑的误解：provider_context是PythonOperator的专属参数，所有继承自Airflow BaseOperator的自定义类无需额外配置，只要在execute方法中声明接收context参数即可直接获取上下文对象。
示例实现：

from airflow.providers.google.cloud.operators.dataproc import DataprocCreateClusterOperator
from airflow.utils.context import Context

class RetryableDataprocCreateClusterOperator(DataprocCreateClusterOperator):
    def execute(self, context: Context):
        # 直接从上下文获取任务实例，查询当前重试次数
        ti = context["task_instance"]
        current_retry = ti.try_number - 1  # try_number从1开始计数，减1得到实际重试次数

        # 重试时动态更新集群配置，比如切换可用区、调整子网等
        if current_retry > 0:
            self.cluster_config["gce_cluster_config"]["zone_uri"] = "us-central1-b"
        
        # 调用父类原生逻辑完成集群创建
        return super().execute(context)

该方案完全保留原生DataprocCreateClusterOperator的所有能力，支持编程动态修改配置，无冗余代码。

方案2：基于TaskFlow API动态传参 #

如果不想自定义Operator，可使用Airflow 2.0+的TaskFlow API实现动态配置+原生Operator调用：

from airflow.decorators import task
from airflow.providers.google.cloud.operators.dataproc import DataprocCreateClusterOperator

@task(provide_context=True)
def build_cluster_config(**context):
    ti = context["ti"]
    current_retry = ti.try_number - 1
    cluster_config = {
        # 填写基础集群配置
        "gce_cluster_config": {
            "zone_uri": "us-central1-a",
            # 其他基础配置
        }
    }
    # 重试时修改配置
    if current_retry > 0:
        cluster_config["gce_cluster_config"]["zone_uri"] = "us-central1-b"
    return cluster_config

# 动态接收配置，调用原生Operator创建集群
create_cluster = DataprocCreateClusterOperator(
    task_id="create_dataproc_cluster",
    project_id="你的GCP项目ID",
    region="us-central1",
    cluster_name="你的集群名称",
    cluster_config=build_cluster_config(),
    retries=2  # 配置最大重试次数
)

可选优化：指定重试触发条件 #

可通过retry_exceptions参数限定仅在资源不足类错误时触发重试，避免无效重试：

from google.api_core.exceptions import FailedPrecondition, ResourceExhausted

create_cluster.retry_exceptions = (FailedPrecondition, ResourceExhausted)

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