feat: Apply TaskGroupWithTimeout for TPU Observability DAGs - part 1 - #262
Merged
alfredyu-cienet merged 6 commits intoJul 1, 2026
Merged
Conversation
yuna-tzeng
requested review from
alfredyu-cienet,
andrewyct,
severus-ho and
xibinliu
as code owners
May 21, 2026 09:29
Collaborator
|
first scan OK |
alfredyu-cienet
changed the title
TaskGroupWithTimeout for TPU Observability DAGs - part 1
Comment on lines
+172
to
+184
| # 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, | ||
| ) |
Collaborator
There was a problem hiding this comment.
I think this should be in pre-test
| ) | ||
|
|
||
| task_id = "wait_for_error" | ||
| wait_for_error = node_pool.wait_for_status.override(task_id=task_id)( |
Collaborator
There was a problem hiding this comment.
Suggested change
| wait_for_error = node_pool.wait_for_status.override(task_id=task_id)( | |
| validate_problematic_node_pool_enter_error_state = node_pool.wait_for_status.override(task_id=task_id)( |
Collaborator
|
ditto #264 (comment) |
alfredyu-cienet
force-pushed
the
dev
branch
4 times, most recently
from
June 8, 2026 08:58
4af284a to
a9fa8fb
Compare
alfredyu-cienet
force-pushed
the
dev
branch
3 times, most recently
from
June 12, 2026 07:31
cd8b06c to
d4c530a
Compare
yuna-tzeng
force-pushed
the
tpu-obs/user/yuna/taskgroup_timeout
branch
2 times, most recently
from
June 12, 2026 09:42
e784d2e to
71043b8
Compare
alfredyu-cienet
force-pushed
the
dev
branch
2 times, most recently
from
June 18, 2026 00:43
00a4856 to
fdcbae9
Compare
yuna-tzeng
force-pushed
the
tpu-obs/user/yuna/taskgroup_timeout
branch
from
June 26, 2026 03:38
71043b8 to
e683b6e
Compare
…nrelated GKE DAGs
Integrates `TaskGroupWithTimeout` directly into all TPU Observability
DAGs that do not depend on Kubernetes JobSets. This provides strict,
shared deadline timeout enforcement across distinct stages of GKE node
pool validation.
Specifically:
1. Configures stage-specific timeouts for all modified DAGs:
- `PRE_TEST_TIMEOUT`: 10 minutes (for GKE provisioning and availability setup)
- `POST_TEST_TIMEOUT`: 10 minutes (for GKE teardown and cleanup)
- `TEST_TIMEOUT`: Evaluates dynamically as `DAGRUN_TIMEOUT - 20 mins`
2. Restructures the following 5 DAGs into three separate `TaskGroupWithTimeout` blocks:
- `dags/tpu_observability/node_pool_status.py`
- `dags/tpu_observability/node_pool_ttr_disk_size.py`
- `dags/tpu_observability/node_pool_ttr_update_label.py`
- `dags/tpu_observability/multi_host_nodepool_rollback_dag.py`
- `dags/tpu_observability/update_node_pool_label.py`
3. Refactors GKE node pool operations into:
- `pre_test`: Handles node pool creation and waiting for provisioning/running.
- `test`: Executes the core test mutation and verification flow.
- `post_test`: Manages cleanup/teardown of GKE node pools (using `is_teardown=True`
to guarantee execution even if upstream stages timeout or fail).
…upWithTimeout Refactors five TPU observability DAGs to use TaskGroupWithTimeout for proper timeout management during pre-test setup, test execution, and post-test teardown. This aligns their implementation pattern with node_pool_status.py. Specifically, this change: - Imports TaskGroupWithTimeout in each DAG. - Defines PRE_TEST_TIMEOUT, TEST_TIMEOUT, and POST_TEST_TIMEOUT at the module level. - Groups tasks into pre_test, test, and post_test TaskGroupWithTimeout blocks. - Chains tasks sequentially within each group and updates the overall DAG chain. Modified DAGs: - dags/tpu_observability/jobset_ttr_drain_restart.py - dags/tpu_observability/jobset_ttr_node_pool_resize.py - dags/tpu_observability/jobset_ttr_pod_delete.py - dags/tpu_observability/jobset_ttr_rollback.py - dags/tpu_observability/jobset_uptime_validation.py
yuna-tzeng
force-pushed
the
tpu-obs/user/yuna/taskgroup_timeout
branch
from
June 29, 2026 03:37
5f57ff6 to
b70dc15
Compare
alfredyu-cienet
approved these changes
Jul 1, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR integrates TaskGroupWithTimeout directly into all GKE Node Pool validation DAGs that are unrelated to JobSets/workloads. This ensures strict stage-specific timeout enforcement across different phases (Pre-test provisioning: 10 mins, Post-test cleanup: 10 mins, and the remaining budget allocated to the main test case).
DAGs Refactored with TaskGroupWithTimeout
The following 5 GKE Node Pool status and label validation DAGs have been successfully refactored into separate pre_test, test, and post_test blocks under TaskGroupWithTimeout:
Deferred DAGs (Unchanged)
Because TaskGroupWithTimeout does not support nested TaskGroups, GKE DAGs that involve JAX workloads and rely on jobset.create_jobset_startup_group() cannot be refactored yet. The helper function internally instantiates a standard TaskGroup (jobset_startup_and_prepare), which triggers a parsing-time AirflowFailException when nested inside a parent TaskGroupWithTimeout.
As a result, the following JobSet and workload-related DAGs are temporarily left unchanged using original TaskGroups, to be refactored once nested TaskGroups are natively supported:
Key Changes & Mechanisms
Timeout Budgets:
Checklist
Before submitting this PR, please make sure (put X in square brackets):