jigneshjain25 commented on code in PR #74753:
URL: https://github.com/apache/airflow/pull/74753#discussion_r4289093410
##########
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:
"Good catch! Tightened _to_flexible_config to accept
FlexibleInstanceGroupConfig | dict[str, Any] and return
FlexibleInstanceGroupConfig, eliminating the mypy union-attr warning."
--
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]
