"""Contains possible Type values of Preemptibility applicable for every secondary worker of Cluster."""

PREEMPTIBLE = "PREEMPTIBLE"

SPOT = "SPOT"

PREEMPTIBILITY_UNSPECIFIED = "PREEMPTIBILITY_UNSPECIFIED"

"""Defines machines types and a rank to which the machines types belong.

Representation for

google.cloud.dataproc.v1#google.cloud.dataproc.v1.InstanceFlexibilityPolicy.InstanceSelection.

:param machine_types: Full machine-type names, e.g. "n1-standard-16".

:param rank: Preference of this instance selection. Lower number means higher preference.

Dataproc will first try to create a VM based on the machine-type with priority rank and fallback

to next rank based on availability. Machine types and instance selections with the same priority have

the same preference.

"""

machine_types: list[str]

"""

Instance flexibility Policy allowing a mixture of VM shapes and provisioning models.

Representation for google.cloud.dataproc.v1#google.cloud.dataproc.v1.InstanceFlexibilityPolicy.

:param instance_selection_list: List of instance selection options that the group will use when

creating new VMs.

"""

"""Create a new Dataproc Cluster.

:param cluster_name: The name of the DataProc cluster to create. (templated)

:param project_id: The ID of the google cloud project in which

to create the cluster. (templated)

:param num_workers: The # of workers to spin up. If set to zero will

spin up cluster in a single node mode

:param min_num_workers: The minimum number of primary worker instances to create.

If more than ``min_num_workers`` VMs are created out of ``num_workers``, the failed VMs will be

deleted, cluster is resized to available VMs and set to RUNNING.

If created VMs are less than ``min_num_workers``, the cluster is placed in ERROR state. The failed

VMs are not deleted.

:param storage_bucket: The storage bucket to use, setting to None lets dataproc

generate a custom one for you

:param init_actions_uris: List of GCS uri's containing

dataproc initialization scripts

:param init_action_timeout: Amount of time executable scripts in

init_actions_uris has to complete

:param metadata: dict of key-value google compute engine metadata entries

to add to all instances

:param image_version: the version of software inside the Dataproc cluster

:param custom_image: custom Dataproc image for more info see

https://cloud.google.com/dataproc/docs/guides/dataproc-images

:param custom_image_project_id: project id for the custom Dataproc image, for more info see

https://cloud.google.com/dataproc/docs/guides/dataproc-images

:param custom_image_family: family for the custom Dataproc image,

family name can be provide using --family flag while creating custom image, for more info see

https://cloud.google.com/dataproc/docs/guides/dataproc-images

:param autoscaling_policy: The autoscaling policy used by the cluster. Only resource names

including projectid and location (region) are valid. Example:

``projects/[projectId]/locations/[dataproc_region]/autoscalingPolicies/[policy_id]``

:param properties: dict of properties to set on

config files (e.g. spark-defaults.conf), see

https://cloud.google.com/dataproc/docs/reference/rest/v1/projects.regions.clusters#SoftwareConfig

:param optional_components: List of optional cluster components, for more info see

https://cloud.google.com/dataproc/docs/reference/rest/v1/ClusterConfig#Component

:param num_masters: The # of master nodes to spin up

:param master_machine_type: Compute engine machine type to use for the primary node

:param master_disk_type: Type of the boot disk for the primary node

(default is ``pd-standard``).

Valid values: ``pd-ssd`` (Persistent Disk Solid State Drive) or

``pd-standard`` (Persistent Disk Hard Disk Drive).

:param master_disk_size: Disk size for the primary node

:param master_accelerator_type: Type of the accelerator card (GPU) to attach to the primary node,

see https://cloud.google.com/dataproc/docs/reference/rest/v1/InstanceGroupConfig#acceleratorconfig

:param master_accelerator_count: Number of accelerator cards (GPUs) to attach to the primary node

:param worker_machine_type: Compute engine machine type to use for the worker nodes

:param worker_disk_type: Type of the boot disk for the worker node

(default is ``pd-standard``).

Valid values: ``pd-ssd`` (Persistent Disk Solid State Drive) or

``pd-standard`` (Persistent Disk Hard Disk Drive).

:param worker_disk_size: Disk size for the worker nodes

:param worker_accelerator_type: Type of the accelerator card (GPU) to attach to the worker nodes,

see https://cloud.google.com/dataproc/docs/reference/rest/v1/InstanceGroupConfig#acceleratorconfig

:param worker_accelerator_count: Number of accelerator cards (GPUs) to attach to the worker nodes

:param num_preemptible_workers: The # of VM instances in the instance group as secondary workers

inside the cluster with Preemptibility enabled by default.

Note, that it is not possible to mix non-preemptible and preemptible secondary workers in

one cluster.

:param preemptibility: The type of Preemptibility applicable for every secondary worker, see

https://cloud.google.com/dataproc/docs/reference/rpc/ \

google.cloud.dataproc.v1#google.cloud.dataproc.v1.InstanceGroupConfig.Preemptibility

:param zone: The zone where the cluster will be located. Set to None to auto-zone. (templated)

:param network_uri: The network uri to be used for machine communication, cannot be

specified with subnetwork_uri

:param subnetwork_uri: The subnetwork uri to be used for machine communication,

cannot be specified with network_uri

:param internal_ip_only: If true, all instances in the cluster will only

have internal IP addresses. This can only be enabled for subnetwork

enabled networks

:param tags: The GCE tags to add to all instances

:param region: The specified region where the dataproc cluster is created.

:param gcp_conn_id: The connection ID to use connecting to Google Cloud.

:param service_account: The service account of the dataproc instances.

:param service_account_scopes: The URIs of service account scopes to be included.

:param idle_delete_ttl: The longest duration that cluster would keep alive while

staying idle. Passing this threshold will cause cluster to be auto-deleted.

A duration in seconds.

:param auto_delete_time: The time when cluster will be auto-deleted.

:param auto_delete_ttl: The life duration of cluster, the cluster will be

auto-deleted at the end of this duration.

A duration in seconds. (If auto_delete_time is set this parameter will be ignored)

:param customer_managed_key: The customer-managed key used for disk encryption

``projects/[PROJECT_STORING_KEYS]/locations/[LOCATION]/keyRings/[KEY_RING_NAME]/cryptoKeys/[KEY_NAME]`` # noqa

:param enable_component_gateway: Provides access to the web interfaces of default and selected optional

components on the cluster.

:param driver_pool_size: The number of driver nodes in the node group.

:param driver_pool_id: The ID for the driver pool. Must be unique within the cluster. Use this ID to

identify the driver group in future operations, such as resizing the node group.

:param secondary_worker_instance_flexibility_policy: Instance flexibility Policy allowing a mixture of VM

shapes and provisioning models.

:param secondary_worker_accelerator_type: Type of the accelerator card (GPU) to attach to the secondary workers,

see https://cloud.google.com/dataproc/docs/reference/rest/v1/InstanceGroupConfig#acceleratorconfig

:param secondary_worker_accelerator_count: Number of accelerator cards (GPUs) to attach to the secondary workers

"""

