From 0791ccb0427f21fb993e2fc5a807b447c7761bc1 Mon Sep 17 00:00:00 2001 From: Jim Tseng Date: Thu, 18 Jun 2026 08:37:15 +0800 Subject: [PATCH 1/4] fix: Skip GCS configuration download in GitHub Actions environment During GitHub Actions (CI) pipeline execution, native Airflow methods and cloud-dependent functions cause blocking errors due to the lack of an active database connection or cloud permissions. --- dags/common/quarantined_tests.py | 31 +++++++++++++++------ dags/tpu_observability/utils/jobset_util.py | 24 ++++++++-------- xlml/apis/gcs.py | 22 +++++++++++---- 3 files changed, 51 insertions(+), 26 deletions(-) diff --git a/dags/common/quarantined_tests.py b/dags/common/quarantined_tests.py index 6d44fcc33..d27a6cfd4 100644 --- a/dags/common/quarantined_tests.py +++ b/dags/common/quarantined_tests.py @@ -54,13 +54,27 @@ def match_quarantine_patterns( return False +def is_ci_environment() -> bool: + """Checks if running in a CI environment. + + Identifies environments where framework-specific runtime dependencies do not + exist. For example, GitHub Actions lacks the actual Airflow environment and + its components (such as database context or cloud credentials), requiring code + to bypass these dependencies and use mocks instead. + + Returns: + bool: True if in CI (GitHub Actions), False otherwise. + """ + return os.getenv("GITHUB_ACTIONS", "false").lower() == "true" + + def safe_get_from_variable(key: str, default_var: str): """ - Check whether the current runtime is GitHub Actions. Skip retrieving variables in GitHub Actions to avoid excessive log output. + Check whether the current runtime is GitHub Actions. Skip retrieving variables + in GitHub Actions to avoid excessive log output. """ value = default_var - is_ci_env = os.getenv("GITHUB_ACTIONS", "false").lower() == "true" - if is_ci_env: + if is_ci_environment(): logging.info("In GitHub Actions, skip getting variables") else: value = Variable.get(key, default_var=default_var) @@ -68,8 +82,8 @@ def safe_get_from_variable(key: str, default_var: str): """ -The quarantine list is defined by a set of UNIX Shell Glob Patterns. -These patterns are used to match and quarantine tests. +The quarantine list is defined by a set of UNIX Shell Glob Patterns. +These patterns are used to match and quarantine tests. The patterns are stored in the Airflow Variable named 'quarantine_patterns'. @@ -90,9 +104,10 @@ def is_quarantined(test_name) -> bool: """ Checks if a test is quarantined using both legacy and current methods. - The legacy method checks if `test_name` is present in `QuarantineTests.tests`. - The current method checks against a runtime quarantine list fetched from Airflow Variables - (key: 'quarantine_list') using `is_in_runtime_quarantine_list()`. + The legacy method checks if `test_name` is present in + `QuarantineTests.tests`. The current method checks against a runtime + quarantine list fetched from Airflow Variables (key: 'quarantine_list') + using `is_in_runtime_quarantine_list()`. The test is considered quarantined if it's found by either method. """ diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index 112751ae8..8e8a82591 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -282,18 +282,18 @@ class JobSet: recreating pods. Defaults to False. """ - namespace: str - max_restarts: int - replicated_job_name: str - replicas: int - backoff_limit: int - completions: int - parallelism: int - tpu_accelerator_type: str - tpu_topology: str - container_name: str - image: str - tpu_cores_per_pod: int + namespace: str = "" + max_restarts: int = -1 + replicated_job_name: str = "" + replicas: int = -1 + backoff_limit: int = -1 + completions: int = -1 + parallelism: int = -1 + tpu_accelerator_type: str = "" + tpu_topology: str = "" + container_name: str = "" + image: str = "" + tpu_cores_per_pod: int = -1 privileged: bool = False dag_id_prefix: str = "" delay_recovery: bool = False diff --git a/xlml/apis/gcs.py b/xlml/apis/gcs.py index e2f483288..cab105ce6 100644 --- a/xlml/apis/gcs.py +++ b/xlml/apis/gcs.py @@ -14,16 +14,17 @@ """Functions for GCS Bucket""" +from absl import logging import os import re import tempfile from typing import List +import yaml -from absl import logging from airflow.decorators import task from airflow.hooks.subprocess import SubprocessHook from airflow.providers.google.cloud.operators.gcs import GCSHook -import yaml +from dags.common.quarantined_tests import is_ci_environment def obtain_file_list(gcs_path: str) -> List[str]: @@ -97,6 +98,14 @@ def load_yaml_from_gcs(gcs_path: str) -> dict: """Loads and parses the DAG configuration YAML file from GCS.""" logging.info(f"Attempting to load config from: {gcs_path}") + # Workflows triggered by fork PRs do not have access to upstream secrets. + # Bypass GCS loading strictly to allow the CI workflow to pass successfully. + if is_ci_environment(): + logging.info( + "In GitHub Actions, skip load_yaml_from_gcs and return an empty dict." + ) + return {} + if not gcs_path.startswith("gs://"): raise ValueError( f"Invalid GCS path: '{gcs_path}'. Path must start with 'gs://'." @@ -106,8 +115,8 @@ def load_yaml_from_gcs(gcs_path: str) -> dict: gcs_path.lower().endswith(".yaml") or gcs_path.lower().endswith(".yml") ): logging.warning( - f"GCS path '{gcs_path}' does not have a typical YAML extension (.yaml or .yml). " - "Proceeding, but be aware this might not be a YAML file." + f"GCS path '{gcs_path}' does not have a typical YAML extension (.yaml " + "or .yml). Proceeding, but be aware this might not be a YAML file." ) with tempfile.TemporaryDirectory() as tmpdir: @@ -123,8 +132,9 @@ def load_yaml_from_gcs(gcs_path: str) -> dict: if not os.path.exists(temp_file_path): logging.error( - f"gcloud storage cp command completed, but '{temp_file_path}' was not created. " - "This often means the copy failed. Check gcloud storage stdout/stderr in logs." + f"gcloud storage cp command completed, but '{temp_file_path}' " + "was not created. This often means the copy failed. " + "Check gcloud storage stdout/stderr in logs." ) raise FileNotFoundError( f"[Errno 2] Failed to download file from GCS path: {gcs_path}" From fdcbae901723e5672139aabb2b1e2e2f01bdebc3 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 2/4] [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 b85625224..c1d3e7cf3 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 461bf802bc75aca3e01a569f7075ed0eb932381c Mon Sep 17 00:00:00 2001 From: Jim1428 Date: Mon, 15 Jun 2026 14:36:57 +0800 Subject: [PATCH 3/4] feat: Test TaskGroupWithTimeout --- .../scheduling_helper/scheduling_helper.py | 1 + .../test_task_group_with_timeout.py | 126 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 dags/common/task_group_with_timeout/test_task_group_with_timeout.py diff --git a/dags/common/scheduling_helper/scheduling_helper.py b/dags/common/scheduling_helper/scheduling_helper.py index 8aef11b56..a4e0409d7 100644 --- a/dags/common/scheduling_helper/scheduling_helper.py +++ b/dags/common/scheduling_helper/scheduling_helper.py @@ -74,6 +74,7 @@ class DayOfWeek(enum.Enum): "jobset_ttr_drain_restart": DefaultTimeout, "tpu_info_metrics_verification": DefaultTimeout, "jobset_ttr_node_reboot": dt.timedelta(minutes=90), + "test_task_group_with_timeout": DefaultTimeout, }, TPU_INTERRUPTION_MOCK_CLUSTER.name: { "validate_interruption_count_gce_bare_metal_preemption": DefaultTimeout, diff --git a/dags/common/task_group_with_timeout/test_task_group_with_timeout.py b/dags/common/task_group_with_timeout/test_task_group_with_timeout.py new file mode 100644 index 000000000..bacae8980 --- /dev/null +++ b/dags/common/task_group_with_timeout/test_task_group_with_timeout.py @@ -0,0 +1,126 @@ +"""A integration DAG to test the behavior of TaskGroupWithTimeout.""" + +import datetime +from datetime import timedelta +import logging +import time + +from airflow import models +from airflow.models.baseoperator import chain +from airflow.operators.python import PythonOperator +from airflow.exceptions import AirflowException +from airflow.utils.task_group import TaskGroup +from airflow.utils.trigger_rule import TriggerRule + +from dags.common.scheduling_helper.scheduling_helper import ( + SchedulingHelper, + get_dag_timeout, +) +from dags.common.task_group_with_timeout.task_group_with_timeout import ( + TaskGroupWithTimeout, +) + +DAG_ID = "test_task_group_with_timeout" +DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) +SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) + + +def simulate_workload(sleep_seconds: int): + """Simulates a long-running workload with explicit sleep.""" + logging.info(f"Executing workload... Sleeping for {sleep_seconds}s.") + time.sleep(sleep_seconds) + logging.info("Workload completed successfully.") + + +# Test 1: Active enforcement against nested TaskGroups. +try: + with TaskGroupWithTimeout( + group_id="isolated_nested_test", timeout=timedelta(seconds=60) + ): + with TaskGroup(group_id="nested_native_group"): + pass +except AirflowException as e: + print(f"Standalone Parsing Check - Caught expected nested error: {e}") + +# Test 2: Active enforcement against Dynamic Task Mapping inside the group. +try: + with TaskGroupWithTimeout( + group_id="isolated_mapping_test", timeout=timedelta(seconds=60) + ): + PythonOperator.partial( + task_id="dummy_map", python_callable=lambda: None + ).expand(op_args=[[]]) +except AirflowException as e: + print(f"Standalone Parsing Check - Caught expected mapping error: {e}") + +# Runtime Timeout Tests +with models.DAG( + dag_id=DAG_ID, + start_date=datetime.datetime(2026, 6, 1), + schedule=None, + dagrun_timeout=DAGRUN_TIMEOUT, + catchup=False, + tags=["integration-test", "timeout-validation"], + description=( + "Validates continuous time-budget pool depletion across consecutive " + "TaskGroups and tests teardown resource protection safeguards." + ), + doc_md=""" + # TaskGroupWithTimeout Integration Test Suite + + ### Description + This DAG serves as a production-grade sandbox to verify the dynamic runtime + capabilities of the custom `TaskGroupWithTimeout` infrastructure component. + + ### Procedures + 1. **Case 1 (Normal Flow):** Consumes 60 seconds from the global pool and succeeds. + 2. **Case 2 (Timeout Flow):** Dynamically inherits the remaining budget. Forces `t3` + to fail due to budget exhaustion. + 3. **Teardown & Final Status:** Ensures external cleanup runs and forces the final + DAG status to SUCCESS via an ALL_DONE end-node. + """, +) as dag: + # Case 1: Normal Flow + with TaskGroupWithTimeout( + group_id="case1_normal_flow", timeout=timedelta(seconds=100) + ) as case1: + c1_t1 = PythonOperator( + task_id="t1", python_callable=simulate_workload, op_args=[20] + ) + c1_t2 = PythonOperator( + task_id="t2", python_callable=simulate_workload, op_args=[20] + ) + c1_t3 = PythonOperator( + task_id="t3", python_callable=simulate_workload, op_args=[20] + ) + + chain(c1_t1, c1_t2, c1_t3) + + # Case 2: Timeout Block & Teardown Bypass + with TaskGroupWithTimeout( + group_id="case2_block_and_teardown", timeout=timedelta(seconds=100) + ) as case2: + c2_t1 = PythonOperator( + task_id="t1", python_callable=simulate_workload, op_args=[20] + ) + c2_t2 = PythonOperator( + task_id="t2", python_callable=simulate_workload, op_args=[40] + ) + c2_t3 = PythonOperator( + task_id="t3", python_callable=simulate_workload, op_args=[60], retries=0 + ) + + c2_teardown = PythonOperator( + task_id="cleanup", python_callable=simulate_workload, op_args=[5] + ).as_teardown() + + dag_final_success = PythonOperator( + task_id="dag_final_success", + python_callable=lambda: logging.info( + "Test pipeline finished. Overriding final status to SUCCESS." + ), + trigger_rule=TriggerRule.ALL_DONE, + ) + + chain(c2_t1, c2_t2, c2_t3, c2_teardown) + chain(case1, case2, dag_final_success) From 82fe1b1820dfd068b6da41bfb5c0f2de538eb083 Mon Sep 17 00:00:00 2001 From: Jim1428 Date: Thu, 25 Jun 2026 13:50:25 +0800 Subject: [PATCH 4/4] fix: Relocate static defenses into DAG context and add Prerequisites --- .../test_task_group_with_timeout.py | 130 ++++++++++++------ 1 file changed, 86 insertions(+), 44 deletions(-) diff --git a/dags/common/task_group_with_timeout/test_task_group_with_timeout.py b/dags/common/task_group_with_timeout/test_task_group_with_timeout.py index bacae8980..c04713964 100644 --- a/dags/common/task_group_with_timeout/test_task_group_with_timeout.py +++ b/dags/common/task_group_with_timeout/test_task_group_with_timeout.py @@ -12,6 +12,7 @@ from airflow.utils.task_group import TaskGroup from airflow.utils.trigger_rule import TriggerRule +from dags import composer_env from dags.common.scheduling_helper.scheduling_helper import ( SchedulingHelper, get_dag_timeout, @@ -32,32 +33,47 @@ def simulate_workload(sleep_seconds: int): logging.info("Workload completed successfully.") -# Test 1: Active enforcement against nested TaskGroups. -try: - with TaskGroupWithTimeout( - group_id="isolated_nested_test", timeout=timedelta(seconds=60) - ): - with TaskGroup(group_id="nested_native_group"): - pass -except AirflowException as e: - print(f"Standalone Parsing Check - Caught expected nested error: {e}") - -# Test 2: Active enforcement against Dynamic Task Mapping inside the group. -try: - with TaskGroupWithTimeout( - group_id="isolated_mapping_test", timeout=timedelta(seconds=60) - ): - PythonOperator.partial( - task_id="dummy_map", python_callable=lambda: None +def verify_static_defenses(): + """Verifies that TaskGroupWithTimeout structural restrictions are active.""" + # Mock a minimal dummy DAG context to allow TaskGroup initialization + mock_dag = models.DAG( + dag_id="mock_static_defense_dag", start_date=datetime.datetime(2026, 6, 1) + ) + + # Test 1: Active enforcement against nested TaskGroups + try: + tg_timeout = TaskGroupWithTimeout( + group_id="parent_timeout_group", + timeout=timedelta(seconds=60), + dag=mock_dag, + ) + nested_native = TaskGroup(group_id="nested_native_group", dag=mock_dag) + # Explicitly trigger the component's add constraint logic + tg_timeout.add(nested_native) + except AirflowException as e: + print(f"Static Parsing Check - Caught expected nested error: {e}") + + # Test 2: Active enforcement against Dynamic Task Mapping inside the group + try: + tg_timeout_map = TaskGroupWithTimeout( + group_id="parent_mapping_group", + timeout=timedelta(seconds=60), + dag=mock_dag, + ) + dummy_mapped_task = PythonOperator.partial( + task_id="dummy_map", python_callable=lambda: None, dag=mock_dag ).expand(op_args=[[]]) -except AirflowException as e: - print(f"Standalone Parsing Check - Caught expected mapping error: {e}") + # Explicitly trigger the component's add constraint logic + tg_timeout_map.add(dummy_mapped_task) + except AirflowException as e: + print(f"Static Parsing Check - Caught expected mapping error: {e}") + # Runtime Timeout Tests with models.DAG( dag_id=DAG_ID, start_date=datetime.datetime(2026, 6, 1), - schedule=None, + schedule=SCHEDULE if composer_env.is_prod_env() else None, dagrun_timeout=DAGRUN_TIMEOUT, catchup=False, tags=["integration-test", "timeout-validation"], @@ -72,55 +88,81 @@ def simulate_workload(sleep_seconds: int): This DAG serves as a production-grade sandbox to verify the dynamic runtime capabilities of the custom `TaskGroupWithTimeout` infrastructure component. + ### Prerequisites + The custom TaskGroupWithTimeout component must be available under dags.common. + ### Procedures - 1. **Case 1 (Normal Flow):** Consumes 60 seconds from the global pool and succeeds. - 2. **Case 2 (Timeout Flow):** Dynamically inherits the remaining budget. Forces `t3` + 1. **Static Validation:** Pre-checks blocking of nested groups and dynamic mapping. + 2. **Case 1 (Normal Flow):** Consumes 60 seconds from the global pool and succeeds. + 3. **Case 2 (Timeout Flow):** Dynamically inherits the remaining budget. Forces `t3` to fail due to budget exhaustion. - 3. **Teardown & Final Status:** Ensures external cleanup runs and forces the final + 4. **Teardown & Final Status:** Ensures external cleanup runs and forces the final DAG status to SUCCESS via an ALL_DONE end-node. """, ) as dag: + verify_static_defenses() + # Case 1: Normal Flow with TaskGroupWithTimeout( - group_id="case1_normal_flow", timeout=timedelta(seconds=100) + group_id="normal_flow", timeout=timedelta(seconds=100) ) as case1: - c1_t1 = PythonOperator( - task_id="t1", python_callable=simulate_workload, op_args=[20] + normal_flow_task_1 = PythonOperator( + task_id="normal_flow_task_1", + python_callable=simulate_workload, + op_args=[20], ) - c1_t2 = PythonOperator( - task_id="t2", python_callable=simulate_workload, op_args=[20] + normal_flow_task_2 = PythonOperator( + task_id="normal_flow_task_2", + python_callable=simulate_workload, + op_args=[20], ) - c1_t3 = PythonOperator( - task_id="t3", python_callable=simulate_workload, op_args=[20] + normal_flow_task_3 = PythonOperator( + task_id="normal_flow_task_3", + python_callable=simulate_workload, + op_args=[20], ) - chain(c1_t1, c1_t2, c1_t3) + chain(normal_flow_task_1, normal_flow_task_2, normal_flow_task_3) # Case 2: Timeout Block & Teardown Bypass with TaskGroupWithTimeout( - group_id="case2_block_and_teardown", timeout=timedelta(seconds=100) + group_id="timeout_flow", timeout=timedelta(seconds=100) ) as case2: - c2_t1 = PythonOperator( - task_id="t1", python_callable=simulate_workload, op_args=[20] + timeout_flow_task_1 = PythonOperator( + task_id="timeout_flow_task_1", + python_callable=simulate_workload, + op_args=[20], ) - c2_t2 = PythonOperator( - task_id="t2", python_callable=simulate_workload, op_args=[40] + timeout_flow_task_2 = PythonOperator( + task_id="timeout_flow_task_2", + python_callable=simulate_workload, + op_args=[40], ) - c2_t3 = PythonOperator( - task_id="t3", python_callable=simulate_workload, op_args=[60], retries=0 + timeout_flow_task_3 = PythonOperator( + task_id="timeout_flow_task_3", + python_callable=simulate_workload, + op_args=[60], + retries=0, ) - c2_teardown = PythonOperator( - task_id="cleanup", python_callable=simulate_workload, op_args=[5] + env_cleanup_teardown = PythonOperator( + task_id="env_cleanup_teardown", + python_callable=simulate_workload, + op_args=[5], ).as_teardown() - dag_final_success = PythonOperator( - task_id="dag_final_success", + dag_status_override = PythonOperator( + task_id="dag_status_override", python_callable=lambda: logging.info( "Test pipeline finished. Overriding final status to SUCCESS." ), trigger_rule=TriggerRule.ALL_DONE, ) - chain(c2_t1, c2_t2, c2_t3, c2_teardown) - chain(case1, case2, dag_final_success) + chain( + timeout_flow_task_1, + timeout_flow_task_2, + timeout_flow_task_3, + env_cleanup_teardown, + ) + chain(case1, case2, dag_status_override)