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/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, 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..7f8a71e29 --- /dev/null +++ b/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, + ) 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."