def __init__(

self,

project_id: str,

num_workers: int | None = None,

min_num_workers: int | None = None,

zone: str | None = None,

network_uri: str | None = None,

subnetwork_uri: str | None = None,

internal_ip_only: bool | None = None,

tags: list[str] | None = None,

storage_bucket: str | None = None,

init_actions_uris: list[str] | None = None,

init_action_timeout: str = "10m",

metadata: dict | None = None,

custom_image: str | None = None,

custom_image_project_id: str | None = None,

custom_image_family: str | None = None,

image_version: str | None = None,

autoscaling_policy: str | None = None,

properties: dict | None = None,

optional_components: list[str] | None = None,

num_masters: int = 1,

master_machine_type: str = "n1-standard-4",

master_disk_type: str = "pd-standard",

master_disk_size: int = 1024,

master_accelerator_type: str | None = None,

master_accelerator_count: int | None = None,

worker_machine_type: str = "n1-standard-4",

worker_disk_type: str = "pd-standard",

worker_disk_size: int = 1024,

worker_accelerator_type: str | None = None,

worker_accelerator_count: int | None = None,

num_preemptible_workers: int = 0,

preemptibility: str = PreemptibilityType.PREEMPTIBLE.value,

service_account: str | None = None,

service_account_scopes: list[str] | None = None,

idle_delete_ttl: int | None = None,

auto_delete_time: datetime | None = None,

auto_delete_ttl: int | None = None,

customer_managed_key: str | None = None,

enable_component_gateway: bool | None = False,

driver_pool_size: int = 0,

driver_pool_id: str | None = None,

secondary_worker_instance_flexibility_policy: InstanceFlexibilityPolicy | None = None,

secondary_worker_accelerator_type: str | None = None,

secondary_worker_accelerator_count: int | None = None,

**kwargs,

) -> None:

