jigneshjain25 commented on code in PR #74753:
URL: https://github.com/apache/airflow/pull/74753#discussion_r4289101403
##########
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:
Great points! All three have been addressed: 1. Added not isinstance(value,
bool) check so booleans are not treated as valid instance counts. 2. Added
check for proto-plus messages so unset submessages don't expose default 0
instance counts. 3. In _validate_update_mask_paths, skipped parse-time conflict
checks if cluster is a templated string or None, deferring evaluation to
runtime as documented. Added unit tests covering all these edge cases.
--
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]
