feat(tpu_observability): Add GKE cluster version manager to automate master and node pool upgrades

.github/workflows/dag-check.yml

@@ -3,7 +3,6 @@ name: DAG Check
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

.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

.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

.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
+          requireChecklist: false # If this is true and there are no checklists detected, the action will fail

.github/workflows/unit-test.yml

@@ -3,7 +3,6 @@ name: Unit Test
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

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,

dags/tpu_observability/gke_cluster_version_manager.py

@@ -0,0 +1,218 @@
+# 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.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
+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:
+  """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."""
+
+  command = (
+      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)
+
+  output_json = json.loads(stdout)
+  valid_versions = output_json.get("validMasterVersions", [])
+
+  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."""
+
+  stdout = describe_cluster(node_pool_info)
+
+  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")
+
+  if current_master != latest_version:
+    logging.info(
+        "Master version (%s) != Target (%s). Upgrading.",
+        current_master,
+        latest_version,
+    )
+    upgrade_cluster_master(node_pool_info, latest_version)
+  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")
+
+  if current_node != latest_version:
+    logging.info(
+        "Node version (%s) != Target (%s). Upgrading.",
+        current_node,
+        latest_version,
+    )
+    upgrade_cluster_node_pool(node_pool_info, latest_version)
+  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."""
+
+  stdout = describe_cluster(node_pool_info)
+
+  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(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:
+  for machine in MachineConfigMap:
+    config = machine.value
+
+    with TaskGroup(group_id=f"v{config.tpu_version.value}"):
+      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,
+          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(cluster_info)
+      curr_vers = find_current_cluster_version(cluster_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,
+          node_up,
+          verify,
+      )

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."