self.project_id = project_id

self.num_masters = num_masters

self.num_workers = num_workers

self.min_num_workers = min_num_workers

self.num_preemptible_workers = num_preemptible_workers

self.preemptibility = self._set_preemptibility_type(preemptibility)

self.storage_bucket = storage_bucket

self.init_actions_uris = init_actions_uris

self.init_action_timeout = init_action_timeout

self.metadata = metadata

self.custom_image = custom_image

self.custom_image_project_id = custom_image_project_id

self.custom_image_family = custom_image_family

self.image_version = image_version

self.properties = properties or {}

self.optional_components = optional_components

self.master_machine_type = master_machine_type

self.master_disk_type = master_disk_type

self.master_disk_size = master_disk_size

self.master_accelerator_type = master_accelerator_type

self.master_accelerator_count = master_accelerator_count

self.autoscaling_policy = autoscaling_policy

self.worker_machine_type = worker_machine_type

self.worker_disk_type = worker_disk_type

self.worker_disk_size = worker_disk_size

self.worker_accelerator_type = worker_accelerator_type

self.worker_accelerator_count = worker_accelerator_count

self.zone = zone

self.network_uri = network_uri

self.subnetwork_uri = subnetwork_uri

self.internal_ip_only = internal_ip_only

self.tags = tags

self.service_account = service_account

self.service_account_scopes = service_account_scopes

self.idle_delete_ttl = idle_delete_ttl

self.auto_delete_time = auto_delete_time

self.auto_delete_ttl = auto_delete_ttl

self.customer_managed_key = customer_managed_key

self.enable_component_gateway = enable_component_gateway

self.single_node = num_workers == 0

self.driver_pool_size = driver_pool_size

self.driver_pool_id = driver_pool_id

self.secondary_worker_instance_flexibility_policy = secondary_worker_instance_flexibility_policy

self.secondary_worker_accelerator_type = secondary_worker_accelerator_type

self.secondary_worker_accelerator_count = secondary_worker_accelerator_count

if self.custom_image and self.image_version:

raise ValueError("The custom_image and image_version can't be both set")

if self.custom_image_family and self.image_version:

raise ValueError("The image_version and custom_image_family can't be both set")

if self.custom_image_family and self.custom_image:

raise ValueError("The custom_image and custom_image_family can't be both set")

if self.single_node and self.num_preemptible_workers > 0:

raise ValueError("Single node cannot have preemptible workers.")

if self.min_num_workers:

if not self.num_workers:

raise ValueError("Must specify num_workers when min_num_workers are provided.")

if self.min_num_workers > self.num_workers:

raise ValueError(

"The value of min_num_workers must be less than or equal to num_workers. "

f"Provided {self.min_num_workers}(min_num_workers) and {self.num_workers}(num_workers)."

)

def _set_preemptibility_type(self, preemptibility: str):

return PreemptibilityType(preemptibility.upper())

def _get_init_action_timeout(self) -> dict:

match = re.fullmatch(r"(\d+)([sm])", self.init_action_timeout)

if match:

val = int(match.group(1))

unit = match.group(2)

if unit == "s":

return {"seconds": val}

elif unit == "m":

return {"seconds": int(timedelta(minutes=val).total_seconds())}

raise AirflowException(

"DataprocClusterCreateOperator init_action_timeout"

" should be expressed in minutes or seconds. i.e. 10m, 30s"

)

def _build_gce_cluster_config(self, cluster_data):

# This variable is created since same string was being used multiple times

config = "gce_cluster_config"

if self.zone:

zone_uri = f"https://www.googleapis.com/compute/v1/projects/{self.project_id}/zones/{self.zone}"

cluster_data[config]["zone_uri"] = zone_uri

if self.metadata:

cluster_data[config]["metadata"] = self.metadata

if self.network_uri:

cluster_data[config]["network_uri"] = self.network_uri

if self.subnetwork_uri:

cluster_data[config]["subnetwork_uri"] = self.subnetwork_uri

if self.internal_ip_only is not None:

if not self.subnetwork_uri and self.internal_ip_only:

raise AirflowException("Set internal_ip_only to true only when you pass a subnetwork_uri.")

cluster_data[config]["internal_ip_only"] = self.internal_ip_only

