如何在Airflow中通过ClusterGenerator配置Dataproc集群公网IP

我正尝试在Airflow DAG中使用DataprocCreateClusterOperator在GCP项目中创建Dataproc集群，当前通过ClusterGenerator生成集群配置。但我希望指定在gcloud CLI中常用的--public-ip-address参数，发现ClusterGenerator并无该输入参数。

请问如何在ClusterGenerator或通用API格式中适配--public-ip-address的配置？

现有代码如下：

from airflow import models
from airflow.providers.google.cloud.operators.dataproc import (
    ClusterGenerator,
    DataprocCreateClusterOperator,
)
from google.api_core.retry import Retry

DATAPROC_CLUSTER_CONFIG = ClusterGenerator(
    project_id=GCP_PROJECT,
    region=GCP_REGION,
    master_machine_type="n2-standard-4",
    master_disk_type="pd-standard",
    master_disk_size=500,
    worker_machine_type="n2-standard-4",
    worker_disk_type="pd-standard",
    worker_disk_size=500,
    num_workers=2,
    image_version="2.2-ubuntu22",
    storage_bucket=DATAPROC_BUCKET,
    properties={
        "dataproc:pip.packages": "sentence-transformers==3.0.1,pydantic==2.8.2"
    },
    internal_ip_only=False,
    enable_component_gateway=True,
).make()

with models.DAG() as dag:
    create_dataproc_cluster = DataprocCreateClusterOperator(
        task_id="create_dataproc_cluster",
        project_id=GCP_PROJECT,
        cluster_config=DATAPROC_CLUSTER_CONFIG,
        region=GCP_REGION,
        cluster_name=DATAPROC_CLUSTER_NAME,
        retry=Retry(maximum=100.0, initial=10.0, multiplier=1.0),
    )

解决方案 #

gcloud的--public-ip-address参数对应Dataproc API中集群配置里的public_ip字段（属于InstanceGroupConfig下的instance_config），可以通过两种方式配置：

方法1：修改ClusterGenerator生成的配置字典 #

ClusterGenerator.make()会返回字典格式的集群配置，直接修改该字典即可添加public_ip设置：

DATAPROC_CLUSTER_CONFIG = ClusterGenerator(
    # 保留原有参数
    project_id=GCP_PROJECT,
    region=GCP_REGION,
    master_machine_type="n2-standard-4",
    master_disk_type="pd-standard",
    master_disk_size=500,
    worker_machine_type="n2-standard-4",
    worker_disk_type="pd-standard",
    worker_disk_size=500,
    num_workers=2,
    image_version="2.2-ubuntu22",
    storage_bucket=DATAPROC_BUCKET,
    properties={
        "dataproc:pip.packages": "sentence-transformers==3.0.1,pydantic==2.8.2"
    },
    internal_ip_only=False,
    enable_component_gateway=True,
).make()

# 为master和worker实例组添加公网IP配置
# True表示分配公网IP，False表示不分配
DATAPROC_CLUSTER_CONFIG['config']['master_config']['instance_config']['public_ip'] = True
DATAPROC_CLUSTER_CONFIG['config']['worker_config']['instance_config']['public_ip'] = True

方法2：直接构建API格式的集群配置 #

如果不想依赖ClusterGenerator，可以直接按照Dataproc API结构构建配置字典，完全自定义参数：

DATAPROC_CLUSTER_CONFIG = {
    "config": {
        "master_config": {
            "num_instances": 1,
            "machine_type_uri": "n2-standard-4",
            "disk_config": {
                "boot_disk_type": "pd-standard",
                "boot_disk_size_gb": 500
            },
            "instance_config": {
                "public_ip": True  # 控制master节点公网IP分配
            }
        },
        "worker_config": {
            "num_instances": 2,
            "machine_type_uri": "n2-standard-4",
            "disk_config": {
                "boot_disk_type": "pd-standard",
                "boot_disk_size_gb": 500
            },
            "instance_config": {
                "public_ip": True  # 控制worker节点公网IP分配
            }
        },
        "software_config": {
            "image_version": "2.2-ubuntu22",
            "properties": {
                "dataproc:pip.packages": "sentence-transformers==3.0.1,pydantic==2.8.2"
            }
        },
        "config_bucket": DATAPROC_BUCKET,
        "endpoint_config": {
            "enable_http_port_access": True  # 对应enable_component_gateway
        },
        "gce_cluster_config": {
            "internal_ip_only": False
        }
    }
}

注意事项 #

• internal_ip_only=True会强制所有实例仅使用内网IP，此时public_ip设置会被忽略
• 当internal_ip_only=False时，可通过public_ip单独控制master/worker实例是否分配公网IP

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