Feat: add jobset_healthiness_validation DAG by chiajunglien · Pull Request #256 · CIeNET-International/ml-auto-solutions-3

chiajunglien · 2026-05-06T08:47:37Z

Description

The primary goal of this PR is to implement the jobset_healthiness_validation DAG. This test provides end-to-end validation of the JobSet operator's monitoring metrics, ensuring that Ready, Active, Suspended, Succeeded, Specified, and Failed statuses are correctly reported in Cloud Monitoring.

Technical Implementation

Lifecycle State Machine Validation:
Startup: Validates Ready, Active, and Specified counts match the initial configuration.
Maintenance: Uses kubectl patch to suspend and resume the JobSet, validating the Suspended metric.
Fault Injection: Injects a failure using a Python script with sys.exit(1) to validate the Failed metric.
Success Validation: Deploys a successful workload to validate the Succeeded metric.
Immutability Handling: Addresses Kubernetes JobSet immutability by implementing a "Delete-and-Recreate" pattern within the failure_test and success_test TaskGroups to allow specification updates

Airflow/Composer

GCP Composer name: emma-test (under GCP project: cienet-cmcs)
GCP Composer version: 2.13.1
GCP Airflow version: 2.10.5
Test results: https://a049b7befc7c4d978535dca36b91a080-dot-us-central1.composer.googleusercontent.com/dags/jobset_healthiness_validation/grid

Required Variables

Cluster Information (This DAG requires an existing cluster)
- PROJECT_ID: The GCP project ID where the cluster resides. (Default to cienet-cmcs)
- CLUSTER_NAME: The name of the target GKE cluster. (Default to tpu-observability-automation-dev)
- LOCATION: The region of the GKE cluster. (Default to us-central1)
Node Pool Configurations
- NODE_POOL_NAME: The base name for the new node pool. (Default to jobset-healthiness-validate-v6e)
- NODE_LOCATIONS: The zone for the nodes in the normal test path. (Default to us-central1-b)
- NUM_NODES: The number of nodes to create in the pool. (Default to 4)
- MACHINE_TYPE: The machine type for the GKE nodes. (Default to ct6e-standard-4t)
- TPU_TOPOLOGY: The TPU topology for the node pool. (Default to 4x4)

Before submitting this PR, please make sure (put X in square brackets):

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 · 2026-05-07T04:09:41Z

  jobset_start_time: XComArg


+class ReplicatedJobStatus(enum.Enum):


can we define it as jobset healthiness, and its value can be the metric name directly

explain "healthiness" detail (by "replica" metrics) only in the doc string

alfredyu-cienet · 2026-05-07T04:10:50Z

+      This DAG automates the process of creating node-pools, ensuring the
+      correct number of "Suspended" replicas appear, then launching a jobset on
+      multiple replicas to ensure the correct number begin running.


please name the metics to be examined here in this section

alfredyu-cienet · 2026-05-07T04:12:11Z

+      First a node-pool is created. The validation test is then run to
+      check if the number of "Suspended" replicas is 0. Once the jobset is
+      running the jobs should quickly enter the "Ready" state. Then using
+      command to suspend entire jobset. The number of found replicas is
+      tested against the number of replicas which should be "Suspended".
+      If they match the DAG is a success.


and please document "how to go through those metrics" here

(not steps, only the concept)

alfredyu-cienet · 2026-05-26T03:03:59Z

let's resolve conflict

alfredyu-cienet · 2026-05-26T04:14:21Z