if self.tags:

cluster_data[config]["tags"] = self.tags

if self.service_account:

cluster_data[config]["service_account"] = self.service_account

if self.service_account_scopes:

cluster_data[config]["service_account_scopes"] = self.service_account_scopes

return cluster_data

def _build_lifecycle_config(self, cluster_data):

# This variable is created since same string was being used multiple times

lifecycle_config = "lifecycle_config"

if self.idle_delete_ttl:

cluster_data[lifecycle_config]["idle_delete_ttl"] = {"seconds": self.idle_delete_ttl}

if self.auto_delete_time:

utc_auto_delete_time = timezone.convert_to_utc(self.auto_delete_time)

cluster_data[lifecycle_config]["auto_delete_time"] = utc_auto_delete_time.strftime(

"%Y-%m-%dT%H:%M:%S.%fZ"

)

elif self.auto_delete_ttl:

cluster_data[lifecycle_config]["auto_delete_ttl"] = {"seconds": int(self.auto_delete_ttl)}

return cluster_data

def _build_driver_pool(self):

driver_pool = {

"node_group": {

"roles": ["DRIVER"],

"node_group_config": {"num_instances": self.driver_pool_size},

},

}

if self.driver_pool_id:

driver_pool["node_group_id"] = self.driver_pool_id

return driver_pool

def _build_cluster_data(self):

if self.zone:

master_type_uri = (

f"projects/{self.project_id}/zones/{self.zone}/machineTypes/{self.master_machine_type}"

)

worker_type_uri = (

f"projects/{self.project_id}/zones/{self.zone}/machineTypes/{self.worker_machine_type}"

)

else:

master_type_uri = self.master_machine_type

worker_type_uri = self.worker_machine_type

cluster_data = {

"gce_cluster_config": {},

"master_config": {

"num_instances": self.num_masters,

"machine_type_uri": master_type_uri,

"disk_config": {

"boot_disk_type": self.master_disk_type,

"boot_disk_size_gb": self.master_disk_size,

},

},

"worker_config": {

"num_instances": self.num_workers,

"machine_type_uri": worker_type_uri,

"disk_config": {

"boot_disk_type": self.worker_disk_type,

"boot_disk_size_gb": self.worker_disk_size,

},

},

"secondary_worker_config": {},

"software_config": {},

"lifecycle_config": {},

"encryption_config": {},

"autoscaling_config": {},

"endpoint_config": {},

}

if self.min_num_workers:

cluster_data["worker_config"]["min_num_instances"] = self.min_num_workers

if self.master_accelerator_type:

cluster_data["master_config"]["accelerators"] = {

"accelerator_type_uri": self.master_accelerator_type,

"accelerator_count": self.master_accelerator_count,

}

if self.worker_accelerator_type:

cluster_data["worker_config"]["accelerators"] = {

"accelerator_type_uri": self.worker_accelerator_type,

"accelerator_count": self.worker_accelerator_count,

}

if self.num_preemptible_workers > 0:

cluster_data["secondary_worker_config"] = {

"num_instances": self.num_preemptible_workers,

"machine_type_uri": worker_type_uri,

"disk_config": {

"boot_disk_type": self.worker_disk_type,

"boot_disk_size_gb": self.worker_disk_size,

},

"is_preemptible": True,

"preemptibility": self.preemptibility.value,

}

if self.worker_accelerator_type:

cluster_data["secondary_worker_config"]["accelerators"] = {

"accelerator_type_uri": self.secondary_worker_accelerator_type,

"accelerator_count": self.secondary_worker_accelerator_count,

}

if self.secondary_worker_instance_flexibility_policy:

cluster_data["secondary_worker_config"]["instance_flexibility_policy"] = {

"instance_selection_list": [

vars(s)

for s in self.secondary_worker_instance_flexibility_policy.instance_selection_list

]

}

if self.storage_bucket:

cluster_data["config_bucket"] = self.storage_bucket

if self.image_version:

cluster_data["software_config"]["image_version"] = self.image_version

elif self.custom_image:

project_id = self.custom_image_project_id or self.project_id

custom_image_url = (

f"https://www.googleapis.com/compute/beta/projects/{project_id}"

f"/global/images/{self.custom_image}"

)

cluster_data["master_config"]["image_uri"] = custom_image_url

if not self.single_node:

cluster_data["worker_config"]["image_uri"] = custom_image_url

