From 88c565e11a56ec5cfeb62238f3d162cea10dae86 Mon Sep 17 00:00:00 2001 From: viiccwen Date: Wed, 15 Jul 2026 20:03:57 +0800 Subject: [PATCH] Fix deprecated node fields in GKE system Dags Cluster-level node settings emit provider deprecation warnings, while Autopilot manages its own nodes and cannot use user-managed node pools. Signed-off-by: viiccwen --- .../google/cloud/operators/kubernetes_engine.py | 10 +++++++++- .../example_kubernetes_engine.py | 2 +- .../example_kubernetes_engine_async.py | 5 ++++- .../example_kubernetes_engine_job.py | 5 ++++- .../example_kubernetes_engine_kueue.py | 2 +- .../example_kubernetes_engine_resource.py | 5 ++++- .../cloud/operators/test_kubernetes_engine.py | 15 ++++++++++++++- 7 files changed, 37 insertions(+), 7 deletions(-) diff --git a/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py b/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py index 49c80445af933..3b09e984eb066 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -447,7 +447,9 @@ def _validate_input(self) -> None: "body['node_config'] or body['initial_code_count'] are specified." ) - if not any([self._body_field("node_config"), self._body_field("initial_node_count")]): + if not self._is_autopilot_cluster() and not any( + [self._body_field("node_config"), self._body_field("initial_node_count")] + ): if not self._body_field("node_pools"): error_messages.append( "Field body['node_pools'] is required if none of fields " @@ -466,6 +468,12 @@ def _body_field(self, field_name: str, default_value: Any = None) -> Any: return self.body.get(field_name, default_value) return getattr(self.body, field_name, default_value) + def _is_autopilot_cluster(self) -> bool: + autopilot = self._body_field("autopilot") + if isinstance(autopilot, dict): + return bool(autopilot.get("enabled")) + return bool(getattr(autopilot, "enabled", False)) + def _alert_deprecated_body_fields(self) -> None: """Generate warning messages if deprecated fields were used in the body.""" deprecated_body_fields_with_replacement = [ diff --git a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine.py b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine.py index 54de4a904bac0..051bcaa310a0b 100644 --- a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine.py +++ b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine.py @@ -51,7 +51,7 @@ CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL # [START howto_operator_gcp_gke_create_cluster_definition] -CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1, "autopilot": {"enabled": True}} +CLUSTER = {"name": CLUSTER_NAME, "autopilot": {"enabled": True}} # [END howto_operator_gcp_gke_create_cluster_definition] diff --git a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py index b82ecd97fc08f..12f61ea5fef82 100644 --- a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py +++ b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_async.py @@ -51,7 +51,10 @@ CLUSTER_NAME_FULL = CLUSTER_NAME_BASE + f"-{ENV_ID}".replace("_", "-") CLUSTER_NAME = CLUSTER_NAME_BASE if len(CLUSTER_NAME_FULL) >= 33 else CLUSTER_NAME_FULL -CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1} +CLUSTER = { + "name": CLUSTER_NAME, + "node_pools": [{"name": "default-pool", "initial_node_count": 1}], +} with DAG( DAG_ID, diff --git a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_job.py b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_job.py index cfb56b5da6e83..112fba3242892 100644 --- a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_job.py +++ b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_job.py @@ -49,7 +49,10 @@ GCP_LOCATION = "europe-north1-a" CLUSTER_NAME = f"gke-job-{ENV_ID}".replace("_", "-") -CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1} +CLUSTER = { + "name": CLUSTER_NAME, + "node_pools": [{"name": "default-pool", "initial_node_count": 1}], +} JOB_NAME = "test-pi" JOB_NAME_DEF = "test-pi-def" diff --git a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_kueue.py b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_kueue.py index c48ae19cea3c6..60127bbd9ab3e 100644 --- a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_kueue.py +++ b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_kueue.py @@ -47,7 +47,7 @@ GCP_LOCATION = "europe-west3" CLUSTER_NAME = f"gke-kueue-{ENV_ID}".replace("_", "-") -CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1, "autopilot": {"enabled": True}} +CLUSTER = {"name": CLUSTER_NAME, "autopilot": {"enabled": True}} flavor_conf = """ apiVersion: kueue.x-k8s.io/v1beta1 diff --git a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_resource.py b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_resource.py index f9dae65f1cb6a..81666fc55ef36 100644 --- a/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_resource.py +++ b/providers/google/tests/system/google/cloud/kubernetes_engine/example_kubernetes_engine_resource.py @@ -44,7 +44,10 @@ GCP_LOCATION = "europe-north1-a" CLUSTER_NAME = f"gke-resource-{ENV_ID}".replace("_", "-") -CLUSTER = {"name": CLUSTER_NAME, "initial_node_count": 1} +CLUSTER = { + "name": CLUSTER_NAME, + "node_pools": [{"name": "default-pool", "initial_node_count": 1}], +} PVC_NAME = "test-pvc-name" PVC_CONF = f""" diff --git a/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py b/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py index 4c92ab65a7943..377b1a146b7ba 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py +++ b/providers/google/tests/unit/google/cloud/operators/test_kubernetes_engine.py @@ -102,6 +102,11 @@ GKE_CLUSTER_CREATE_BODY_OBJECT = Cluster( name=GKE_CLUSTER_NAME, node_pools=[NodePool(name="a_node_pool", initial_node_count=1)] ) +GKE_AUTOPILOT_CLUSTER_CREATE_BODY_DICT = { + "name": GKE_CLUSTER_NAME, + "autopilot": {"enabled": True}, +} +GKE_AUTOPILOT_CLUSTER_CREATE_BODY_OBJECT = Cluster(name=GKE_CLUSTER_NAME, autopilot={"enabled": True}) GKE_CLUSTER_CREATE_BODY_DICT_DEPRECATED = {"name": GKE_CLUSTER_NAME, "initial_node_count": 1} GKE_CLUSTER_CREATE_BODY_OBJECT_DEPRECATED = Cluster(name=GKE_CLUSTER_NAME, initial_node_count=1) @@ -419,7 +424,15 @@ def test_template_fields(self): ) assert set(GKECreateClusterOperator.template_fields) == expected_template_fields - @pytest.mark.parametrize("body", [GKE_CLUSTER_CREATE_BODY_DICT, GKE_CLUSTER_CREATE_BODY_OBJECT]) + @pytest.mark.parametrize( + "body", + [ + GKE_CLUSTER_CREATE_BODY_DICT, + GKE_CLUSTER_CREATE_BODY_OBJECT, + GKE_AUTOPILOT_CLUSTER_CREATE_BODY_DICT, + GKE_AUTOPILOT_CLUSTER_CREATE_BODY_OBJECT, + ], + ) def test_body(self, body): op = GKECreateClusterOperator( task_id=TEST_TASK_ID,