Adds a new DAG to test "Jobset Ready Healthiness" metric by QuinnMMcGarry · Pull Request #156 · CIeNET-International/ml-auto-solutions-3

QuinnMMcGarry · 2025-12-24T01:10:29Z

Description

This change adds a new DAG that automatically checks if the correct number of replicas appear for a jobset's "Ready" status.

Tests

DAG Name: jobset_healthiness_ready
Airflow test runs: https://b371d513bb74418681c3aab7b27bbf12-dot-us-east1.composer.googleusercontent.com/dags/jobset_healthiness_ready/grid?tab=logs&dag_run_id=manual__2025-12-23T20%3A36%3A17.106112%2B00%3A00&task_id=v6e.cleanup_node_pool__1

Airflow / Composer Environment

GCP Project: cienet-cmcs
Composer Environment: tpu-obs-testing
Composer Version: 2.13.1

Required Variables

Cluster Information (This DAG requires an existing cluster)
- PROJECT_ID: cienet-cmcs
- CLUSTER_NAME: tpu-observability-automation-prod
- LOCATION: us-central1
- NODE_LOCATIONS: us-central1-b
Node Pool / Machine Configuration
- NUM_NODE_POOL: 2
- NODE_POOL_NAMES: jobset-healthiness-ready-v6e / jobset-healthiness-ready-v6e-2
- NUM_NODES: 4
- MACHINE_TYPE: ct6e-standard-4t
- TPU_TOPOLOGY: 4x4

Checklist

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run one-shot tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed.

alfredyu-cienet · 2025-12-24T08:26:04Z