elif self.custom_image_family:

project_id = self.custom_image_project_id or self.project_id

custom_image_url = (

"https://www.googleapis.com/compute/beta/projects/"

f"{project_id}/global/images/family/{self.custom_image_family}"

)

cluster_data["master_config"]["image_uri"] = custom_image_url

if not self.single_node:

cluster_data["worker_config"]["image_uri"] = custom_image_url

if self.driver_pool_size > 0:

cluster_data["auxiliary_node_groups"] = [self._build_driver_pool()]

cluster_data = self._build_gce_cluster_config(cluster_data)

if self.single_node:

self.properties["dataproc:dataproc.allow.zero.workers"] = "true"

if self.properties:

cluster_data["software_config"]["properties"] = self.properties

if self.optional_components:

cluster_data["software_config"]["optional_components"] = self.optional_components

cluster_data = self._build_lifecycle_config(cluster_data)

if self.init_actions_uris:

init_actions_dict = [

{"executable_file": uri, "execution_timeout": self._get_init_action_timeout()}

for uri in self.init_actions_uris

]

cluster_data["initialization_actions"] = init_actions_dict

if self.customer_managed_key:

cluster_data["encryption_config"] = {"gce_pd_kms_key_name": self.customer_managed_key}

if self.autoscaling_policy:

cluster_data["autoscaling_config"] = {"policy_uri": self.autoscaling_policy}

if self.enable_component_gateway:

cluster_data["endpoint_config"] = {"enable_http_port_access": self.enable_component_gateway}

return cluster_data

def make(self):

"""

Act as a helper method for easier migration.

:return: Dict representing Dataproc cluster.

"""

"""Create a new cluster on Google Cloud Dataproc.

The operator will wait until the creation is successful or an error occurs

in the creation process.

If the cluster already exists and ``use_if_exists`` is True, then the operator will:

- if cluster state is ERROR then delete it if specified and raise error

- if cluster state is CREATING wait for it and then check for ERROR state

- if cluster state is DELETING wait for it and then create new cluster

Please refer to

https://cloud.google.com/dataproc/docs/reference/rest/v1/projects.regions.clusters

for a detailed explanation on the different parameters. Most of the configuration

parameters detailed in the link are available as a parameter to this operator.

.. seealso::

For more information on how to use this operator, take a look at the guide:

:ref:`howto/operator:DataprocCreateClusterOperator`

:param project_id: The ID of the Google cloud project in which

to create the cluster. (templated)

:param cluster_name: Name of the cluster to create

:param labels: Labels that will be assigned to created cluster. Please, notice that

adding labels to ClusterConfig object in cluster_config parameter will not lead

to adding labels to the cluster. Labels for the clusters could be only set by passing

values to parameter of DataprocCreateCluster operator.

:param cluster_config: Required. The cluster config to create.

If a dict is provided, it must be of the same form as the protobuf message

:class:`~google.cloud.dataproc_v1.types.ClusterConfig`

:param virtual_cluster_config: Optional. The virtual cluster config, used when creating a Dataproc

cluster that does not directly control the underlying compute resources, for example, when creating a

`Dataproc-on-GKE cluster

<https://cloud.google.com/dataproc/docs/concepts/jobs/dataproc-gke#create-a-dataproc-on-gke-cluster>`

:param region: The specified region where the dataproc cluster is created.

:param delete_on_error: If true the cluster will be deleted if created with ERROR state. Default

value is true.

:param use_if_exists: If true use existing cluster

:param request_id: Optional. A unique id used to identify the request. If the server receives two

``DeleteClusterRequest`` requests with the same id, then the second request will be ignored and the

first ``google.longrunning.Operation`` created and stored in the backend is returned.

:param retry: A retry object used to retry requests. If ``None`` is specified, requests will not be

retried.

:param timeout: The amount of time, in seconds, to wait for the request to complete. Note that if

``retry`` is specified, the timeout applies to each individual attempt.

:param metadata: Additional metadata that is provided to the method.

:param gcp_conn_id: The connection ID to use connecting to Google Cloud.

:param impersonation_chain: Optional service account to impersonate using short-term

credentials, or chained list of accounts required to get the access_token

of the last account in the list, which will be impersonated in the request.

If set as a string, the account must grant the originating account

the Service Account Token Creator IAM role.

If set as a sequence, the identities from the list must grant

Service Account Token Creator IAM role to the directly preceding identity, with first

account from the list granting this role to the originating account (templated).

:param deferrable: Run operator in the deferrable mode.

:param polling_interval_seconds: Time (seconds) to wait between calls to check the run status.

"""

