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/dags/tpu_observability/jobset_ttr_drain_restart.py b/dags/tpu_observability/jobset_ttr_drain_restart.py index 597b092f9..13d53dfc1 100644 --- a/dags/tpu_observability/jobset_ttr_drain_restart.py +++ b/dags/tpu_observability/jobset_ttr_drain_restart.py @@ -21,6 +21,11 @@ from airflow.models.baseoperator import chain from airflow.utils.task_group import TaskGroup from airflow.utils.trigger_rule import TriggerRule +from airflow.utils.task_group import TaskGroup +from dags.common.task_group_with_timeout import TaskGroupWithTimeout + + +from airflow.decorators import task from dags import composer_env from dags.common.scheduling_helper.scheduling_helper import ( @@ -42,6 +47,10 @@ DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + @task def check_nodes_number( @@ -140,75 +149,102 @@ def check_nodes_number( selector = jobset.generate_node_pool_selector(DAG_ID) jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - node_pool_selector=selector, - ) - - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - node_pool_selector=selector, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - select_node = node_pool.draw_random_node.override(task_id="select_node")( - node_pool=cluster_info - ) - - drained_node = node_pool.operate_node.override(task_id="drained_node")( - node_pool=cluster_info, - operation=NodeOperationSpec.Drain(), - node_name=select_node, - ) - - check_nodes_number = check_nodes_number.override( - task_id="check_nodes_number" - )( - pool=cluster_info, - drained_node_number=1, - ) - - uncordon_node = node_pool.operate_node.override(task_id="uncordon_node")( - node_pool=cluster_info, - operation=NodeOperationSpec.Uncordon(), - node_name=select_node, - ) - - wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override( - task_id="wait_for_metric_upload" - )( - node_pool=cluster_info, - jobset_name=jobset_name, - ) - - cleanup_workload = jobset.end_workload.override( - task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ).as_teardown( - setups=startup.jobset_start_time - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool" + )( + node_pool=cluster_info, + node_pool_selector=selector, + ) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + node_pool_selector=selector, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + select_node = node_pool.draw_random_node.override( + task_id="select_node" + )(node_pool=cluster_info) + + drained_node = node_pool.operate_node.override(task_id="drained_node")( + node_pool=cluster_info, + operation=NodeOperationSpec.Drain(), + node_name=select_node, + ) + + check_nodes_number_task = check_nodes_number.override( + task_id="check_nodes_number" + )( + pool=cluster_info, + drained_node_number=1, + ) + + uncordon_node = node_pool.operate_node.override( + task_id="uncordon_node" + )( + node_pool=cluster_info, + operation=NodeOperationSpec.Uncordon(), + node_name=select_node, + ) + + wait_for_metric_upload = ( + jobset.wait_for_jobset_ttr_to_be_found.override( + task_id="wait_for_metric_upload" + )( + node_pool=cluster_info, + jobset_name=jobset_name, + ) + ) + + chain( + *startup.tasks, + select_node, + drained_node, + check_nodes_number_task, + uncordon_node, + wait_for_metric_upload, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ).as_teardown( + setups=startup.jobset_start_time + ) + + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) + + chain( + cleanup_workload, + cleanup_node_pool, + ) chain( selector, jobset_name, - create_node_pool, - *startup.tasks, - select_node, - drained_node, - check_nodes_number, - uncordon_node, - wait_for_metric_upload, - cleanup_workload, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/jobset_ttr_node_pool_resize.py b/dags/tpu_observability/jobset_ttr_node_pool_resize.py index 60140d19b..d3d0a045b 100644 --- a/dags/tpu_observability/jobset_ttr_node_pool_resize.py +++ b/dags/tpu_observability/jobset_ttr_node_pool_resize.py @@ -34,12 +34,18 @@ from dags.tpu_observability.utils import jobset_util as jobset from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils.jobset_util import Workload +from dags.common.task_group_with_timeout import TaskGroupWithTimeout + DAG_ID = "jobset_ttr_node_pool_resize" DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) _DISK_SIZE_INCREMENT = 100 +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -108,76 +114,101 @@ jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - node_pool_selector=selector, - ) - - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - node_pool_selector=selector, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - node_pool_resize_start_time = node_pool.update.override( - task_id="node_pool_resize" - )( - node_pool=cluster_info, - spec=node_pool.NodePoolUpdateSpec.DiskSize( - delta=_DISK_SIZE_INCREMENT - ), - ) - - wait_for_recovery = jobset.wait_for_jobset_recovered.override( - task_id="wait_for_recovery" - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ) - - verify_duration = jobset.verify_recovery_duration.override( - task_id="verify_recovery_duration" - )( - start_time=node_pool_resize_start_time, - end_time=wait_for_recovery, - ) - - wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override( - task_id="wait_for_jobset_ttr_to_be_found", - )( - node_pool=cluster_info, - jobset_name=jobset_name, - start_time=node_pool_resize_start_time, - ) - - cleanup_workload = jobset.end_workload.override( - task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ).as_teardown( - setups=startup.jobset_start_time - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool" + )( + node_pool=cluster_info, + node_pool_selector=selector, + ) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + node_pool_selector=selector, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + node_pool_resize_start_time = node_pool.update.override( + task_id="node_pool_resize" + )( + node_pool=cluster_info, + spec=node_pool.NodePoolUpdateSpec.DiskSize( + delta=_DISK_SIZE_INCREMENT + ), + ) + + wait_for_recovery = jobset.wait_for_jobset_recovered.override( + task_id="wait_for_recovery" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ) + + verify_duration = jobset.verify_recovery_duration.override( + task_id="verify_recovery_duration" + )( + start_time=node_pool_resize_start_time, + end_time=wait_for_recovery, + ) + + wait_for_metric_upload = ( + jobset.wait_for_jobset_ttr_to_be_found.override( + task_id="wait_for_jobset_ttr_to_be_found", + )( + node_pool=cluster_info, + jobset_name=jobset_name, + start_time=node_pool_resize_start_time, + ) + ) + + chain( + *startup.tasks, + node_pool_resize_start_time, + wait_for_recovery, + verify_duration, + wait_for_metric_upload, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ).as_teardown( + setups=startup.jobset_start_time + ) + + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) + + chain( + cleanup_workload, + cleanup_node_pool, + ) chain( selector, jobset_name, - create_node_pool, - *startup.tasks, - node_pool_resize_start_time, - wait_for_recovery, - verify_duration, - wait_for_metric_upload, - cleanup_workload, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/jobset_ttr_pod_delete.py b/dags/tpu_observability/jobset_ttr_pod_delete.py index dcd78d564..66aeedaab 100644 --- a/dags/tpu_observability/jobset_ttr_pod_delete.py +++ b/dags/tpu_observability/jobset_ttr_pod_delete.py @@ -34,11 +34,23 @@ from dags.tpu_observability.utils import jobset_util as jobset from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils.jobset_util import Workload +from dags.tpu_observability.configs.common import ( + MachineConfigMap, + GCS_CONFIG_PATH, + GCS_JOBSET_CONFIG_PATH, +) +from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout +from dags.common.task_group_with_timeout import TaskGroupWithTimeout + DAG_ID = "jobset_ttr_pod_delete" DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -105,75 +117,100 @@ selector = jobset.generate_node_pool_selector(DAG_ID) jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - node_pool_selector=selector, - ) - - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - node_pool_selector=selector, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - deletion_start_time = jobset.delete_one_random_pod.override( - task_id="delete_random_pod" - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ) - - wait_for_recovery = jobset.wait_for_jobset_recovered.override( - task_id="wait_for_recovery" - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ) - - verify_duration = jobset.verify_recovery_duration.override( - task_id="verify_recovery_duration" - )( - start_time=deletion_start_time, - end_time=wait_for_recovery, - ) - - wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override( - task_id="wait_for_jobset_ttr_to_be_found", - )( - node_pool=cluster_info, - jobset_name=jobset_name, - start_time=deletion_start_time, - ) - - cleanup_workload = jobset.end_workload.override( - task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ).as_teardown( - setups=startup.jobset_start_time - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool" + )( + node_pool=cluster_info, + node_pool_selector=selector, + ) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + node_pool_selector=selector, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + deletion_start_time = jobset.delete_one_random_pod.override( + task_id="delete_random_pod" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ) + + wait_for_recovery = jobset.wait_for_jobset_recovered.override( + task_id="wait_for_recovery" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ) + + verify_duration = jobset.verify_recovery_duration.override( + task_id="verify_recovery_duration" + )( + start_time=deletion_start_time, + end_time=wait_for_recovery, + ) + + wait_for_metric_upload = ( + jobset.wait_for_jobset_ttr_to_be_found.override( + task_id="wait_for_jobset_ttr_to_be_found", + )( + node_pool=cluster_info, + jobset_name=jobset_name, + start_time=deletion_start_time, + ) + ) + + chain( + *startup.tasks, + deletion_start_time, + wait_for_recovery, + verify_duration, + wait_for_metric_upload, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ).as_teardown( + setups=startup.jobset_start_time + ) + + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) + + chain( + cleanup_workload, + cleanup_node_pool, + ) chain( selector, jobset_name, - create_node_pool, - *startup.tasks, - deletion_start_time, - wait_for_recovery, - verify_duration, - wait_for_metric_upload, - cleanup_workload, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/jobset_ttr_rollback.py b/dags/tpu_observability/jobset_ttr_rollback.py index 761c473ff..28a1c11bb 100644 --- a/dags/tpu_observability/jobset_ttr_rollback.py +++ b/dags/tpu_observability/jobset_ttr_rollback.py @@ -34,11 +34,23 @@ from dags.tpu_observability.utils import jobset_util as jobset from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils.jobset_util import Workload +from dags.tpu_observability.configs.common import ( + MachineConfigMap, + GCS_CONFIG_PATH, + GCS_JOBSET_CONFIG_PATH, +) +from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout +from dags.common.task_group_with_timeout import TaskGroupWithTimeout + DAG_ID = "jobset_rollback_ttr" DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -106,71 +118,96 @@ selector = jobset.generate_node_pool_selector(DAG_ID) jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - node_pool_selector=selector, - ) - - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - node_pool_selector=selector, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - rollback_node_pool = node_pool.rollback.override( - task_id="rollback_node_pool" - )(node_pool=cluster_info) - - wait_for_recovery = jobset.wait_for_jobset_recovered.override( - task_id="wait_for_recovery" - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ) - - verify_duration = jobset.verify_recovery_duration.override( - task_id="verify_recovery_duration" - )( - start_time=rollback_node_pool, - end_time=wait_for_recovery, - ) - - wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override( - task_id="wait_for_jobset_ttr_to_be_found", - )( - node_pool=cluster_info, - jobset_name=jobset_name, - start_time=rollback_node_pool, - ) - - cleanup_workload = jobset.end_workload.override( - task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ).as_teardown( - setups=startup.jobset_start_time - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool" + )( + node_pool=cluster_info, + node_pool_selector=selector, + ) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + node_pool_selector=selector, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + rollback_node_pool = node_pool.rollback.override( + task_id="rollback_node_pool" + )(node_pool=cluster_info) + + wait_for_recovery = jobset.wait_for_jobset_recovered.override( + task_id="wait_for_recovery" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ) + + verify_duration = jobset.verify_recovery_duration.override( + task_id="verify_recovery_duration" + )( + start_time=rollback_node_pool, + end_time=wait_for_recovery, + ) + + wait_for_metric_upload = ( + jobset.wait_for_jobset_ttr_to_be_found.override( + task_id="wait_for_jobset_ttr_to_be_found", + )( + node_pool=cluster_info, + jobset_name=jobset_name, + start_time=rollback_node_pool, + ) + ) + + chain( + *startup.tasks, + rollback_node_pool, + wait_for_recovery, + verify_duration, + wait_for_metric_upload, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ).as_teardown( + setups=startup.jobset_start_time + ) + + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) + + chain( + cleanup_workload, + cleanup_node_pool, + ) chain( selector, jobset_name, - create_node_pool, - *startup.tasks, - rollback_node_pool, - wait_for_recovery, - verify_duration, - wait_for_metric_upload, - cleanup_workload, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/jobset_uptime_validation.py b/dags/tpu_observability/jobset_uptime_validation.py index 3cc07f3a2..1024063d0 100644 --- a/dags/tpu_observability/jobset_uptime_validation.py +++ b/dags/tpu_observability/jobset_uptime_validation.py @@ -36,11 +36,18 @@ from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils.jobset_util import Workload from dags.tpu_observability.utils.time_util import TimeUtil +from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout +from dags.common.task_group_with_timeout import TaskGroupWithTimeout + DAG_ID = "jobset_uptime_validation" DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + @task def get_current_time() -> TimeUtil: @@ -117,67 +124,88 @@ def get_current_time() -> TimeUtil: selector = jobset.generate_node_pool_selector(DAG_ID) jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - node_pool_selector=selector, - ) - - startup = jobset.create_jobset_startup_tasks( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - node_pool_selector=selector, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - wait_for_jobset_uptime_data = jobset.wait_for_jobset_uptime_data.override( - task_id="wait_for_jobset_uptime_data" - )( - node_pool=cluster_info, - jobset_name=jobset_name, - jobset_apply_time=startup.jobset_start_time, - ) - - clean_up_workload = jobset.end_workload.override( - task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_config=jobset_config, - jobset_name=jobset_name, - ).as_teardown( - setups=startup.jobset_start_time - ) - - jobset_clear_time = get_current_time.override( - task_id="get_current_time" - )() - - ensure_no_jobset_uptime_data = ( - jobset.ensure_no_jobset_uptime_data.override( - task_id="ensure_no_jobset_uptime_data" - ) - )( - node_pool=cluster_info, - jobset_name=jobset_name, - jobset_clear_time=jobset_clear_time, - # Wait 5 minutes to confirm no data has been detected. - wait_time_seconds=300, - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=cluster_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool" + )( + node_pool=cluster_info, + node_pool_selector=selector, + ) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + startup = jobset.create_jobset_startup_tasks( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + node_pool_selector=selector, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + wait_for_jobset_uptime_data = ( + jobset.wait_for_jobset_uptime_data.override( + task_id="wait_for_jobset_uptime_data" + )( + node_pool=cluster_info, + jobset_name=jobset_name, + jobset_apply_time=startup.jobset_start_time, + ) + ) + + clean_up_workload = jobset.end_workload.override( + task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ).as_teardown( + setups=startup.jobset_start_time + ) + + jobset_clear_time = get_current_time.override( + task_id="get_current_time" + )() + + ensure_no_jobset_uptime_data = ( + jobset.ensure_no_jobset_uptime_data.override( + task_id="ensure_no_jobset_uptime_data" + ) + )( + node_pool=cluster_info, + jobset_name=jobset_name, + jobset_clear_time=jobset_clear_time, + # Wait 5 minutes to confirm no data has been detected. + wait_time_seconds=300, + ) + + chain( + *startup.tasks, + wait_for_jobset_uptime_data, + clean_up_workload, + jobset_clear_time, + ensure_no_jobset_uptime_data, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) chain( selector, jobset_name, - create_node_pool, - *startup.tasks, - wait_for_jobset_uptime_data, - clean_up_workload, - jobset_clear_time, - ensure_no_jobset_uptime_data, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/multi_host_nodepool_rollback_dag.py b/dags/tpu_observability/multi_host_nodepool_rollback_dag.py index b7c7284b7..f3f5e4272 100644 --- a/dags/tpu_observability/multi_host_nodepool_rollback_dag.py +++ b/dags/tpu_observability/multi_host_nodepool_rollback_dag.py @@ -30,6 +30,7 @@ SchedulingHelper, get_dag_timeout, ) +from dags.common.task_group_with_timeout import TaskGroupWithTimeout from dags.tpu_observability.configs.common import ( GCS_CONFIG_PATH, MachineConfigMap, @@ -40,6 +41,10 @@ DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -100,38 +105,58 @@ tpu_topology=config.tpu_topology, ) - create_node_pool = node_pool.create.override(owner=test_owner.QUINN_M)( - node_pool=node_pool_info, - ) - - wait_node_pool_available = node_pool.wait_for_availability( - node_pool=node_pool_info, availability=True - ) - - rollback_node_pool = node_pool.rollback(node_pool=node_pool_info) - - wait_node_pool_unavailable = node_pool.wait_for_availability( - node_pool=node_pool_info, availability=False - ) - - # A successful rollback means the availability will return to True. - # The end of the rollback marks the start the availability, so - # the client side should see the state change, and update the metric. - wait_node_pool_recovered = node_pool.wait_for_availability( - node_pool=node_pool_info, availability=True - ) - - cleanup_node_pool = node_pool.delete.override( - trigger_rule=TriggerRule.ALL_DONE - )(node_pool=node_pool_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool", + owner=test_owner.QUINN_M, + )(node_pool=node_pool_info) + + wait_node_pool_available = node_pool.wait_for_availability.override( + task_id="wait_node_pool_available" + )(node_pool=node_pool_info, availability=True) + + chain(create_node_pool, wait_node_pool_available) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + rollback_node_pool = node_pool.rollback.override( + task_id="rollback_node_pool" + )(node_pool=node_pool_info) + + wait_node_pool_unavailable = node_pool.wait_for_availability.override( + task_id="wait_node_pool_unavailable" + )(node_pool=node_pool_info, availability=False) + + # A successful rollback means the availability will return to True. + # The end of the rollback marks the start the availability, so + # the client side should see the state change, and update the metric. + wait_node_pool_recovered = node_pool.wait_for_availability.override( + task_id="wait_node_pool_recovered" + )(node_pool=node_pool_info, availability=True) + + chain( + rollback_node_pool, + wait_node_pool_unavailable, + wait_node_pool_recovered, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", + trigger_rule=TriggerRule.ALL_DONE, + )(node_pool=node_pool_info) chain( - create_node_pool, - wait_node_pool_available, - rollback_node_pool, - wait_node_pool_unavailable, - wait_node_pool_recovered, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/node_pool_status.py b/dags/tpu_observability/node_pool_status.py index e97aec249..9e12f89c5 100644 --- a/dags/tpu_observability/node_pool_status.py +++ b/dags/tpu_observability/node_pool_status.py @@ -28,6 +28,7 @@ SchedulingHelper, get_dag_timeout, ) +from dags.common.task_group_with_timeout import TaskGroupWithTimeout from dags.tpu_observability.configs.common import ( GCS_CONFIG_PATH, MachineConfigMap, @@ -40,6 +41,10 @@ DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -100,107 +105,119 @@ f"{node_pool_info.node_pool_name}-x" ) - task_id = "create_node_pool" - create_node_pool = node_pool.create.override( - task_id=task_id, - owner=test_owner.YUNA_T, - )( - node_pool=node_pool_info, - ) - - task_id = "wait_for_provisioning" - wait_for_provisioning = node_pool.wait_for_status.override( - task_id=task_id - )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING) - - task_id = "wait_for_running" - wait_for_running = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=node_pool_info, status=node_pool.Status.RUNNING - ) - - task_id = "select_random_node" - select_random_node = node_pool.draw_random_node.override(task_id=task_id)( - node_pool=node_pool_info - ) - - task_id = "delete_node" - delete_node = node_pool.operate_node.override(task_id=task_id)( - node_pool=node_pool_info, - operation=NodeOperationSpec.Delete(), - node_name=select_random_node, - ) - - # TODO: add a check that the node count decreases after deletion. - # kubectl is not available here since this DAG has no workload or jobset. - - task_id = "wait_for_repair" - wait_for_repair = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=node_pool_info, status=node_pool.Status.RECONCILING - ) - - task_id = "wait_for_recovered" - wait_for_recovered = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=node_pool_info, status=node_pool.Status.RUNNING - ) - - task_id = "delete_node_pool" - delete_node_pool = node_pool.delete.override(task_id=task_id)( - node_pool=node_pool_info - ) - - task_id = "wait_for_stopping" - wait_for_stopping = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=node_pool_info, status=node_pool.Status.STOPPING - ) - - task_id = "cleanup_node_pool" - cleanup_node_pool = node_pool.delete.override( - task_id=task_id, trigger_rule=TriggerRule.ALL_DONE - )(node_pool=node_pool_info).as_teardown( - setups=create_node_pool, - ) - - # Intentionally create a node pool with problematic configurations - # to validate that it enters the ERROR state. - task_id = "create_problematic_node_pool_info" - create_problematic_node_pool_info = node_pool.create.override( - task_id=task_id, - owner=test_owner.YUNA_T, - )( - node_pool=problematic_node_pool_info, - # The failure is intentionally ignored because we want to validate - # that the status of the node pool (which fails to be created) is - # "ERROR". - ignore_failure=True, - ) - - task_id = "wait_for_error" - wait_for_error = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=problematic_node_pool_info, status=node_pool.Status.ERROR - ) - - task_id = "cleanup_wrong_node_pool" - cleanup_wrong_node_pool = node_pool.delete.override( - task_id=task_id, trigger_rule=TriggerRule.ALL_DONE - )(node_pool=problematic_node_pool_info).as_teardown( - setups=create_problematic_node_pool_info, - ) - - chain( - create_node_pool, - wait_for_provisioning, - wait_for_running, - select_random_node, - delete_node, - wait_for_repair, - wait_for_recovered, - delete_node_pool, - wait_for_stopping, - cleanup_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + task_id = "create_node_pool" + create_node_pool = node_pool.create.override( + task_id=task_id, + owner=test_owner.YUNA_T, + )( + node_pool=node_pool_info, + ) + + # Intentionally create a node pool with problematic configurations + # to validate that it enters the ERROR state. + task_id = "create_problematic_node_pool_info" + create_problematic_node_pool_info = node_pool.create.override( + task_id=task_id, + owner=test_owner.YUNA_T, + )( + node_pool=problematic_node_pool_info, + # The failure is intentionally ignored because we want to validate + # that the status of the node pool (which fails to be created) is + # "ERROR". + ignore_failure=True, + ) + + task_id = "wait_for_provisioning" + wait_for_provisioning = node_pool.wait_for_status.override( + task_id=task_id + )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING) + + task_id = "wait_for_running" + wait_for_running = node_pool.wait_for_status.override(task_id=task_id)( + node_pool=node_pool_info, status=node_pool.Status.RUNNING + ) + + chain(create_node_pool, wait_for_provisioning, wait_for_running) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + task_id = "select_random_node" + select_random_node = node_pool.draw_random_node.override( + task_id=task_id + )(node_pool=node_pool_info) + + task_id = "delete_node" + delete_node = node_pool.operate_node.override(task_id=task_id)( + node_pool=node_pool_info, + operation=NodeOperationSpec.Delete(), + node_name=select_random_node, + ) + + task_id = "wait_for_repair" + wait_for_repair = node_pool.wait_for_status.override(task_id=task_id)( + node_pool=node_pool_info, status=node_pool.Status.RECONCILING + ) + + task_id = "wait_for_recovered" + wait_for_recovered = node_pool.wait_for_status.override( + task_id=task_id + )(node_pool=node_pool_info, status=node_pool.Status.RUNNING) + + task_id = "delete_node_pool" + delete_node_pool = node_pool.delete.override(task_id=task_id)( + node_pool=node_pool_info + ) + + task_id = "wait_for_stopping" + wait_for_stopping = node_pool.wait_for_status.override(task_id=task_id)( + node_pool=node_pool_info, status=node_pool.Status.STOPPING + ) + + task_id = "wait_for_error" + validate_problematic_node_pool_enter_error_state = ( + node_pool.wait_for_status.override(task_id=task_id)( + node_pool=problematic_node_pool_info, + status=node_pool.Status.ERROR, + ) + ) + + chain( + select_random_node, + delete_node, + wait_for_repair, + wait_for_recovered, + delete_node_pool, + wait_for_stopping, + ) + + chain( + create_problematic_node_pool_info, + validate_problematic_node_pool_enter_error_state, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + task_id = "cleanup_node_pool" + cleanup_node_pool = node_pool.delete.override( + task_id=task_id, trigger_rule=TriggerRule.ALL_DONE + )(node_pool=node_pool_info) + + task_id = "cleanup_wrong_node_pool" + cleanup_wrong_node_pool = node_pool.delete.override( + task_id=task_id, trigger_rule=TriggerRule.ALL_DONE + )(node_pool=problematic_node_pool_info) chain( - create_problematic_node_pool_info, - wait_for_error, - cleanup_wrong_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/node_pool_ttr_disk_size.py b/dags/tpu_observability/node_pool_ttr_disk_size.py index c44604bcf..61caf5e8e 100644 --- a/dags/tpu_observability/node_pool_ttr_disk_size.py +++ b/dags/tpu_observability/node_pool_ttr_disk_size.py @@ -27,6 +27,7 @@ SchedulingHelper, get_dag_timeout, ) +from dags.common.task_group_with_timeout import TaskGroupWithTimeout from dags.tpu_observability.configs.common import ( GCS_CONFIG_PATH, MachineConfigMap, @@ -38,6 +39,10 @@ SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) _DISK_SIZE_INCREMENT = 50 +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + with models.DAG( dag_id=DAG_ID, start_date=datetime.datetime(2025, 6, 26), @@ -89,45 +94,58 @@ tpu_topology=config.tpu_topology, ) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=node_pool_info - ) - - wait_for_provisioning = node_pool.wait_for_status.override( - task_id="wait_for_provisioning" - )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING) - - wait_for_running = node_pool.wait_for_status.override( - task_id="wait_for_running" - )(node_pool=node_pool_info, status=node_pool.Status.RUNNING) - - update_start_time = node_pool.update.override(task_id="update_node_pool")( - node_pool=node_pool_info, - spec=node_pool.NodePoolUpdateSpec.DiskSize( - delta=_DISK_SIZE_INCREMENT - ), - ) - - wait_for_recovered = node_pool.wait_for_status.override( - task_id="wait_for_recovered" - )(node_pool=node_pool_info, status=node_pool.Status.RUNNING) - - wait_for_ttr = node_pool.wait_for_ttr( - node_pool=node_pool_info, operation_start_time=update_start_time - ) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=node_pool_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool" + )(node_pool=node_pool_info) + + wait_for_provisioning = node_pool.wait_for_status.override( + task_id="wait_for_provisioning" + )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING) + + wait_for_running = node_pool.wait_for_status.override( + task_id="wait_for_running" + )(node_pool=node_pool_info, status=node_pool.Status.RUNNING) + + chain(create_node_pool, wait_for_provisioning, wait_for_running) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + update_start_time = node_pool.update.override( + task_id="update_node_pool" + )( + node_pool=node_pool_info, + spec=node_pool.NodePoolUpdateSpec.DiskSize( + delta=_DISK_SIZE_INCREMENT + ), + ) + + wait_for_recovered = node_pool.wait_for_status.override( + task_id="wait_for_recovered" + )(node_pool=node_pool_info, status=node_pool.Status.RUNNING) + + wait_for_ttr = node_pool.wait_for_ttr( + node_pool=node_pool_info, operation_start_time=update_start_time + ) + + chain(update_start_time, wait_for_recovered, wait_for_ttr) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=node_pool_info) chain( - create_node_pool, - wait_for_provisioning, - wait_for_running, - update_start_time, - wait_for_recovered, - wait_for_ttr, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/node_pool_ttr_update_label.py b/dags/tpu_observability/node_pool_ttr_update_label.py index d916d0d2a..3faac6c6d 100644 --- a/dags/tpu_observability/node_pool_ttr_update_label.py +++ b/dags/tpu_observability/node_pool_ttr_update_label.py @@ -27,6 +27,7 @@ SchedulingHelper, get_dag_timeout, ) +from dags.common.task_group_with_timeout import TaskGroupWithTimeout from dags.tpu_observability.configs.common import ( GCS_CONFIG_PATH, MachineConfigMap, @@ -37,6 +38,10 @@ DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + with models.DAG( dag_id=DAG_ID, start_date=datetime.datetime(2025, 9, 30), @@ -87,50 +92,62 @@ tpu_topology=config.tpu_topology, ) - task_id = "create_node_pool" - create_node_pool = node_pool.create.override(task_id=task_id)( - node_pool=node_pool_info, - ) - - task_id = "wait_for_provisioning" - wait_for_provisioning = node_pool.wait_for_status.override( - task_id=task_id - )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING) - - task_id = "wait_for_running" - wait_for_running = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=node_pool_info, status=node_pool.Status.RUNNING - ) - - task_id = "update_node_pool_label" - update_node_pool_label = node_pool.update.override(task_id=task_id)( - node_pool=node_pool_info, - spec=node_pool.NodePoolUpdateSpec.Label(delta=labels_to_update), - ) - - task_id = "wait_for_recovered" - wait_for_recovered = node_pool.wait_for_status.override(task_id=task_id)( - node_pool=node_pool_info, status=node_pool.Status.RUNNING - ) - - task_id = "wait_for_ttr" - wait_for_ttr = node_pool.wait_for_ttr.override(task_id=task_id)( - node_pool=node_pool_info, operation_start_time=update_node_pool_label - ) - - task_id = "cleanup_node_pool" - cleanup_node_pool = node_pool.delete.override( - task_id=task_id, trigger_rule=TriggerRule.ALL_DONE - )(node_pool=node_pool_info).as_teardown( - setups=create_node_pool, - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + task_id = "create_node_pool" + create_node_pool = node_pool.create.override(task_id=task_id)( + node_pool=node_pool_info, + ) + + task_id = "wait_for_provisioning" + wait_for_provisioning = node_pool.wait_for_status.override( + task_id=task_id + )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING) + + task_id = "wait_for_running" + wait_for_running = node_pool.wait_for_status.override(task_id=task_id)( + node_pool=node_pool_info, status=node_pool.Status.RUNNING + ) + + chain(create_node_pool, wait_for_provisioning, wait_for_running) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + task_id = "update_node_pool_label" + update_node_pool_label = node_pool.update.override(task_id=task_id)( + node_pool=node_pool_info, + spec=node_pool.NodePoolUpdateSpec.Label(delta=labels_to_update), + ) + + task_id = "wait_for_recovered" + wait_for_recovered = node_pool.wait_for_status.override( + task_id=task_id + )(node_pool=node_pool_info, status=node_pool.Status.RUNNING) + + task_id = "wait_for_ttr" + wait_for_ttr = node_pool.wait_for_ttr.override(task_id=task_id)( + node_pool=node_pool_info, + operation_start_time=update_node_pool_label, + ) + + chain(update_node_pool_label, wait_for_recovered, wait_for_ttr) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + task_id = "cleanup_node_pool" + cleanup_node_pool = node_pool.delete.override( + task_id=task_id, trigger_rule=TriggerRule.ALL_DONE + )(node_pool=node_pool_info) chain( - create_node_pool, - wait_for_provisioning, - wait_for_running, - update_node_pool_label, - wait_for_recovered, - wait_for_ttr, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/update_node_pool_label.py b/dags/tpu_observability/update_node_pool_label.py index cc9b93cad..e2a32de98 100644 --- a/dags/tpu_observability/update_node_pool_label.py +++ b/dags/tpu_observability/update_node_pool_label.py @@ -29,6 +29,7 @@ SchedulingHelper, get_dag_timeout, ) +from dags.common.task_group_with_timeout import TaskGroupWithTimeout from dags.tpu_observability.configs.common import ( GCS_CONFIG_PATH, MachineConfigMap, @@ -39,6 +40,10 @@ DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) +PRE_TEST_TIMEOUT = datetime.timedelta(minutes=10) +POST_TEST_TIMEOUT = datetime.timedelta(minutes=10) +TEST_TIMEOUT = DAGRUN_TIMEOUT - PRE_TEST_TIMEOUT - POST_TEST_TIMEOUT + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -88,43 +93,59 @@ tpu_topology=config.tpu_topology, ) - create_node_pool = node_pool.create.override( - task_id="create_node_pool", - owner=test_owner.YUNA_T, - )( - node_pool=node_pool_info, - ) - - wait_for_availability = node_pool.wait_for_availability.override( - task_id="wait_for_initial_availability" - )(node_pool=node_pool_info, availability=True) - - update_node_pool_label = node_pool.update.override( - task_id="update_node_pool_label" - )( - node_pool=node_pool_info, - spec=node_pool.NodePoolUpdateSpec.Label(delta=labels_to_update), - ) - - wait_for_unavailable = node_pool.wait_for_availability.override( - task_id="wait_for_unavailability_after_update" - )(node_pool=node_pool_info, availability=False) - - wait_node_pool_recovered = node_pool.wait_for_availability.override( - task_id="wait_for_recovery" - )(node_pool=node_pool_info, availability=True) - - cleanup_node_pool = node_pool.delete.override( - task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE - )(node_pool=node_pool_info).as_teardown( - setups=[create_node_pool], - ) + with TaskGroupWithTimeout( + group_id="pre_test", + timeout=PRE_TEST_TIMEOUT, + ) as pre_test: + create_node_pool = node_pool.create.override( + task_id="create_node_pool", + owner=test_owner.YUNA_T, + )( + node_pool=node_pool_info, + ) + + wait_for_availability = node_pool.wait_for_availability.override( + task_id="wait_for_initial_availability" + )(node_pool=node_pool_info, availability=True) + + chain(create_node_pool, wait_for_availability) + + with TaskGroupWithTimeout( + group_id="test", + timeout=TEST_TIMEOUT, + ) as test: + update_node_pool_label = node_pool.update.override( + task_id="update_node_pool_label" + )( + node_pool=node_pool_info, + spec=node_pool.NodePoolUpdateSpec.Label(delta=labels_to_update), + ) + + wait_for_unavailable = node_pool.wait_for_availability.override( + task_id="wait_for_unavailability_after_update" + )(node_pool=node_pool_info, availability=False) + + wait_node_pool_recovered = node_pool.wait_for_availability.override( + task_id="wait_for_recovery" + )(node_pool=node_pool_info, availability=True) + + chain( + update_node_pool_label, + wait_for_unavailable, + wait_node_pool_recovered, + ) + + with TaskGroupWithTimeout( + group_id="post_test", + timeout=POST_TEST_TIMEOUT, + is_teardown=True, + ) as post_test: + cleanup_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE + )(node_pool=node_pool_info) chain( - create_node_pool, - wait_for_availability, - update_node_pool_label, - wait_for_unavailable, - wait_node_pool_recovered, - cleanup_node_pool, + pre_test, + test, + post_test, ) diff --git a/dags/tpu_observability/utils/node_pool_util.py b/dags/tpu_observability/utils/node_pool_util.py index 7ab8708b1..7cb1521ce 100644 --- a/dags/tpu_observability/utils/node_pool_util.py +++ b/dags/tpu_observability/utils/node_pool_util.py @@ -553,7 +553,7 @@ def _query_status_metric(node_pool: Info) -> Status: return Status.from_str(latest_status) -@task.sensor(poke_interval=60, timeout=600, mode="poke") +@task.sensor(poke_interval=60, timeout=600, mode="poke", retries=0) def wait_for_status( node_pool: Info, status: Status, @@ -608,7 +608,7 @@ def rollback(node_pool: Info) -> None: return TimeUtil.from_datetime(current_time_utc) -@task.sensor(poke_interval=30, timeout=1200, mode="poke") +@task.sensor(poke_interval=30, timeout=1200, mode="poke", retries=0) def wait_for_availability( node_pool: Info, availability: bool, @@ -684,7 +684,7 @@ def wait_for_availability( return availability == state -@task.sensor(poke_interval=30, timeout=3600, mode="poke") +@task.sensor(poke_interval=30, timeout=3600, mode="poke", retries=0) def wait_for_ttr( node_pool: Info, operation_start_time: TimeUtil, 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."