+
+  for machine in MachineConfigMap:
+    config = machine.value
+    cluster_info = node_pool.Info(


new configuration way: PR-1097

I'll rebase the working branch later (Friday or next Monday)

alfredyu-cienet · 2025-12-24T10:21:04Z

  Each workload is a JSON-escaped string, ready to be used as a shell argument.
  """

+  READY_TPU = json.dumps(


From the script context, it appears that this workload script keeps the jobset in “Ready” for ~1000s, while JAX_TPU_BENCHMARK is designed to fully utilize the TPU.

If that's the case, can we rename the workload script to something more descriptive?

alfredyu-cienet · 2025-12-24T10:24:32Z

+# Can be moved to jobset_util.py
+@task.sensor(poke_interval=30, timeout=900, mode="reschedule")
+def validate_replica_number(


it seems wait_for_jobset_status_occurrence is not used, perhaps we can modify and use that one

alfredyu-cienet · 2025-12-24T10:28:02Z

+with models.DAG(  # pylint: disable=unexpected-keyword-arg
+    dag_id="jobset_healthiness_ready",
+    start_date=datetime.datetime(2025, 8, 10),
+    schedule="00 03 * * *",


Prevent scheduling on non-PROD Composer environments.

Schedule near—but not overlapping with—our other DAGs.

For 2, do we/should we have a task to create a schedule for DAGs we which are running or in development?

alfredyu-cienet · 2025-12-29T01:20:49Z

code base outdated

clue#1: https://github.com/CIeNET-International/ml-auto-solutions-3/pull/156/commits
clue#2:

QuinnMMcGarry · 2025-12-29T19:28:55Z

code base outdated

clue#1: https://github.com/CIeNET-International/ml-auto-solutions-3/pull/156/commits clue#2:

I used git rebase --onto origin/tpu-obs/dev 7da7f1803bc5323c0d5fad35d0995c05daf46a3c (my original PR hash), resolved the conflicts, then saved my changes and force pushed them

alfredyu-cienet · 2025-12-30T01:49:19Z

This has happened too many times, and I’ve already outlined multiple ways to double-check. This is critical—you are required to double-check using the methods already provided (PR "Commits" page, git graph, git CLI). Do not submit another revision until you have verified your branch end-to-end.

By “correct,” I mean:

the base codebase/branch is correct
EACH commit is correct

Must perform these checks BEFORE and AFTER submitting any new revision.
Must perform these checks BEFORE and AFTER submitting any new revision.
Must perform these checks BEFORE and AFTER submitting any new revision.

git fetch origin
git checkout tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready
git reset --hard origin/tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready
git log tpu-obs/dev..HEAD --stat

QuinnMMcGarry · 2025-12-30T16:09:46Z

This has happened too many times, and I’ve already outlined multiple ways to double-check. This is critical—you are required to double-check using the methods already provided (PR "Commits" page, git graph, git CLI). Do not submit another revision until you have verified your branch end-to-end.

By “correct,” I mean:

the base codebase/branch is correct

EACH commit is correct

Must perform these checks BEFORE and AFTER submitting any new revision.

Must perform these checks BEFORE and AFTER submitting any new revision.
Must perform these checks BEFORE and AFTER submitting any new revision.
git fetch origin
git checkout tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready
git reset --hard origin/tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready
git log tpu-obs/dev..HEAD --stat

You have told me when my commits are incorrect and how I should be committing them, but not necessarily why they are incorrect. In this case I can see there is a commit with an incorrect author and date, I should have seen that it was out of place and tried to fix that. However, it was after rebasing using the rebase --onto method you showed me, so I am not sure how it occurred.

Also, for example, what about my current commit? It is rebased onto the correct branch, and I can see that the code changes are correct. The commit itself looks fine to me, but it has removed the original workflow commit from the PR itself. Is that okay? I am not sure what caused it to happen.

For that matter, I am also not extremely familiar with git tools, so to me it has made more sense to be able to submit the actual code to be reviewed as soon as it is ready. During which time I can figure out the git issues I am having. I am not ignoring the code changes made in the /dev branch.

alfredyu-cienet · 2025-12-31T09:21:49Z

it was after rebasing using the rebase --onto method you showed me, so I am not sure how it occurred.

It's hard to diagnose how it happened without reproducible details.

A plain, non-interactive rebase (including rebase --onto) typically won’t introduce an unrelated commit into the middle of the branch. So seeing the commit with the [cienet-private] prefix sitting in the middle strongly suggests there were additional history-rewriting steps beyond the rebase --onto itself (for example: interactive rebase, cherry-pick, merge, or amend).

Also, for example, what about my current commit? It is rebased onto the correct branch, and I can see that the code changes are correct. The commit itself looks fine to me, but it has removed the original workflow commit from the PR itself. Is that okay?

Squashing can be OK depending on the goal. For now, I’d suggest avoiding squash in your PRs while you’re still ramping up on Git — keeping the intermediate commits gives you a clear trail to debug metadata/history issues. If you squash every time, you lose that evidence and it becomes hard to learn which operation caused the problem. You can squash later once the workflow is stable.

For that matter, I am also not extremely familiar with git tools, so to me it has made more sense to be able to submit the actual code to be reviewed as soon as it is ready

I understand the intent, but in a collaborative repo the PR needs to be in a reviewable state before it’s submitted. It’s not just about the code diff — branches, commits, and metadata also need to be correct so we don’t disrupt others or make debugging harder. Review bandwidth is limited.

It’s totally fair if you’re still getting comfortable with Git. What helps is performing the lightweight checks we discussed more frequently around Git operations. Put differently: when we have no clue in production, we add logs—the same principle applies here. Please instrument your workflow so it’s observable.

alfredyu-cienet · 2026-01-16T09:02:11Z

Before I left about two weeks ago, I repointed the old working branch to Google’s master branch in case a production hotfix was needed.

Now that I’m back and we have a new working branch, please rebase your branch correctly—verify the base branch and merge target, and ensure the commit history and metadata are clean.

You may encounter failures depending on which commands you use—there is no one-size-fits-all command. The key is to maintain sufficient traceability so you and others can troubleshoot and fix issues properly.

Please also note that we review each PR with its associated metadata as a whole to keep the process efficient and avoid back-and-forth or duplicate effort. Authors are expected to submit complete, self-contained PRs; partial reviews shift the author’s work onto reviewers and fragment discussion. End-to-end reviews also preserve a single source of truth—feedback stays in one place rather than splitting into conflicting threads.

QuinnMMcGarry · 2026-01-18T20:05:31Z

Before I left about two weeks ago, I repointed the old working branch to Google’s master branch in case a production hotfix was needed.

Now that I’m back and we have a new working branch, please rebase your branch correctly—verify the base branch and merge target, and ensure the commit history and metadata are clean.

You may encounter failures depending on which commands you use—there is no one-size-fits-all command. The key is to maintain sufficient traceability so you and others can troubleshoot and fix issues properly.

Please also note that we review each PR with its associated metadata as a whole to keep the process efficient and avoid back-and-forth or duplicate effort. Authors are expected to submit complete, self-contained PRs; partial reviews shift the author’s work onto reviewers and fragment discussion. End-to-end reviews also preserve a single source of truth—feedback stays in one place rather than splitting into conflicting threads.

I rebased onto origin/dev, and have ensured Git CLI, GitGraph, and GitHub UI, all show a total of 4 commits, in the correct order, with the Description/Author/Hash. And all Git checks pass.

Does the Git committing look fine as it is?

alfredyu-cienet · 2026-01-19T09:56:13Z

+      ) as create_node_pool:
+        create_first_node_pool = node_pool.create.override(
+            task_id="node_pool_1",
+            retries=2,


please reconsider all the "retries" setups

based on both consistency (compare with the ones we automated) and necessity

(discuss here in this thread)

Realistically, any given run probably does not need any retries. But, on a larger scale, it would not be surprising to see an occasional failure happen due to some hardware or software issue outside our control which could have been fixed by a retry. That being said, if 1 retry does not work, the chance that 2 retries will work is even lower, and probably unneeded

so it looks like this "retries" config does not help much on stability,

on the other hand, our work is more like an end to end test, one of the purposes is to find possible regression issues and/or HW/SW outage, so let's not do "retries" by default

Understood. Retries removed

- Adds an entire new DAG to test this metric automatically - Adds the "validate_replica_number" function - Adds a "READY_TPU" Jobset in jobset_util.py -Fix automatic cluster naming script -Created "validate_jobset_replica_number" funtion in jobset_util.py -Changed workload name to "IDLE_READY_TPU_20M" -Changed schedule time and added statement to only run on PROD environment Changed to new nodepool creation format Remove unused DAG flow Fix format/linting Linting fix Delete dags/dags/tpu_observability directory Delete dags/airflow_monitoring.py

-selector task -new global variable syntax -chain instead of ">>"

…et_util.py file Refactor for new DAG formatting

new jobset storage method. Since "jobset_confifg.replicas" cannot be accessed in DAG itself.

alfredyu-cienet · 2026-05-13T02:06:41Z

covered by #256

QuinnMMcGarry requested review from alfredyu-cienet, andrewyct, severus-ho and xibinliu as code owners December 24, 2025 01:10

alfredyu-cienet reviewed Dec 24, 2025

View reviewed changes

alfredyu-cienet added the P2 label Dec 24, 2025

alfredyu-cienet force-pushed the tpu-obs/dev branch from 0535667 to fdcca93 Compare December 26, 2025 08:14

QuinnMMcGarry force-pushed the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch 2 times, most recently from aea57e1 to 7da7f18 Compare December 28, 2025 23:23

QuinnMMcGarry force-pushed the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch 2 times, most recently from acc6fdb to fba8836 Compare December 29, 2025 19:26

QuinnMMcGarry force-pushed the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch from fba8836 to caf272f Compare December 30, 2025 15:01

alfredyu-cienet force-pushed the tpu-obs/dev branch from fdcca93 to 846aa8d Compare January 5, 2026 07:09

QuinnMMcGarry changed the base branch from tpu-obs/dev to dev January 14, 2026 15:58

alfredyu-cienet added P1 and removed P2 labels Jan 16, 2026

QuinnMMcGarry force-pushed the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch 2 times, most recently from bb2b8ee to 0973a13 Compare January 18, 2026 19:38

QuinnMMcGarry force-pushed the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch from e1962d2 to 363b1b9 Compare January 19, 2026 03:02

alfredyu-cienet reviewed Jan 19, 2026

View reviewed changes

alfredyu-cienet force-pushed the dev branch 2 times, most recently from 1c93601 to 1ef36b1 Compare January 22, 2026 03:24

alfredyu-cienet force-pushed the dev branch from 1ef36b1 to 0f0e4e2 Compare January 27, 2026 08:47

alfredyu-cienet force-pushed the dev branch 3 times, most recently from 65e3f42 to cb3b7b7 Compare March 13, 2026 10:08

alfredyu-cienet force-pushed the dev branch from cb3b7b7 to c0fcb71 Compare March 17, 2026 07:07

alfredyu-cienet force-pushed the dev branch from c0fcb71 to 01faecf Compare March 30, 2026 03:12

alfredyu-cienet force-pushed the dev branch from 01faecf to 41ab6ab Compare April 8, 2026 04:53

alfredyu-cienet force-pushed the dev branch 3 times, most recently from d5e60d3 to 1ba128e Compare April 22, 2026 08:27

[cienet-private] Enable workflow checks for the primary working branches

0fe6cba

alfredyu-cienet force-pushed the dev branch from 1ba128e to 0fe6cba Compare April 22, 2026 08:44

QuinnMMcGarry added 7 commits April 27, 2026 09:36

Add "IDLE_READY_TPU_20M" workload to jobset_util.py

bf27f08

Removed "retries" from nodepool creation and deletion

fc8c6d8

Refactor to use new

fae8899

-selector task -new global variable syntax -chain instead of ">>"

Added "jobset_rollback_ttr" to REGISTERED_DAGS in scheduling_helper.py

20f3952

Reintroduce the "validate_jobset_replica_number" function in the jobs…

9d16301

…et_util.py file Refactor for new DAG formatting

Change the "validate_jobset_replica_number" function to work using the

522f718

new jobset storage method. Since "jobset_confifg.replicas" cannot be accessed in DAG itself.

chiajunglien force-pushed the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch from 533eb9b to 522f718 Compare April 27, 2026 01:38

chiajunglien added 2 commits April 27, 2026 11:29

Refactor

5fadcd4

fix

4032b8e

alfredyu-cienet force-pushed the dev branch 4 times, most recently from b81c30b to a12e062 Compare May 5, 2026 01:16

alfredyu-cienet force-pushed the dev branch 2 times, most recently from badfd13 to 721afcd Compare May 12, 2026 23:33

alfredyu-cienet closed this May 13, 2026

alfredyu-cienet deleted the tpu-obs/user/QuinnMMcGarry/jobset-healthiness-ready branch May 13, 2026 02:06

Uh oh!

Conversation

QuinnMMcGarry commented Dec 24, 2025

Description

Tests

Airflow / Composer Environment

Required Variables

Checklist

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alfredyu-cienet commented Dec 29, 2025

Uh oh!

QuinnMMcGarry commented Dec 29, 2025

Uh oh!

alfredyu-cienet commented Dec 30, 2025

Must perform these checks BEFORE and AFTER submitting any new revision. Must perform these checks BEFORE and AFTER submitting any new revision. Must perform these checks BEFORE and AFTER submitting any new revision.

Uh oh!

QuinnMMcGarry commented Dec 30, 2025

Must perform these checks BEFORE and AFTER submitting any new revision.

Uh oh!

alfredyu-cienet commented Dec 31, 2025

Uh oh!

alfredyu-cienet commented Jan 16, 2026

Uh oh!

QuinnMMcGarry commented Jan 18, 2026

Uh oh!

alfredyu-cienet Jan 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alfredyu-cienet commented May 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Must perform these checks BEFORE and AFTER submitting any new revision.
Must perform these checks BEFORE and AFTER submitting any new revision.
Must perform these checks BEFORE and AFTER submitting any new revision.

alfredyu-cienet Jan 19, 2026 •

edited

Loading