template_fields: Sequence[str] = (

"project_id",

"region",

"cluster_config",

"virtual_cluster_config",

"cluster_name",

"labels",

"impersonation_chain",

)

template_fields_renderers = {"cluster_config": "json", "virtual_cluster_config": "json"}

operator_extra_links = (DataprocClusterLink(),)

def __init__(

self,

*,

cluster_name: str,

region: str,

project_id: str = PROVIDE_PROJECT_ID,

cluster_config: dict | Cluster | None = None,

virtual_cluster_config: dict | None = None,

labels: dict | None = None,

request_id: str | None = None,

delete_on_error: bool = True,

use_if_exists: bool = True,

retry: AsyncRetry | _MethodDefault = DEFAULT,

timeout: float = 1 * 60 * 60,

metadata: Sequence[tuple[str, str]] = (),

gcp_conn_id: str = "google_cloud_default",

impersonation_chain: str | Sequence[str] | None = None,

deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),

polling_interval_seconds: int = 10,

**kwargs,

) -> None:

# TODO: remove one day

if cluster_config is None and virtual_cluster_config is None:

warnings.warn(

f"Passing cluster parameters by keywords to `{type(self).__name__}` will be deprecated. "

"Please provide cluster_config object using `cluster_config` parameter. "

"You can use `airflow.dataproc.ClusterGenerator.generate_cluster` "

"method to obtain cluster object.",

AirflowProviderDeprecationWarning,

stacklevel=2,

)

# Remove result of apply defaults

if "params" in kwargs:

del kwargs["params"]

# Create cluster object from kwargs

if project_id is None:

raise AirflowException(

"project_id argument is required when building cluster from keywords parameters"

)

kwargs["project_id"] = project_id

cluster_config = ClusterGenerator(**kwargs).make()

# Remove from kwargs cluster params passed for backward compatibility

cluster_params = inspect.signature(ClusterGenerator.__init__).parameters

for arg in cluster_params:

if arg in kwargs:

del kwargs[arg]

super().__init__(**kwargs)

if deferrable and polling_interval_seconds <= 0:

raise ValueError("Invalid value for polling_interval_seconds. Expected value greater than 0")

self.cluster_config = cluster_config

self.cluster_name = cluster_name

self.labels = labels

self.project_id = project_id

self.region = region

self.request_id = request_id

self.retry = retry

self.timeout = timeout

self.metadata = metadata

self.gcp_conn_id = gcp_conn_id

self.delete_on_error = delete_on_error

self.use_if_exists = use_if_exists

self.impersonation_chain = impersonation_chain

self.virtual_cluster_config = virtual_cluster_config

self.deferrable = deferrable

self.polling_interval_seconds = polling_interval_seconds

def _create_cluster(self, hook: DataprocHook):

return hook.create_cluster(

project_id=self.project_id,

region=self.region,

cluster_name=self.cluster_name,

labels=self.labels,

cluster_config=self.cluster_config,

virtual_cluster_config=self.virtual_cluster_config,

request_id=self.request_id,

retry=self.retry,

timeout=self.timeout,

metadata=self.metadata,

)

def _delete_cluster(self, hook):

self.log.info("Deleting the cluster")

hook.delete_cluster(region=self.region, cluster_name=self.cluster_name, project_id=self.project_id)

def _get_cluster(self, hook: DataprocHook) -> Cluster:

return hook.get_cluster(

project_id=self.project_id,

region=self.region,

cluster_name=self.cluster_name,

retry=self.retry,

timeout=self.timeout,

metadata=self.metadata,

)

def _handle_error_state(self, hook: DataprocHook, cluster: Cluster) -> None:

if cluster.status.state != cluster.status.State.ERROR:

return

self.log.info("Cluster is in ERROR state")

self.log.info("Gathering diagnostic information.")

operation = hook.diagnose_cluster(

region=self.region, cluster_name=self.cluster_name, project_id=self.project_id

)

operation.result()

gcs_uri = str(operation.operation.response.value)

self.log.info("Diagnostic information for cluster %s available at: %s", self.cluster_name, gcs_uri)

if self.delete_on_error:

self._delete_cluster(hook)

# The delete op is asynchronous and can cause further failure if the cluster finishes