+    final_expected = jobset_config.replicas
+
+  metric_mapping = {
+      "ready": "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge",


define the enum READY = "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge", or have a method under JobSetHealthiness to return the metric name

alfredyu-cienet · 2026-05-27T02:57:49Z

+@task.sensor(poke_interval=60, timeout=3600, mode="poke")
+def wait_for_jobset_metrics(
+    metric_name: JobSetHealthiness,
+    expected_value: any,


few things related to this one

expected_value: int (for float)

take a look at fix: Separate config construction from Airflow task execution #259, basically the jobset config can be settled in parsing phase and we can do something like this
jobset_config = jobset.build_jobset_from_gcs_yaml(...) wait_for_jobset_metrics(..., expected_value=jobset_config.replicas, ...)

this repo (our fork) has the "GCS_SECRET" set, all you need to do is rebase onto Kim's PR/branch to make it work

that part has merged and dev has been updated, so you can simply rebase FYI

Fixed. Thanks for reminding.

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.

chiajunglien requested review from alfredyu-cienet, andrewyct, severus-ho and xibinliu as code owners May 6, 2026 08:47

chiajunglien added the WIP label May 6, 2026

chiajunglien changed the title ~~feat(tpu): add JobSet healthiness validation DAG and enhance utils~~ Feat: add jobset_healthiness_validation DAG May 7, 2026

chiajunglien removed the WIP label May 7, 2026

alfredyu-cienet reviewed May 7, 2026

View reviewed changes

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

alfredyu-cienet mentioned this pull request May 13, 2026

Adds a new DAG to test "Jobset Ready Healthiness" metric #156

Closed

4 tasks

alfredyu-cienet force-pushed the dev branch 2 times, most recently from 8eb5556 to acf36a1 Compare May 19, 2026 03:54

alfredyu-cienet force-pushed the dev branch from acf36a1 to a529792 Compare May 22, 2026 07:15

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch 2 times, most recently from 916a162 to 1692422 Compare May 26, 2026 03:27

alfredyu-cienet reviewed May 26, 2026

View reviewed changes

alfredyu-cienet force-pushed the dev branch from a529792 to 83f8b37 Compare May 26, 2026 07:41

alfredyu-cienet reviewed May 27, 2026

View reviewed changes

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch from 6feaf0b to 5072531 Compare May 28, 2026 06:23

alfredyu-cienet force-pushed the dev branch from 83f8b37 to ea29ef3 Compare May 28, 2026 07:44

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch from 5072531 to 2ff5a17 Compare May 29, 2026 02:06

chiajunglien added the WIP label May 29, 2026

alfredyu-cienet force-pushed the dev branch 2 times, most recently from f996781 to d3601a9 Compare June 4, 2026 08:06

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch 2 times, most recently from af68025 to 3bd44a5 Compare June 5, 2026 08:54

chiajunglien removed the WIP label Jun 5, 2026

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch from 870840a to 9846760 Compare June 5, 2026 09:04

chiajunglien added the WIP label Jun 15, 2026

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch 3 times, most recently from 5d07788 to fd20cf7 Compare June 15, 2026 03:17

chiajunglien removed the WIP label Jun 15, 2026

alfredyu-cienet force-pushed the dev branch from f8824db to 09c9e42 Compare June 17, 2026 08:08

Jim-0518 and others added 2 commits June 18, 2026 08:37

fix: Skip GCS configuration download in GitHub Actions environment

0791ccb

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.

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

fdcbae9

alfredyu-cienet force-pushed the dev branch from 00a4856 to fdcbae9 Compare June 18, 2026 00:43

chiajunglien added 5 commits June 22, 2026 07:41

feat(tpu): add JobSet healthiness validation DAG and enhance utils

44a9518

Fix: Add validation of success status

b7f2b53

fix

234847f

refactor: encapsulate JobSet metric mapping into JobSetHealthiness enum

99acfc4

fix: Apply TaskGroupWithTimeout into DAG.

c5f8abf

chiajunglien force-pushed the tpu-obs/user/emma/jobset-healthiness-metrics branch from fd20cf7 to c5f8abf Compare June 22, 2026 07:44

alfredyu-cienet force-pushed the dev branch 5 times, most recently from d242c8d to f76e189 Compare July 3, 2026 08:25

alfredyu-cienet force-pushed the dev branch 4 times, most recently from a577785 to 71b8d8a Compare July 14, 2026 03:12

alfredyu-cienet force-pushed the dev branch 4 times, most recently from b588728 to 0e938b1 Compare July 23, 2026 01:43

alfredyu-cienet force-pushed the dev branch 2 times, most recently from 54f9a78 to 8e4ea82 Compare July 27, 2026 07:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Feat: add jobset_healthiness_validation DAG - #256

Feat: add jobset_healthiness_validation DAG#256
chiajunglien wants to merge 7 commits into
devfrom
tpu-obs/user/emma/jobset-healthiness-metrics

chiajunglien commented May 6, 2026 •

edited

Loading

Uh oh!

alfredyu-cienet May 7, 2026

Uh oh!

alfredyu-cienet May 7, 2026

Uh oh!

alfredyu-cienet May 7, 2026

Uh oh!

alfredyu-cienet commented May 26, 2026

Uh oh!

alfredyu-cienet May 26, 2026

Uh oh!

Uh oh!

alfredyu-cienet May 27, 2026

Uh oh!

alfredyu-cienet May 28, 2026

Uh oh!

chiajunglien Jun 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		jobset_start_time: XComArg


		class ReplicatedJobStatus(enum.Enum):

Uh oh!

Conversation

chiajunglien commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Technical Implementation

Airflow/Composer

Required Variables

Uh oh!

alfredyu-cienet May 7, 2026

Choose a reason for hiding this comment

Uh oh!

alfredyu-cienet May 7, 2026

Choose a reason for hiding this comment

Uh oh!

alfredyu-cienet May 7, 2026

Choose a reason for hiding this comment

Uh oh!

alfredyu-cienet commented May 26, 2026

Uh oh!

alfredyu-cienet May 26, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alfredyu-cienet May 27, 2026

Choose a reason for hiding this comment

Uh oh!

alfredyu-cienet May 28, 2026

Choose a reason for hiding this comment

Uh oh!

chiajunglien Jun 5, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

chiajunglien commented May 6, 2026 •

edited

Loading