Colin-Cheng commented on code in PR #74753:
URL: https://github.com/apache/airflow/pull/74753#discussion_r4288563339
##########
providers/google/docs/operators/cloud/dataproc.rst:
##########
@@ -186,6 +186,43 @@ You can use deferrable mode for this action in order to
run the operator asynchr
:start-after: [START how_to_cloud_dataproc_create_cluster_operator_async]
:end-before: [END how_to_cloud_dataproc_create_cluster_operator_async]
+High Obtainability (HO) mode clusters
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+High Obtainability (HO) mode enables creating Dataproc clusters using abstract
hardware requirements
+(:class:`~airflow.providers.google.cloud.operators.dataproc.FlexibleInstanceGroupConfig`:
number of
+instances, core count, and hardware class) instead of concrete machine types.
The Dataproc backend resolves
+these requirements into a ranked list of widely available machine families to
minimize stockouts and improve
+cluster creation reliability.
+
+High Obtainability clusters have the following characteristics:
+
+* **Mandatory AutoZone**: HO clusters require AutoZone to dynamically select
an optimal zone; do not specify a ``zone``.
+* **Cluster-wide flexible config**: Master and primary workers must both use
flexible configurations (or
+ ``num_workers=0`` for a single-node cluster).
Review Comment:
`num_workers=0` was a `ClusterGenerator` parameter. With a `cluster_config`
dictionary, a single-node HO cluster is configured by omitting
`worker_flexible_config` and setting `dataproc:dataproc.allow.zero.workers` to
`"true"`:
```suggestion
* **Cluster-wide flexible config**: Master and primary workers must both use
flexible configurations (or omit
``worker_flexible_config`` and set
``dataproc:dataproc.allow.zero.workers`` to ``"true"`` in
``software_config.properties`` for a single-node cluster).
```
##########
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py:
##########
@@ -117,6 +118,206 @@ class InstanceFlexibilityPolicy:
instance_selection_list: list[InstanceSelection]
+class HardwareClass(Enum):
+ """
+ Hardware class of a High Obtainability (HO) mode instance group.
+
+ Representation for
+
google.cloud.dataproc.v1#google.cloud.dataproc.v1.FlexibleInstanceGroupConfig.HardwareClass.
+ """
+
+ STANDARD = "STANDARD"
+ PERFORMANCE = "PERFORMANCE"
+
+
+FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS = (2, 4, 8, 16, 32)
+
+
+@dataclass
+class FlexibleInstanceGroupConfig:
+ """
+ Abstract (High Obtainability mode) configuration of a node group.
+
+ Representation for
+
google.cloud.dataproc.v1#google.cloud.dataproc.v1.FlexibleInstanceGroupConfig.
+
+ When a cluster is created with flexible instance group configs, Dataproc
+ High Obtainability (HO) mode is enabled: instead of concrete machine types
+ the caller specifies abstract requirements (number of cores and hardware
+ class) and the Dataproc backend resolves them into a ranked list of widely
+ available machine families to maximize obtainability.
+
+ ``num_instances`` is mutable (via ``DataprocUpdateClusterOperator`` and
+ autoscaling). ``core_count``, ``hardware_class`` and ``boot_disk_size_gb``
+ define the initial cluster shape and are create-time only.
+
+ :param num_instances: The number of VM instances in the node group.
+ :param core_count: Number of cores per instance. Must be one of
+ ``2, 4, 8, 16, 32``.
+ :param hardware_class: Hardware class of the instances: ``STANDARD`` or
+ ``PERFORMANCE`` (a string or a :class:`HardwareClass`).
+ :param boot_disk_size_gb: Optional. Size of the boot disk in GB.
+ """
+
+ num_instances: int
+ core_count: int
+ hardware_class: str | HardwareClass = HardwareClass.STANDARD
+ boot_disk_size_gb: int | None = None
+
+ def __post_init__(self):
+ if isinstance(self.hardware_class, str):
+ try:
+ self.hardware_class =
HardwareClass(self.hardware_class.upper())
+ except ValueError as e:
+ raise ValueError(
+ f"Invalid hardware_class: {self.hardware_class!r}. "
+ f"Must be one of {[c.value for c in HardwareClass]}."
+ ) from e
+ elif not isinstance(self.hardware_class, HardwareClass):
+ raise ValueError(
+ f"Invalid hardware_class: {self.hardware_class!r}. Must be a
str or HardwareClass."
+ )
+ if (
+ not isinstance(self.num_instances, int)
+ or isinstance(self.num_instances, bool)
+ or self.num_instances < 0
+ ):
+ raise ValueError(
+ "num_instances of a flexible config must be a non-negative
integer,"
+ f" got {self.num_instances!r}."
+ )
+ if (
+ not isinstance(self.core_count, int)
+ or isinstance(self.core_count, bool)
+ or self.core_count not in FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS
+ ):
+ raise ValueError(
+ "core_count of a flexible config must be one of"
+ f" {list(FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS)}, got"
+ f" {self.core_count!r}."
+ )
+ if self.boot_disk_size_gb is not None and (
+ not isinstance(self.boot_disk_size_gb, int)
+ or isinstance(self.boot_disk_size_gb, bool)
+ or self.boot_disk_size_gb <= 0
+ ):
+ raise ValueError(
+ "boot_disk_size_gb of a flexible config must be a positive
integer,"
+ f" got {self.boot_disk_size_gb!r}."
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ """Return the ``FlexibleInstanceGroupConfig`` as an API-shaped dict."""
+ data: dict[str, Any] = {
+ "num_instances": self.num_instances,
+ "core_count": self.core_count,
+ "hardware_class": HardwareClass(self.hardware_class).value,
+ }
+ if self.boot_disk_size_gb is not None:
+ data["boot_disk_size_gb"] = self.boot_disk_size_gb
+ return data
+
+
+def _to_flexible_config(
+ value: FlexibleInstanceGroupConfig | dict[str, Any] | None, name: str
+) -> FlexibleInstanceGroupConfig | None:
+ """Coerces a user supplied flexible config (dataclass or dict) to a
dataclass."""
+ if value is None:
+ return None
Review Comment:
`_to_flexible_config` is annotated as returning `FlexibleInstanceGroupConfig
| None`, which causes `mypy` `[union-attr]` errors in
`_validate_and_normalize_flexible_cluster_config` on `master_fc.to_dict()`,
`master_fc.num_instances`, `worker_fc.to_dict()`, and
`secondary_worker_fc.to_dict()`.
Since `_to_flexible_config` is only called for keys in `ho_keys` (which
already filters out `None`), we can tighten its signature so `mypy` knows it
always returns `FlexibleInstanceGroupConfig`:
```suggestion
def _to_flexible_config(
value: FlexibleInstanceGroupConfig | dict[str, Any], name: str
) -> FlexibleInstanceGroupConfig:
"""Coerces a user supplied flexible config (dataclass or dict) to a
dataclass."""
```
##########
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py:
##########
@@ -117,6 +118,206 @@ class InstanceFlexibilityPolicy:
instance_selection_list: list[InstanceSelection]
+class HardwareClass(Enum):
+ """
+ Hardware class of a High Obtainability (HO) mode instance group.
+
+ Representation for
+
google.cloud.dataproc.v1#google.cloud.dataproc.v1.FlexibleInstanceGroupConfig.HardwareClass.
+ """
+
+ STANDARD = "STANDARD"
+ PERFORMANCE = "PERFORMANCE"
+
+
+FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS = (2, 4, 8, 16, 32)
+
+
+@dataclass
+class FlexibleInstanceGroupConfig:
+ """
+ Abstract (High Obtainability mode) configuration of a node group.
+
+ Representation for
+
google.cloud.dataproc.v1#google.cloud.dataproc.v1.FlexibleInstanceGroupConfig.
+
+ When a cluster is created with flexible instance group configs, Dataproc
+ High Obtainability (HO) mode is enabled: instead of concrete machine types
+ the caller specifies abstract requirements (number of cores and hardware
+ class) and the Dataproc backend resolves them into a ranked list of widely
+ available machine families to maximize obtainability.
+
+ ``num_instances`` is mutable (via ``DataprocUpdateClusterOperator`` and
+ autoscaling). ``core_count``, ``hardware_class`` and ``boot_disk_size_gb``
+ define the initial cluster shape and are create-time only.
+
+ :param num_instances: The number of VM instances in the node group.
+ :param core_count: Number of cores per instance. Must be one of
+ ``2, 4, 8, 16, 32``.
+ :param hardware_class: Hardware class of the instances: ``STANDARD`` or
+ ``PERFORMANCE`` (a string or a :class:`HardwareClass`).
+ :param boot_disk_size_gb: Optional. Size of the boot disk in GB.
+ """
+
+ num_instances: int
+ core_count: int
+ hardware_class: str | HardwareClass = HardwareClass.STANDARD
+ boot_disk_size_gb: int | None = None
+
+ def __post_init__(self):
+ if isinstance(self.hardware_class, str):
+ try:
+ self.hardware_class =
HardwareClass(self.hardware_class.upper())
+ except ValueError as e:
+ raise ValueError(
+ f"Invalid hardware_class: {self.hardware_class!r}. "
+ f"Must be one of {[c.value for c in HardwareClass]}."
+ ) from e
+ elif not isinstance(self.hardware_class, HardwareClass):
+ raise ValueError(
+ f"Invalid hardware_class: {self.hardware_class!r}. Must be a
str or HardwareClass."
+ )
+ if (
+ not isinstance(self.num_instances, int)
+ or isinstance(self.num_instances, bool)
+ or self.num_instances < 0
+ ):
+ raise ValueError(
+ "num_instances of a flexible config must be a non-negative
integer,"
+ f" got {self.num_instances!r}."
+ )
+ if (
+ not isinstance(self.core_count, int)
+ or isinstance(self.core_count, bool)
+ or self.core_count not in FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS
+ ):
+ raise ValueError(
+ "core_count of a flexible config must be one of"
+ f" {list(FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS)}, got"
+ f" {self.core_count!r}."
+ )
+ if self.boot_disk_size_gb is not None and (
+ not isinstance(self.boot_disk_size_gb, int)
+ or isinstance(self.boot_disk_size_gb, bool)
+ or self.boot_disk_size_gb <= 0
+ ):
+ raise ValueError(
+ "boot_disk_size_gb of a flexible config must be a positive
integer,"
+ f" got {self.boot_disk_size_gb!r}."
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ """Return the ``FlexibleInstanceGroupConfig`` as an API-shaped dict."""
+ data: dict[str, Any] = {
+ "num_instances": self.num_instances,
+ "core_count": self.core_count,
+ "hardware_class": HardwareClass(self.hardware_class).value,
+ }
+ if self.boot_disk_size_gb is not None:
+ data["boot_disk_size_gb"] = self.boot_disk_size_gb
+ return data
+
+
+def _to_flexible_config(
+ value: FlexibleInstanceGroupConfig | dict[str, Any] | None, name: str
+) -> FlexibleInstanceGroupConfig | None:
+ """Coerces a user supplied flexible config (dataclass or dict) to a
dataclass."""
+ if value is None:
+ return None
+ if isinstance(value, FlexibleInstanceGroupConfig):
+ return value
+ if isinstance(value, dict):
+ try:
+ return FlexibleInstanceGroupConfig(**value)
+ except TypeError as e:
+ raise ValueError(
+ f"Invalid {name}: {e}. Expected keys: num_instances,
core_count,"
+ " hardware_class, boot_disk_size_gb."
+ ) from e
+ raise ValueError(
+ f"{name} must be a FlexibleInstanceGroupConfig or a dict, got
{type(value).__name__}."
+ )
+
+
+def flexible_scale_cluster_update(
+ num_workers: int | None = None,
+ num_secondary_workers: int | None = None,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+ """
+ Build ``cluster`` and ``update_mask`` to scale a High Obtainability
cluster.
+
+ Returns the arguments to pass to ``DataprocUpdateClusterOperator`` so that
+ the resize is expressed through the authoritative HO path
+ (``config.worker_flexible_config.num_instances`` /
+ ``config.secondary_worker_flexible_config.num_instances``). Only the
+ flexible update-mask paths are used; the legacy
+ ``config.worker_config.num_instances`` path is never combined with them so
+ that the server never sees conflicting instance counts.
+
+ Example usage:
+
+ .. code-block:: python
+
+ cluster, update_mask = flexible_scale_cluster_update(num_workers=5)
+ scale = DataprocUpdateClusterOperator(
+ task_id="scale_ho_cluster",
+ project_id="my-project",
+ region="us-central1",
+ cluster_name="my-ho-cluster",
+ cluster=cluster,
+ update_mask=update_mask,
+ graceful_decommission_timeout={"seconds": 600},
+ )
+
+ :param num_workers: The new number of primary workers.
+ :param num_secondary_workers: The new number of secondary workers.
+ :return: A ``(cluster, update_mask)`` tuple.
+ """
+ if num_workers is None and num_secondary_workers is None:
+ raise ValueError("At least one of num_workers or num_secondary_workers
must be set.")
+ config: dict[str, Any] = {}
+ paths: list[str] = []
+ if num_workers is not None:
+ if not isinstance(num_workers, int) or isinstance(num_workers, bool)
or num_workers < 0:
+ raise ValueError("num_workers must be a non-negative integer.")
+ config["worker_flexible_config"] = {"num_instances": num_workers}
+ paths.append("config.worker_flexible_config.num_instances")
+ if num_secondary_workers is not None:
+ if (
+ not isinstance(num_secondary_workers, int)
+ or isinstance(num_secondary_workers, bool)
+ or num_secondary_workers < 0
+ ):
+ raise ValueError("num_secondary_workers must be a non-negative
integer.")
+ config["secondary_worker_flexible_config"] = {"num_instances":
num_secondary_workers}
+ paths.append("config.secondary_worker_flexible_config.num_instances")
+ return {"config": config}, {"paths": paths}
+
+
+def _get_num_instances(cluster: Any, path: str) -> int | None:
+ """
+ Read the ``num_instances`` at ``path`` from a dict or ``Cluster`` message.
+
+ Returns ``None`` when the path is not present (e.g. the field is missing in
+ the dict, or the installed client library does not know it yet) or when
+ ``cluster`` is a templated string.
+ """
+ value: Any = cluster
+ for key in path.split("."):
+ if isinstance(value, dict):
+ value = value.get(key)
+ elif isinstance(value, str):
+ return None
+ else:
+ try:
+ value = getattr(value, key)
+ except (AttributeError, KeyError):
+ return None
+ if value is None:
+ return None
+ return int(value) if isinstance(value, int) else None
Review Comment:
A few edge cases here and in `_validate_update_mask_paths`:
1. `isinstance(value, int)` is `True` for `bool`, so line 318 should also
check `not isinstance(value, bool)`.
2. For `proto-plus` `Cluster` messages, `getattr(value, key)` on an unset
submessage returns a default empty submessage with `num_instances == 0` rather
than `None`. Checking `key not in value` on proto messages avoids reading a
proto3 default `0` from an unset group.
3. In `_validate_update_mask_paths` (around line 2647), if `cluster` is a
templated string (`isinstance(cluster, str)`), `_get_num_instances` returns
`None` for both paths and `if legacy_count is None ...` will raise `ValueError`
at parse time instead of skipping validation as documented.
```suggestion
value: Any = cluster
for key in path.split("."):
if isinstance(value, dict):
value = value.get(key)
elif isinstance(value, str):
return None
else:
if hasattr(type(value), "pb") and key not in value:
return None
try:
value = getattr(value, key)
except (AttributeError, KeyError):
return None
if value is None:
return None
return int(value) if isinstance(value, int) and not isinstance(value,
bool) else None
```
##########
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py:
##########
@@ -117,6 +118,206 @@ class InstanceFlexibilityPolicy:
instance_selection_list: list[InstanceSelection]
+class HardwareClass(Enum):
+ """
+ Hardware class of a High Obtainability (HO) mode instance group.
+
+ Representation for
+
google.cloud.dataproc.v1#google.cloud.dataproc.v1.FlexibleInstanceGroupConfig.HardwareClass.
+ """
+
+ STANDARD = "STANDARD"
+ PERFORMANCE = "PERFORMANCE"
+
+
+FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS = (2, 4, 8, 16, 32)
+
+
+@dataclass
+class FlexibleInstanceGroupConfig:
+ """
+ Abstract (High Obtainability mode) configuration of a node group.
+
+ Representation for
+
google.cloud.dataproc.v1#google.cloud.dataproc.v1.FlexibleInstanceGroupConfig.
+
+ When a cluster is created with flexible instance group configs, Dataproc
+ High Obtainability (HO) mode is enabled: instead of concrete machine types
+ the caller specifies abstract requirements (number of cores and hardware
+ class) and the Dataproc backend resolves them into a ranked list of widely
+ available machine families to maximize obtainability.
+
+ ``num_instances`` is mutable (via ``DataprocUpdateClusterOperator`` and
+ autoscaling). ``core_count``, ``hardware_class`` and ``boot_disk_size_gb``
+ define the initial cluster shape and are create-time only.
+
+ :param num_instances: The number of VM instances in the node group.
+ :param core_count: Number of cores per instance. Must be one of
+ ``2, 4, 8, 16, 32``.
+ :param hardware_class: Hardware class of the instances: ``STANDARD`` or
+ ``PERFORMANCE`` (a string or a :class:`HardwareClass`).
+ :param boot_disk_size_gb: Optional. Size of the boot disk in GB.
+ """
+
+ num_instances: int
+ core_count: int
+ hardware_class: str | HardwareClass = HardwareClass.STANDARD
+ boot_disk_size_gb: int | None = None
+
+ def __post_init__(self):
+ if isinstance(self.hardware_class, str):
+ try:
+ self.hardware_class =
HardwareClass(self.hardware_class.upper())
+ except ValueError as e:
+ raise ValueError(
+ f"Invalid hardware_class: {self.hardware_class!r}. "
+ f"Must be one of {[c.value for c in HardwareClass]}."
+ ) from e
+ elif not isinstance(self.hardware_class, HardwareClass):
+ raise ValueError(
+ f"Invalid hardware_class: {self.hardware_class!r}. Must be a
str or HardwareClass."
+ )
+ if (
+ not isinstance(self.num_instances, int)
+ or isinstance(self.num_instances, bool)
+ or self.num_instances < 0
+ ):
+ raise ValueError(
+ "num_instances of a flexible config must be a non-negative
integer,"
+ f" got {self.num_instances!r}."
+ )
+ if (
+ not isinstance(self.core_count, int)
+ or isinstance(self.core_count, bool)
+ or self.core_count not in FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS
+ ):
+ raise ValueError(
+ "core_count of a flexible config must be one of"
+ f" {list(FLEXIBLE_CONFIG_ALLOWED_CORE_COUNTS)}, got"
+ f" {self.core_count!r}."
+ )
+ if self.boot_disk_size_gb is not None and (
+ not isinstance(self.boot_disk_size_gb, int)
+ or isinstance(self.boot_disk_size_gb, bool)
+ or self.boot_disk_size_gb <= 0
+ ):
+ raise ValueError(
+ "boot_disk_size_gb of a flexible config must be a positive
integer,"
+ f" got {self.boot_disk_size_gb!r}."
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ """Return the ``FlexibleInstanceGroupConfig`` as an API-shaped dict."""
+ data: dict[str, Any] = {
+ "num_instances": self.num_instances,
+ "core_count": self.core_count,
+ "hardware_class": HardwareClass(self.hardware_class).value,
+ }
+ if self.boot_disk_size_gb is not None:
+ data["boot_disk_size_gb"] = self.boot_disk_size_gb
+ return data
+
+
+def _to_flexible_config(
+ value: FlexibleInstanceGroupConfig | dict[str, Any] | None, name: str
+) -> FlexibleInstanceGroupConfig | None:
+ """Coerces a user supplied flexible config (dataclass or dict) to a
dataclass."""
+ if value is None:
+ return None
+ if isinstance(value, FlexibleInstanceGroupConfig):
+ return value
+ if isinstance(value, dict):
+ try:
+ return FlexibleInstanceGroupConfig(**value)
+ except TypeError as e:
+ raise ValueError(
+ f"Invalid {name}: {e}. Expected keys: num_instances,
core_count,"
+ " hardware_class, boot_disk_size_gb."
+ ) from e
+ raise ValueError(
+ f"{name} must be a FlexibleInstanceGroupConfig or a dict, got
{type(value).__name__}."
+ )
+
+
+def flexible_scale_cluster_update(
+ num_workers: int | None = None,
+ num_secondary_workers: int | None = None,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+ """
+ Build ``cluster`` and ``update_mask`` to scale a High Obtainability
cluster.
+
+ Returns the arguments to pass to ``DataprocUpdateClusterOperator`` so that
+ the resize is expressed through the authoritative HO path
+ (``config.worker_flexible_config.num_instances`` /
+ ``config.secondary_worker_flexible_config.num_instances``). Only the
+ flexible update-mask paths are used; the legacy
+ ``config.worker_config.num_instances`` path is never combined with them so
+ that the server never sees conflicting instance counts.
+
+ Example usage:
+
+ .. code-block:: python
+
+ cluster, update_mask = flexible_scale_cluster_update(num_workers=5)
+ scale = DataprocUpdateClusterOperator(
+ task_id="scale_ho_cluster",
+ project_id="my-project",
+ region="us-central1",
+ cluster_name="my-ho-cluster",
+ cluster=cluster,
+ update_mask=update_mask,
+ graceful_decommission_timeout={"seconds": 600},
+ )
+
+ :param num_workers: The new number of primary workers.
+ :param num_secondary_workers: The new number of secondary workers.
+ :return: A ``(cluster, update_mask)`` tuple.
+ """
+ if num_workers is None and num_secondary_workers is None:
+ raise ValueError("At least one of num_workers or num_secondary_workers
must be set.")
+ config: dict[str, Any] = {}
+ paths: list[str] = []
+ if num_workers is not None:
+ if not isinstance(num_workers, int) or isinstance(num_workers, bool)
or num_workers < 0:
+ raise ValueError("num_workers must be a non-negative integer.")
Review Comment:
Primary workers cannot be scaled to `0` (consistent with
`worker_flexible_config.num_instances` requiring a positive integer in
`_validate_and_normalize_flexible_cluster_config`), whereas
`num_secondary_workers` can be `0`:
```suggestion
if not isinstance(num_workers, int) or isinstance(num_workers, bool)
or num_workers <= 0:
raise ValueError("num_workers must be a positive integer.")
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