# deleting between catching AlreadyExists and checking state

self._wait_for_cluster_in_deleting_state(hook)

raise AirflowException("Cluster was created in an ERROR state then deleted.")

raise AirflowException("Cluster was created but is in ERROR state")

def _wait_for_cluster_in_deleting_state(self, hook: DataprocHook) -> None:

time_left = self.timeout

for time_to_sleep in exponential_sleep_generator(initial=10, maximum=120):

if time_left < 0:

raise AirflowException(f"Cluster {self.cluster_name} is still DELETING state, aborting")

time.sleep(time_to_sleep)

time_left -= time_to_sleep

try:

self._get_cluster(hook)

except NotFound:

break

def _wait_for_cluster_in_creating_state(self, hook: DataprocHook) -> Cluster:

time_left = self.timeout

cluster = self._get_cluster(hook)

for time_to_sleep in exponential_sleep_generator(initial=10, maximum=120):

if cluster.status.state != cluster.status.State.CREATING:

break

if time_left < 0:

raise AirflowException(f"Cluster {self.cluster_name} is still CREATING state, aborting")

time.sleep(time_to_sleep)

time_left -= time_to_sleep

cluster = self._get_cluster(hook)

return cluster

def _start_cluster(self, hook: DataprocHook):

op: operation.Operation = hook.start_cluster(

region=self.region,

project_id=self.project_id,

cluster_name=self.cluster_name,

retry=self.retry,

timeout=self.timeout,

metadata=self.metadata,

)

return hook.wait_for_operation(timeout=self.timeout, result_retry=self.retry, operation=op)

def execute(self, context: Context) -> dict:

self.log.info("Creating cluster: %s", self.cluster_name)

hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)

# Save data required to display extra link no matter what the cluster status will be

project_id = self.project_id or hook.project_id

if project_id:

DataprocClusterLink.persist(

context=context,

operator=self,

cluster_id=self.cluster_name,

project_id=project_id,

region=self.region,

)

try:

# First try to create a new cluster

operation = self._create_cluster(hook)

if not self.deferrable:

cluster = hook.wait_for_operation(

timeout=self.timeout, result_retry=self.retry, operation=operation

)

self.log.info("Cluster created.")

return Cluster.to_dict(cluster)

else:

cluster = hook.get_cluster(

project_id=self.project_id, region=self.region, cluster_name=self.cluster_name

)

if cluster.status.state == cluster.status.State.RUNNING:

self.log.info("Cluster created.")

return Cluster.to_dict(cluster)

else:

self.defer(

trigger=DataprocClusterTrigger(

cluster_name=self.cluster_name,

project_id=self.project_id,

region=self.region,

gcp_conn_id=self.gcp_conn_id,

impersonation_chain=self.impersonation_chain,

polling_interval_seconds=self.polling_interval_seconds,

delete_on_error=self.delete_on_error,

),

method_name="execute_complete",

)

except AlreadyExists:

if not self.use_if_exists:

raise

self.log.info("Cluster already exists.")

cluster = self._get_cluster(hook)

except AirflowException as ae:

# There still could be a cluster created here in an ERROR state which

# should be deleted immediately rather than consuming another retry attempt

# (assuming delete_on_error is true (default))

# This reduces overall the number of task attempts from 3 to 2 to successful cluster creation

# assuming the underlying GCE issues have resolved within that window. Users can configure

# a higher number of retry attempts in powers of two with 30s-60s wait interval

try:

cluster = self._get_cluster(hook)

self._handle_error_state(hook, cluster)

except AirflowException as ae_inner:

# We could get any number of failures here, including cluster not found and we

# can just ignore to ensure we surface the original cluster create failure

self.log.exception(ae_inner)

finally:

raise ae

# Check if cluster is not in ERROR state

self._handle_error_state(hook, cluster)

if cluster.status.state == cluster.status.State.CREATING:

# Wait for cluster to be created

cluster = self._wait_for_cluster_in_creating_state(hook)

self._handle_error_state(hook, cluster)

elif cluster.status.state == cluster.status.State.DELETING:

# Wait for cluster to be deleted

self._wait_for_cluster_in_deleting_state(hook)

# Create new cluster

cluster = self._create_cluster(hook)

self._handle_error_state(hook, cluster)

elif cluster.status.state == cluster.status.State.STOPPED:

# if the cluster exists and already stopped, then start the cluster

