From 8eb55567352f1c566326c8f5af1d3442da6b515b Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 1/9] [cienet-private] Enable workflow checks for the primary working branches --- .github/workflows/dag-check.yml | 1 - .github/workflows/pyink-check.yml | 3 +- .github/workflows/pylint-check.yml | 3 +- .github/workflows/require-checklist.yml | 5 ++- .github/workflows/unit-test.yml | 1 - scripts/code-style.sh | 41 +++++++++++++++++++------ 6 files changed, 40 insertions(+), 14 deletions(-) diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index 96fcb76bf..ac3982d0e 100644 --- a/.github/workflows/dag-check.yml +++ b/.github/workflows/dag-check.yml @@ -3,7 +3,6 @@ name: DAG Check on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..c5cb1cd83 100644 --- a/.github/workflows/pyink-check.yml +++ b/.github/workflows/pyink-check.yml @@ -2,11 +2,12 @@ name: Formatter on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: branches: [master] + workflow_dispatch: {} + jobs: format_check: runs-on: ubuntu-latest diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..02f4cb31e 100644 --- a/.github/workflows/pylint-check.yml +++ b/.github/workflows/pylint-check.yml @@ -2,12 +2,13 @@ name: Linter on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: branches: [master] + workflow_dispatch: {} + jobs: linting_check: runs-on: ubuntu-latest diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..4da288575 100644 --- a/.github/workflows/require-checklist.yml +++ b/.github/workflows/require-checklist.yml @@ -2,10 +2,13 @@ name: Require Checklist on: pull_request: types: [opened, edited, synchronize] + + workflow_dispatch: {} + jobs: check_pr_body: runs-on: ubuntu-latest steps: - uses: mheap/require-checklist-action@v2 with: - requireChecklist: false # If this is true and there are no checklists detected, the action will fail \ No newline at end of file + requireChecklist: false # If this is true and there are no checklists detected, the action will fail diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 4f1723cb2..b771ca7e9 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -3,7 +3,6 @@ name: Unit Test on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/scripts/code-style.sh b/scripts/code-style.sh index 36cfa13e9..ea50709ac 100755 --- a/scripts/code-style.sh +++ b/scripts/code-style.sh @@ -19,14 +19,37 @@ set -e FOLDERS_TO_FORMAT=("dags" "xlml") -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff -done - -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pylint "./$folder" --fail-under=9.6 -done +HEAD_SHA="$(git rev-parse HEAD)" +BASE_BRANCH="dev" + +if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then + git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || { + echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check." + exit 0 + } +fi + +CHANGED_PY_FILES="$( + git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \ + | grep '\.py$' \ + | while read -r f; do + for folder in "${FOLDERS_TO_FORMAT[@]}"; do + if [[ "$f" == "$folder/"* ]]; then + echo "$f" + break + fi + done + done \ + | sort -u +)" + +if [[ -z "${CHANGED_PY_FILES}" ]]; then + echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}" + exit 1 +fi + +pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff + +pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123 echo "Successfully clean up all codes." From 2bc9cad3f544c53afd1f1a999208c2d7949a3bfb Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Tue, 28 Apr 2026 06:51:09 +0000 Subject: [PATCH 2/9] feat(tpu_observability): implement GKE cluster version manager to automate master and node pool upgrades --- .../gke_cluster_version_manager.py | 200 ++++++++++++++++++ .../tpu_observability/utils/node_pool_util.py | 50 ++++- 2 files changed, 248 insertions(+), 2 deletions(-) create mode 100644 dags/tpu_observability/gke_cluster_version_manager.py diff --git a/dags/tpu_observability/gke_cluster_version_manager.py b/dags/tpu_observability/gke_cluster_version_manager.py new file mode 100644 index 000000000..df36f4797 --- /dev/null +++ b/dags/tpu_observability/gke_cluster_version_manager.py @@ -0,0 +1,200 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A DAG to upgrade GKE cluster to latest available version.""" + +import datetime +import json +import logging +import re + +from airflow import models +from airflow.decorators import task +from airflow.utils.task_group import TaskGroup +from dags import composer_env +from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH +from dags.tpu_observability.utils import node_pool_util as node_pool +from dags.tpu_observability.utils import subprocess_util as subprocess + +DAG_ID = "gke_cluster_version_manager" + + +@task +def find_available_version(node_pool_info: node_pool.Info) -> str: + """Finds the latest available GKE version.""" + region = node_pool_info.region + if not region: + raise ValueError("Region not found in node_pool_info") + + command = ( + f"gcloud container get-server-config --region={region} --format='json'" + ) + logging.info("Running command: %s", command) + stdout = subprocess.run_exec(command) + + output_json = json.loads(stdout) + valid_versions = output_json.get("validMasterVersions", []) + + # Filter: ^1\.(3[2-9]|[4-9][0-9]) + pattern = re.compile(r"^1\.(3[2-9]|[4-9][0-9])") + matching_versions = [v for v in valid_versions if pattern.match(v)] + + if not matching_versions: + raise ValueError("No matching GKE versions found") + + latest_version = matching_versions[0] + logging.info("Found latest available version: %s", latest_version) + return latest_version + + +@task +def find_current_cluster_version(node_pool_info: node_pool.Info) -> dict: + """Finds the current version of the cluster.""" + cluster_name = node_pool_info.cluster_name + region = node_pool_info.region + if not cluster_name or not region: + raise ValueError("cluster_name or region not found in node_pool_info") + + command = ( + f"gcloud container clusters describe {cluster_name} " + f"--region={region} --format='json'" + ) + logging.info("Running command: %s", command) + stdout = subprocess.run_exec(command) + + output_json = json.loads(stdout) + current_master_version = output_json.get("currentMasterVersion") + current_node_version = output_json.get("currentNodeVersion") + + logging.info("Current Master Version: %s", current_master_version) + logging.info("Current Node Version: %s", current_node_version) + + return { + "currentMasterVersion": current_master_version, + "currentNodeVersion": current_node_version, + } + + +@task +def upgrade_master( + latest_version: str, current_versions: dict, node_pool_info: node_pool.Info +): + """Upgrades the master to the target version if needed.""" + current_master = current_versions.get("currentMasterVersion") + cluster_name = node_pool_info.cluster_name + region = node_pool_info.region + + if current_master != latest_version: + logging.info( + "Master version (%s) != Target (%s). Upgrading.", + current_master, + latest_version, + ) + command = ( + f"gcloud container clusters upgrade {cluster_name} --master " + f"--cluster-version={latest_version} --region={region} --quiet" + ) + logging.info("Running command: %s", command) + subprocess.run_exec(command) + else: + logging.info("Master is already at target version. Skipping.") + + +@task +def upgrade_nodes( + latest_version: str, current_versions: dict, node_pool_info: node_pool.Info +): + """Upgrades all node pools to the target version if needed.""" + current_node = current_versions.get("currentNodeVersion") + cluster_name = node_pool_info.cluster_name + region = node_pool_info.region + + if current_node != latest_version: + logging.info( + "Node version (%s) != Target (%s). Upgrading.", + current_node, + latest_version, + ) + command = ( + f"gcloud container clusters upgrade {cluster_name} " + f"--region={region} --cluster-version={latest_version} --quiet" + ) + logging.info("Running command: %s", command) + subprocess.run_exec(command) + else: + logging.info("Nodes are already at target version. Skipping.") + + +@task +def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): + """Verifies that the upgrade was successful.""" + cluster_name = node_pool_info.cluster_name + region = node_pool_info.region + + command = ( + f"gcloud container clusters describe {cluster_name} " + f"--region={region} --format='json'" + ) + logging.info("Running command: %s", command) + stdout = subprocess.run_exec(command) + + output_json = json.loads(stdout) + current_master_version = output_json.get("currentMasterVersion") + current_node_version = output_json.get("currentNodeVersion") + + logging.info("Verifying versions against target: %s", target_version) + logging.info("Post-upgrade Master Version: %s", current_master_version) + logging.info("Post-upgrade Node Version: %s", current_node_version) + + if ( + current_master_version != target_version + or current_node_version != target_version + ): + raise ValueError( + f"Verification failed! Master: {current_master_version}, " + f"Node: {current_node_version}, Target: {target_version}" + ) + logging.info("Upgrade verified successfully!") + + +with models.DAG( + dag_id=DAG_ID, + start_date=datetime.datetime(2025, 8, 1), + schedule=None, + catchup=False, + tags=["gke", "upgrade"], + description="DAG to upgrade GKE cluster to latest available version", +) as dag: + for machine in MachineConfigMap: + config = machine.value + + with TaskGroup(group_id=f"v{config.tpu_version.value}"): + node_pool_info = node_pool.build_node_pool_info_from_gcs_yaml.override( + task_id="build_node_pool_info_from_gcs_yaml" + )( + gcs_path=GCS_CONFIG_PATH, + dag_name=DAG_ID, + is_prod=composer_env.is_prod_env(), + machine_type=config.machine_version.value, + tpu_topology=config.tpu_topology, + ) + + avail_ver = find_available_version(node_pool_info) + curr_vers = find_current_cluster_version(node_pool_info) + + master_up = upgrade_master(avail_ver, curr_vers, node_pool_info) + node_up = upgrade_nodes(avail_ver, curr_vers, node_pool_info) + verify = verify_upgrade(avail_ver, node_pool_info) + + curr_vers >> master_up >> node_up >> verify diff --git a/dags/tpu_observability/utils/node_pool_util.py b/dags/tpu_observability/utils/node_pool_util.py index c9ea6fdcd..155c9e45e 100644 --- a/dags/tpu_observability/utils/node_pool_util.py +++ b/dags/tpu_observability/utils/node_pool_util.py @@ -214,9 +214,17 @@ def create( if _node_pool_exists(node_pool): logging.info( - f"Node pool {node_pool.node_pool_name} already exists. Skipping." + f"Node pool {node_pool.node_pool_name} already exists. Deleting before" + " recreation." ) - return + delete_command = ( + f"gcloud container node-pools delete {node_pool.node_pool_name} " + f"--project={node_pool.project_id} " + f"--cluster={node_pool.cluster_name} " + f"--location={node_pool.location} " + "--quiet" + ) + subprocess.run_exec(delete_command) command = ( f"gcloud container node-pools create {node_pool.node_pool_name} " @@ -932,3 +940,41 @@ def update(node_pool: Info, spec: NodePoolUpdateSpec) -> TimeUtil: subprocess.run_exec(update_cmd) return operation_start_time + + +def describe_cluster(node_pool_info: Info) -> str: + """Describes the GKE cluster using gcloud command.""" + command = ( + f"gcloud container clusters describe {node_pool_info.cluster_name} " + f"--project={node_pool_info.project_id} " + f"--region={node_pool_info.region} " + "--format='json'" + ) + return subprocess.run_exec(command) + + +def upgrade_cluster_master(node_pool_info: Info, latest_version: str) -> str: + """Upgrades the master of the GKE cluster.""" + command = ( + f"gcloud container clusters upgrade {node_pool_info.cluster_name} " + "--master " + f"--cluster-version={latest_version} " + f"--project={node_pool_info.project_id} " + f"--region={node_pool_info.region} --quiet" + ) + return subprocess.run_exec(command) + + +def upgrade_cluster_node_pool( + node_pool_info: Info, latest_version: str, node_pool_name: str = "default-pool" +) -> str: + """Upgrades a specific node pool of the GKE cluster.""" + command = ( + f"gcloud container clusters upgrade {node_pool_info.cluster_name} " + f"--project={node_pool_info.project_id} " + f"--region={node_pool_info.region} " + f"--cluster-version={latest_version} " + f"--node-pool={node_pool_name} " + "--quiet" + ) + return subprocess.run_exec(command) From d57ea5b3c8cf8badd2707ed5e06937ff099c8b06 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Tue, 28 Apr 2026 06:53:15 +0000 Subject: [PATCH 3/9] format --- dags/tpu_observability/utils/node_pool_util.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dags/tpu_observability/utils/node_pool_util.py b/dags/tpu_observability/utils/node_pool_util.py index 155c9e45e..402b4fe08 100644 --- a/dags/tpu_observability/utils/node_pool_util.py +++ b/dags/tpu_observability/utils/node_pool_util.py @@ -966,7 +966,9 @@ def upgrade_cluster_master(node_pool_info: Info, latest_version: str) -> str: def upgrade_cluster_node_pool( - node_pool_info: Info, latest_version: str, node_pool_name: str = "default-pool" + node_pool_info: Info, + latest_version: str, + node_pool_name: str = "default-pool", ) -> str: """Upgrades a specific node pool of the GKE cluster.""" command = ( From 23e02c2a33e962f632d0c00dff328778953e29d0 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Tue, 28 Apr 2026 08:11:15 +0000 Subject: [PATCH 4/9] scheduling_helper update --- dags/common/scheduling_helper/scheduling_helper.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dags/common/scheduling_helper/scheduling_helper.py b/dags/common/scheduling_helper/scheduling_helper.py index e0cff02dc..18cb74a40 100644 --- a/dags/common/scheduling_helper/scheduling_helper.py +++ b/dags/common/scheduling_helper/scheduling_helper.py @@ -64,6 +64,7 @@ class DayOfWeek(enum.Enum): "jobset_uptime_validation": dt.timedelta(minutes=90), "jobset_ttr_drain_restart": DefaultTimeout, "tpu_info_metrics_verification": DefaultTimeout, + "gke_cluster_version_manager": DefaultTimeout, }, TPU_INTERRUPTION_MOCK_CLUSTER.name: { "validate_interruption_count_gce_bare_metal_preemption": DefaultTimeout, From 95ce4db078ef49e5d94ce24a2668f3299a16c7a5 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Tue, 28 Apr 2026 08:23:53 +0000 Subject: [PATCH 5/9] fix --- .../gke_cluster_version_manager.py | 46 +++---------------- 1 file changed, 7 insertions(+), 39 deletions(-) diff --git a/dags/tpu_observability/gke_cluster_version_manager.py b/dags/tpu_observability/gke_cluster_version_manager.py index df36f4797..8870535c4 100644 --- a/dags/tpu_observability/gke_cluster_version_manager.py +++ b/dags/tpu_observability/gke_cluster_version_manager.py @@ -33,12 +33,11 @@ @task def find_available_version(node_pool_info: node_pool.Info) -> str: """Finds the latest available GKE version.""" - region = node_pool_info.region - if not region: - raise ValueError("Region not found in node_pool_info") command = ( - f"gcloud container get-server-config --region={region} --format='json'" + f"gcloud container get-server-config --region={node_pool_info.region} " + f"--project={node_pool_info.project_id} " + "--format='json'" ) logging.info("Running command: %s", command) stdout = subprocess.run_exec(command) @@ -46,7 +45,6 @@ def find_available_version(node_pool_info: node_pool.Info) -> str: output_json = json.loads(stdout) valid_versions = output_json.get("validMasterVersions", []) - # Filter: ^1\.(3[2-9]|[4-9][0-9]) pattern = re.compile(r"^1\.(3[2-9]|[4-9][0-9])") matching_versions = [v for v in valid_versions if pattern.match(v)] @@ -61,17 +59,8 @@ def find_available_version(node_pool_info: node_pool.Info) -> str: @task def find_current_cluster_version(node_pool_info: node_pool.Info) -> dict: """Finds the current version of the cluster.""" - cluster_name = node_pool_info.cluster_name - region = node_pool_info.region - if not cluster_name or not region: - raise ValueError("cluster_name or region not found in node_pool_info") - command = ( - f"gcloud container clusters describe {cluster_name} " - f"--region={region} --format='json'" - ) - logging.info("Running command: %s", command) - stdout = subprocess.run_exec(command) + stdout = node_pool.describe_cluster(node_pool_info) output_json = json.loads(stdout) current_master_version = output_json.get("currentMasterVersion") @@ -92,8 +81,6 @@ def upgrade_master( ): """Upgrades the master to the target version if needed.""" current_master = current_versions.get("currentMasterVersion") - cluster_name = node_pool_info.cluster_name - region = node_pool_info.region if current_master != latest_version: logging.info( @@ -101,12 +88,7 @@ def upgrade_master( current_master, latest_version, ) - command = ( - f"gcloud container clusters upgrade {cluster_name} --master " - f"--cluster-version={latest_version} --region={region} --quiet" - ) - logging.info("Running command: %s", command) - subprocess.run_exec(command) + node_pool.upgrade_cluster_master(node_pool_info, latest_version) else: logging.info("Master is already at target version. Skipping.") @@ -117,8 +99,6 @@ def upgrade_nodes( ): """Upgrades all node pools to the target version if needed.""" current_node = current_versions.get("currentNodeVersion") - cluster_name = node_pool_info.cluster_name - region = node_pool_info.region if current_node != latest_version: logging.info( @@ -126,12 +106,7 @@ def upgrade_nodes( current_node, latest_version, ) - command = ( - f"gcloud container clusters upgrade {cluster_name} " - f"--region={region} --cluster-version={latest_version} --quiet" - ) - logging.info("Running command: %s", command) - subprocess.run_exec(command) + node_pool.upgrade_cluster_node_pool(node_pool_info, latest_version) else: logging.info("Nodes are already at target version. Skipping.") @@ -139,15 +114,8 @@ def upgrade_nodes( @task def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): """Verifies that the upgrade was successful.""" - cluster_name = node_pool_info.cluster_name - region = node_pool_info.region - command = ( - f"gcloud container clusters describe {cluster_name} " - f"--region={region} --format='json'" - ) - logging.info("Running command: %s", command) - stdout = subprocess.run_exec(command) + stdout = node_pool.describe_cluster(node_pool_info) output_json = json.loads(stdout) current_master_version = output_json.get("currentMasterVersion") From 6dc157d30c65319f17f90c0e509a3808a4136cb0 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 15 May 2026 03:35:55 +0000 Subject: [PATCH 6/9] Move function describe_cluster, upgrade_cluster_master and upgrade_cluster_node_pool to gke_cluster_version_manager.py --- .../gke_cluster_version_manager.py | 50 ++++++++++++++++-- .../tpu_observability/utils/node_pool_util.py | 52 +------------------ 2 files changed, 48 insertions(+), 54 deletions(-) diff --git a/dags/tpu_observability/gke_cluster_version_manager.py b/dags/tpu_observability/gke_cluster_version_manager.py index 8870535c4..0ff492675 100644 --- a/dags/tpu_observability/gke_cluster_version_manager.py +++ b/dags/tpu_observability/gke_cluster_version_manager.py @@ -30,6 +30,48 @@ DAG_ID = "gke_cluster_version_manager" +def describe_cluster(node_pool_info: node_pool.Info) -> str: + """Describes the GKE cluster using gcloud command.""" + command = ( + f"gcloud container clusters describe {node_pool_info.cluster_name} " + f"--project={node_pool_info.project_id} " + f"--region={node_pool_info.region} " + "--format='json'" + ) + return subprocess.run_exec(command) + + +def upgrade_cluster_master( + node_pool_info: node_pool.Info, latest_version: str +) -> str: + """Upgrades the master of the GKE cluster.""" + command = ( + f"gcloud container clusters upgrade {node_pool_info.cluster_name} " + "--master " + f"--cluster-version={latest_version} " + f"--project={node_pool_info.project_id} " + f"--region={node_pool_info.region} --quiet" + ) + return subprocess.run_exec(command) + + +def upgrade_cluster_node_pool( + node_pool_info: node_pool.Info, + latest_version: str, + node_pool_name: str = "default-pool", +) -> str: + """Upgrades a specific node pool of the GKE cluster.""" + command = ( + f"gcloud container clusters upgrade {node_pool_info.cluster_name} " + f"--project={node_pool_info.project_id} " + f"--region={node_pool_info.region} " + f"--cluster-version={latest_version} " + f"--node-pool={node_pool_name} " + "--quiet" + ) + return subprocess.run_exec(command) + + @task def find_available_version(node_pool_info: node_pool.Info) -> str: """Finds the latest available GKE version.""" @@ -60,7 +102,7 @@ def find_available_version(node_pool_info: node_pool.Info) -> str: def find_current_cluster_version(node_pool_info: node_pool.Info) -> dict: """Finds the current version of the cluster.""" - stdout = node_pool.describe_cluster(node_pool_info) + stdout = describe_cluster(node_pool_info) output_json = json.loads(stdout) current_master_version = output_json.get("currentMasterVersion") @@ -88,7 +130,7 @@ def upgrade_master( current_master, latest_version, ) - node_pool.upgrade_cluster_master(node_pool_info, latest_version) + upgrade_cluster_master(node_pool_info, latest_version) else: logging.info("Master is already at target version. Skipping.") @@ -106,7 +148,7 @@ def upgrade_nodes( current_node, latest_version, ) - node_pool.upgrade_cluster_node_pool(node_pool_info, latest_version) + upgrade_cluster_node_pool(node_pool_info, latest_version) else: logging.info("Nodes are already at target version. Skipping.") @@ -115,7 +157,7 @@ def upgrade_nodes( def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): """Verifies that the upgrade was successful.""" - stdout = node_pool.describe_cluster(node_pool_info) + stdout = describe_cluster(node_pool_info) output_json = json.loads(stdout) current_master_version = output_json.get("currentMasterVersion") diff --git a/dags/tpu_observability/utils/node_pool_util.py b/dags/tpu_observability/utils/node_pool_util.py index 402b4fe08..c9ea6fdcd 100644 --- a/dags/tpu_observability/utils/node_pool_util.py +++ b/dags/tpu_observability/utils/node_pool_util.py @@ -214,17 +214,9 @@ def create( if _node_pool_exists(node_pool): logging.info( - f"Node pool {node_pool.node_pool_name} already exists. Deleting before" - " recreation." + f"Node pool {node_pool.node_pool_name} already exists. Skipping." ) - delete_command = ( - f"gcloud container node-pools delete {node_pool.node_pool_name} " - f"--project={node_pool.project_id} " - f"--cluster={node_pool.cluster_name} " - f"--location={node_pool.location} " - "--quiet" - ) - subprocess.run_exec(delete_command) + return command = ( f"gcloud container node-pools create {node_pool.node_pool_name} " @@ -940,43 +932,3 @@ def update(node_pool: Info, spec: NodePoolUpdateSpec) -> TimeUtil: subprocess.run_exec(update_cmd) return operation_start_time - - -def describe_cluster(node_pool_info: Info) -> str: - """Describes the GKE cluster using gcloud command.""" - command = ( - f"gcloud container clusters describe {node_pool_info.cluster_name} " - f"--project={node_pool_info.project_id} " - f"--region={node_pool_info.region} " - "--format='json'" - ) - return subprocess.run_exec(command) - - -def upgrade_cluster_master(node_pool_info: Info, latest_version: str) -> str: - """Upgrades the master of the GKE cluster.""" - command = ( - f"gcloud container clusters upgrade {node_pool_info.cluster_name} " - "--master " - f"--cluster-version={latest_version} " - f"--project={node_pool_info.project_id} " - f"--region={node_pool_info.region} --quiet" - ) - return subprocess.run_exec(command) - - -def upgrade_cluster_node_pool( - node_pool_info: Info, - latest_version: str, - node_pool_name: str = "default-pool", -) -> str: - """Upgrades a specific node pool of the GKE cluster.""" - command = ( - f"gcloud container clusters upgrade {node_pool_info.cluster_name} " - f"--project={node_pool_info.project_id} " - f"--region={node_pool_info.region} " - f"--cluster-version={latest_version} " - f"--node-pool={node_pool_name} " - "--quiet" - ) - return subprocess.run_exec(command) From a70a50fda5fe98320422b763517ef2581e2bef35 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Wed, 20 May 2026 02:51:09 +0000 Subject: [PATCH 7/9] format fix --- dags/tpu_observability/gke_cluster_version_manager.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/dags/tpu_observability/gke_cluster_version_manager.py b/dags/tpu_observability/gke_cluster_version_manager.py index 0ff492675..7582abfd3 100644 --- a/dags/tpu_observability/gke_cluster_version_manager.py +++ b/dags/tpu_observability/gke_cluster_version_manager.py @@ -21,6 +21,7 @@ from airflow import models from airflow.decorators import task +from airflow.models.baseoperator import chain from airflow.utils.task_group import TaskGroup from dags import composer_env from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH @@ -207,4 +208,8 @@ def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): node_up = upgrade_nodes(avail_ver, curr_vers, node_pool_info) verify = verify_upgrade(avail_ver, node_pool_info) - curr_vers >> master_up >> node_up >> verify + chain( + master_up, + node_up, + verify, + ) From 5eb9215f55f22b279763d337f74faedd846ec81e Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Wed, 20 May 2026 07:11:38 +0000 Subject: [PATCH 8/9] add dag timeout --- dags/tpu_observability/gke_cluster_version_manager.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/dags/tpu_observability/gke_cluster_version_manager.py b/dags/tpu_observability/gke_cluster_version_manager.py index 7582abfd3..9838b3c07 100644 --- a/dags/tpu_observability/gke_cluster_version_manager.py +++ b/dags/tpu_observability/gke_cluster_version_manager.py @@ -27,8 +27,11 @@ from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils import subprocess_util as subprocess +from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout DAG_ID = "gke_cluster_version_manager" +DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) +SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) def describe_cluster(node_pool_info: node_pool.Info) -> str: @@ -181,9 +184,9 @@ def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): with models.DAG( dag_id=DAG_ID, - start_date=datetime.datetime(2025, 8, 1), - schedule=None, - catchup=False, + start_date=datetime.datetime(2026, 5, 20), + schedule=SCHEDULE if composer_env.is_prod_env() else None, + dagrun_timeout=DAGRUN_TIMEOUT, tags=["gke", "upgrade"], description="DAG to upgrade GKE cluster to latest available version", ) as dag: From f3794a8cf81208dd6f1af48dc36715e078967045 Mon Sep 17 00:00:00 2001 From: yuna-tzeng Date: Fri, 22 May 2026 03:50:56 +0000 Subject: [PATCH 9/9] fix --- .../tpu_observability/gke_cluster_version_manager.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/dags/tpu_observability/gke_cluster_version_manager.py b/dags/tpu_observability/gke_cluster_version_manager.py index 9838b3c07..7f8a71e29 100644 --- a/dags/tpu_observability/gke_cluster_version_manager.py +++ b/dags/tpu_observability/gke_cluster_version_manager.py @@ -194,7 +194,7 @@ def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): config = machine.value with TaskGroup(group_id=f"v{config.tpu_version.value}"): - node_pool_info = node_pool.build_node_pool_info_from_gcs_yaml.override( + cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override( task_id="build_node_pool_info_from_gcs_yaml" )( gcs_path=GCS_CONFIG_PATH, @@ -204,12 +204,12 @@ def verify_upgrade(target_version: str, node_pool_info: node_pool.Info): tpu_topology=config.tpu_topology, ) - avail_ver = find_available_version(node_pool_info) - curr_vers = find_current_cluster_version(node_pool_info) + avail_ver = find_available_version(cluster_info) + curr_vers = find_current_cluster_version(cluster_info) - master_up = upgrade_master(avail_ver, curr_vers, node_pool_info) - node_up = upgrade_nodes(avail_ver, curr_vers, node_pool_info) - verify = verify_upgrade(avail_ver, node_pool_info) + master_up = upgrade_master(avail_ver, curr_vers, cluster_info) + node_up = upgrade_nodes(avail_ver, curr_vers, cluster_info) + verify = verify_upgrade(avail_ver, cluster_info) chain( master_up,