feat: automate image builds with different libtpu version daily and enhance JobSet DAG validations by yuna-tzeng · Pull Request #246 · CIeNET-International/ml-auto-solutions-3

yuna-tzeng · 2026-04-23T08:55:04Z

Description

This PR introduces an automated pipeline for building and pushing libtpu images and significantly enhances the validation process within our TPU observability DAGs.

Automated libtpu CI/CD (Cloud Build)
Implemented a Cloud Build pipeline to automate the daily building and pushing of libtpu images to GCP.

Introduced a dual-track image strategy:
Stable: Reliable builds for production consistency.
Nightly: Latest builds to catch regressions early.
Enhanced DAG Validation Logic
Refactored tpu_info_format_validation_dag and related JobSet DAGs to sequentially validate both stable and nightly images within a single execution cycle.

Infrastructure Optimization: Updated the workflow to create the TPU node pool once and reuse it for multiple image validations, reducing overall test duration and resource overhead.

Expanded tpu-info table verification to include:
TPU Inbound Buffer Transfer Latency
TPU gRPC TCP Minimum RTT
TPU gRPC TCP Delivery Rate

Tests

Dag Name : tpu_info_format_validation_dag
Airflow test results: https://39b01025d8504fa5afa7226c9e54d883-dot-us-east1.composer.googleusercontent.com/dags/tpu_info_format_validation_dag/grid?dag_run_id=manual__2026-04-23T06%3A35%3A19.595000%2B00%3A00
Other Related DAG : jobset_ttr_node_pool_resize, jobset_rollback_ttr, jobset_ttr_pod_delete, jobset_ttr_kill_process, jobset_uptime_validation, tpu_info_metrics_verification

Airflow/Composer

GCP Composer name: yuna-ml-auto (under GCP project: cloud-ml-auto-solutions)
GCP Composer version: 2.13.1

Required file in GCS bucket ml-auto-solutions-dag-configs/tpu_observability

dag_config.yaml
jobset_config.yaml

Checklist

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-04-24T07:59:37Z

      resources, including the JobSet and the temporary node pools.
      """,
 ) as dag:
+  docker_images = {


since we are adding one additional test

please provide the total execution time, for reviewing the dagrun_timeout

According to my test result, I expect the execution time to be around 30 minutes as well.

for the one you provide, I'd say that's pretty close to 30, might need to set a longer timeout
as for the other ones, please double check

Most libtpu DAGs I checked finish within 30 minutes.
I don't have the resources to test the node pool creation flow right now, so I’m letting it run in the test env for a few days to gather more timing data.
If you're worried about the 30-minute limit, I think we can go with 40 minutes instead.
The node pool creation time is the biggest variable here.

let's do 45

update tpu_info_format_validation_dags and tpu_info_metrics_dag DAG timeout to 45

alfredyu-cienet · 2026-05-15T08:12:07Z

+  TPU_OBS_LIBTPU_NIGHTLY = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest"
+  TPU_OBS_LIBTPU_STABLE = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest"


Suggested change

TPU_OBS_LIBTPU_NIGHTLY = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest"

TPU_OBS_LIBTPU_STABLE = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest"

LIBTPU_NIGHTLY = (

"us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:"

"latest"

)

LIBTPU_STABLE = (

"us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:"

"latest"

)

alfredyu-cienet · 2026-05-15T08:15:33Z

      resources, including the JobSet and the temporary node pools.
      """,
 ) as dag:
+  docker_images = {


let's do 45

alfredyu-cienet · 2026-05-15T08:17:32Z

    # TODO(b/481177412): Re-enable and validate latency metrics.
    # Current Monitoring API aggregation differs from tpu-info, making it
    # unsuitable as a Source of Truth. Investigation for a valid verification
    # method is ongoing.
-    # BufferTransferLatencyStrategy(),
-    # HostToDeviceTransferLatenciesStrategy(),
-    # DeviceToHostTransferLatenciesStrategy(),
-    # CollectiveEndToEndLatencyLatenciesStrategy(),
+    BufferTransferLatencyStrategy(),
+    HostToDeviceTransferLatenciesStrategy(),
+    DeviceToHostTransferLatenciesStrategy(),
+    CollectiveEndToEndLatencyLatenciesStrategy(),


question, does Google aware of we disabled those tests in automation?

No, I don't think they looked that closely into our DAG

I uncommented this section because Victor mentioned recently that there were updates in the latest tpu-info, I want to validate it with automation.
https://b.corp.google.com/issues/481177412#comment5

I provided him with a manual test report comparing the differences between versions 0.10.0 and 0.11.0.
https://docs.google.com/document/d/1hF8yziLPsc7gcCgp9v5pcopkEOus_K4v_451N7rglbU/edit?pli=1&tab=t.0

However, the P99 metrics still show some discrepancies, which I also brought up to Victor during our meeting last week and he will take a look.

- Add TPU_OBS_LIBTPU_ image to vm_resource for environment testing. - Update tpu_info_utils logic to support parsing the latest tpu-info output format. - Add a new DAG to validate compatibility across different TPU images.

…testing Updated the tpu-info format validation DAG to include new network and latency metrics, and refactored the workflow to test multiple libtpu images. Changes: - Added new expected tables to `verify_table_amount`: - "TPU Inbound Buffer Transfer Latency" - "TPU gRPC TCP Minimum RTT" - "TPU gRPC TCP Delivery Rate" - Note: "Host Compute Latency Status" is currently commented out as the metric is broken even in manual testing. - Refactored DAG to run sequential validations for both 'stable' and 'nightly' libtpu Docker images daily. - Optimized node pool usage by creating the infrastructure once and running multiple workloads before cleanup.

Explicitly specify the stable Docker image when building jobsets in: - jobset_ttr_kill_process.py - jobset_ttr_node_pool_resize.py - jobset_ttr_pod_delete_exist.py - jobset_ttr_pod_delete.py - jobset_ttr_rollback.py - jobset_uptime_validation.py This ensures we are using a known stable image for these validation tests.

…imeout to 45

yuna-tzeng requested review from alfredyu-cienet, andrewyct, severus-ho and xibinliu as code owners April 23, 2026 08:55

yuna-tzeng changed the title ~~Tpu obs/user/yuna/cloud build libtpu nightly~~ feat: automate image builds with different libtpu version daily and enhance JobSet DAG validations Apr 24, 2026

alfredyu-cienet reviewed Apr 24, 2026

View reviewed changes

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 3 times, most recently from 721afcd to 8eb5556 Compare May 14, 2026 07:49

yuna-tzeng force-pushed the tpu-obs/user/yuna/cloud_build_libtpu_nightly branch from 94e6e2b to a54d4c8 Compare May 14, 2026 09:38

alfredyu-cienet reviewed May 15, 2026

View reviewed changes

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

yuna-tzeng force-pushed the tpu-obs/user/yuna/cloud_build_libtpu_nightly branch from 83622e9 to 1b87ff6 Compare May 20, 2026 02:27

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

yuna-tzeng force-pushed the tpu-obs/user/yuna/cloud_build_libtpu_nightly branch from 1b87ff6 to dbdc16b Compare May 28, 2026 07:22

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

ea29ef3

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

yuna-tzeng added 8 commits May 28, 2026 07:47

Add nightly Dockerfile

fae3f4b

Add nightly and stable libtpu dockerfile

836da2b

Add argument to Dockerfile-nigtly

13470ee

chore: Change "Latencies" related metric tolerance_percent to 10.0

3d7b51c

format fix

0d06dbc

yuna-tzeng added 3 commits May 28, 2026 08:18

fix format

aab432a

update tpu_info_format_validation_dags and tpu_info_metrics_dag DAG t…

976fbe8

…imeout to 45

Update libtpu image global parameter in vm_resource

8a83f06

yuna-tzeng force-pushed the tpu-obs/user/yuna/cloud_build_libtpu_nightly branch from dbdc16b to 8a83f06 Compare May 28, 2026 08:19

image parameter fix

eef9d27

alfredyu-cienet force-pushed the dev branch 4 times, most recently from 4af284a to a9fa8fb Compare June 8, 2026 08:58

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

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

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 3 times, most recently from d43f051 to a577785 Compare July 9, 2026 08:48

fix

9647123

alfredyu-cienet force-pushed the dev branch 4 times, most recently from ccd7f71 to b588728 Compare July 21, 2026 01:14

alfredyu-cienet force-pushed the dev branch 3 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: automate image builds with different libtpu version daily and enhance JobSet DAG validations - #246

feat: automate image builds with different libtpu version daily and enhance JobSet DAG validations#246
yuna-tzeng wants to merge 19 commits into
devfrom
tpu-obs/user/yuna/cloud_build_libtpu_nightly

yuna-tzeng commented Apr 23, 2026 •

edited

Loading

Uh oh!

alfredyu-cienet Apr 24, 2026

Uh oh!

yuna-tzeng Apr 24, 2026

Uh oh!

alfredyu-cienet Apr 24, 2026

Uh oh!

yuna-tzeng Apr 24, 2026

Uh oh!

alfredyu-cienet May 15, 2026

Uh oh!

yuna-tzeng May 20, 2026

Uh oh!

alfredyu-cienet May 15, 2026

Uh oh!

alfredyu-cienet May 15, 2026

Uh oh!

alfredyu-cienet May 15, 2026

Uh oh!

yuna-tzeng May 20, 2026

Uh oh!

yuna-tzeng May 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		TPU_OBS_LIBTPU_NIGHTLY = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest"
		TPU_OBS_LIBTPU_STABLE = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest"

-  TPU_OBS_LIBTPU_NIGHTLY = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:latest"
-  TPU_OBS_LIBTPU_STABLE = "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:latest"
+  LIBTPU_NIGHTLY = (
+      "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-nightly:"
+      "latest"
+  )
+  LIBTPU_STABLE = (
+      "us-west1-docker.pkg.dev/cienet-cmcs/tpu-obs-cloud-build/libtpu-stable:"
+      "latest"
+  )

Uh oh!

Conversation

yuna-tzeng commented Apr 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Airflow/Composer

Required file in GCS bucket ml-auto-solutions-dag-configs/tpu_observability

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!

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!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yuna-tzeng commented Apr 23, 2026 •

edited

Loading