self._start_cluster(hook)

return Cluster.to_dict(cluster)

def execute_complete(self, context: Context, event: dict[str, Any]) -> Any:

"""

Act as a callback for when the trigger fires - returns immediately.

Relies on trigger to throw an exception, otherwise it assumes execution was successful.

"""

cluster_state = event["cluster_state"]

cluster_name = event["cluster_name"]

if cluster_state == ClusterStatus.State.ERROR:

raise AirflowException(f"Cluster is in ERROR state:\n{cluster_name}")

self.log.info("%s completed successfully.", self.task_id)

"""Scale, up or down, a cluster on Google Cloud Dataproc.

The operator will wait until the cluster is re-scaled.

Example usage:

.. code-block:: python

t1 = DataprocClusterScaleOperator(

task_id="dataproc_scale",

project_id="my-project",

cluster_name="cluster-1",

num_workers=10,

num_preemptible_workers=10,

graceful_decommission_timeout="1h",

)

.. seealso::

For more detail on about scaling clusters have a look at the reference:

https://cloud.google.com/dataproc/docs/concepts/configuring-clusters/scaling-clusters

:param cluster_name: The name of the cluster to scale. (templated)

:param project_id: The ID of the google cloud project in which

the cluster runs. (templated)

:param region: The region for the dataproc cluster. (templated)

:param num_workers: The new number of workers

:param num_preemptible_workers: The new number of preemptible workers

:param graceful_decommission_timeout: Timeout for graceful YARN decommissioning.

Maximum value is 1d

:param gcp_conn_id: The connection ID to use connecting to Google Cloud.

:param impersonation_chain: Optional service account to impersonate using short-term

credentials, or chained list of accounts required to get the access_token

of the last account in the list, which will be impersonated in the request.

If set as a string, the account must grant the originating account

the Service Account Token Creator IAM role.

If set as a sequence, the identities from the list must grant

Service Account Token Creator IAM role to the directly preceding identity, with first

account from the list granting this role to the originating account (templated).

"""

template_fields: Sequence[str] = ("cluster_name", "project_id", "region", "impersonation_chain")

operator_extra_links = (DataprocLink(),)

def __init__(

self,

*,

cluster_name: str,

project_id: str = PROVIDE_PROJECT_ID,

region: str = "global",

num_workers: int = 2,

num_preemptible_workers: int = 0,

graceful_decommission_timeout: str | None = None,

gcp_conn_id: str = "google_cloud_default",

impersonation_chain: str | Sequence[str] | None = None,

**kwargs,

) -> None:

super().__init__(**kwargs)

self.project_id = project_id

self.region = region

self.cluster_name = cluster_name

self.num_workers = num_workers

self.num_preemptible_workers = num_preemptible_workers

self.graceful_decommission_timeout = graceful_decommission_timeout

self.gcp_conn_id = gcp_conn_id

self.impersonation_chain = impersonation_chain

def _build_scale_cluster_data(self) -> dict:

scale_data = {

"config": {

"worker_config": {"num_instances": self.num_workers},

"secondary_worker_config": {"num_instances": self.num_preemptible_workers},

}

}

return scale_data

@property

def _graceful_decommission_timeout_object(self) -> dict[str, int] | None:

if not self.graceful_decommission_timeout:

return None

timeout = None

match = re.fullmatch(r"(\d+)([smdh])", self.graceful_decommission_timeout)

if match:

val = int(match.group(1))

unit = match.group(2)

if unit == "s":

timeout = val

elif unit == "m":

timeout = int(timedelta(minutes=val).total_seconds())

elif unit == "h":

timeout = int(timedelta(hours=val).total_seconds())

elif unit == "d":

timeout = int(timedelta(days=val).total_seconds())

if not timeout:

raise AirflowException(

"DataprocClusterScaleOperator "

" should be expressed in day, hours, minutes or seconds. "

" i.e. 1d, 4h, 10m, 30s"

)

return {"seconds": timeout}

def execute(self, context: Context) -> None:

"""Scale, up or down, a cluster on Google Cloud Dataproc."""

self.log.info("Scaling cluster: %s", self.cluster_name)

scaling_cluster_data = self._build_scale_cluster_data()

update_mask = ["config.worker_config.num_instances", "config.secondary_worker_config.num_instances"]

hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)

# Save data required to display extra link no matter what the cluster status will be

DataprocLink.persist(

context=context,

task_instance=self,